82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
||
"""JSONL 去重脚本 — 清理 sbert_docs.jsonl 中的重复 ID
|
||
|
||
逻辑：
|
||
1. 读取 sbert_docs.jsonl，按 ID 去重（保留第一条）
|
||
2. 同步截断 sbert_vectors.jsonl 到对应长度
|
||
3. 重建 FAISS 索引
|
||
4. 更新 sbert_meta.json
|
||
"""
|
||
import json
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
DATA_DIR = Path.home() / "projects" / "zhiyi" / "data"
|
||
DOCS_FILE = DATA_DIR / "sbert_docs.jsonl"
|
||
VECTORS_FILE = DATA_DIR / "sbert_vectors.jsonl"
|
||
META_FILE = DATA_DIR / "sbert_meta.json"
|
||
|
||
print("=== JSONL 去重 ===")
|
||
|
||
# 1. 读取 docs 并去重
|
||
seen_ids = []
|
||
unique_docs = []
|
||
dup_count = 0
|
||
|
||
with open(DOCS_FILE) as f:
|
||
for i, line in enumerate(f):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
doc = json.loads(line)
|
||
doc_id = doc.get('id', f'line_{i}')
|
||
if doc_id in seen_ids:
|
||
print(f" 重复 ID: {doc_id[:30]} (line {i+1})")
|
||
dup_count += 1
|
||
else:
|
||
seen_ids.append(doc_id)
|
||
unique_docs.append(doc)
|
||
|
||
print(f"\n去重结果：{len(seen_ids) + dup_count} 行 → {len(unique_docs)} 唯一文档（移除 {dup_count} 个重复）")
|
||
|
||
if dup_count == 0:
|
||
print("无需去重，文件已干净。")
|
||
exit(0)
|
||
|
||
# 2. 备份
|
||
backup_docs = DOCS_FILE.with_suffix('.jsonl.bak')
|
||
shutil.copy2(DOCS_FILE, backup_docs)
|
||
print(f"已备份: {backup_docs}")
|
||
|
||
# 3. 重写 docs.jsonl（去重后）
|
||
with open(DOCS_FILE, 'w') as f:
|
||
for doc in unique_docs:
|
||
f.write(json.dumps(doc, ensure_ascii=False) + '\n')
|
||
print(f"已重写: {DOCS_FILE}")
|
||
|
||
# 4. 截断 vectors.jsonl 到对应长度
|
||
vec_count = 0
|
||
with open(VECTORS_FILE) as f:
|
||
vec_count = sum(1 for line in f if line.strip())
|
||
|
||
if vec_count > len(unique_docs):
|
||
print(f"向量文件: {vec_count} 行 → 截断到 {len(unique_docs)} 行")
|
||
backup_vec = VECTORS_FILE.with_suffix('.jsonl.bak')
|
||
shutil.copy2(VECTORS_FILE, backup_vec)
|
||
with open(VECTORS_FILE) as f:
|
||
lines = [line for line in f if line.strip()]
|
||
with open(VECTORS_FILE, 'w') as f:
|
||
f.write(''.join(lines[:len(unique_docs)]))
|
||
print(f"已截断: {VECTORS_FILE}")
|
||
elif vec_count == len(unique_docs):
|
||
print(f"向量文件: {vec_count} 行，无需修改")
|
||
else:
|
||
print(f"警告: 向量文件 ({vec_count}) 比文档 ({len(unique_docs)}) 少！需要重建索引。")
|
||
|
||
# 5. 更新 meta
|
||
meta = {'model_name': 'moka-ai/m3e-base', 'doc_count': len(unique_docs)}
|
||
with open(META_FILE, 'w') as f:
|
||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||
print(f"已更新: {META_FILE}")
|
||
|
||
print("\n去重完成！下一步：重建 FAISS 索引") |