116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Полная переиндексация ВСЕХ данных в ChromaDB. Работает до победного."""
|
|
import json, os, sys, time, urllib.request, glob
|
|
from chromadb import PersistentClient
|
|
|
|
EF_URL = "http://localhost:8081/embedding"
|
|
DB_DIR = os.path.expanduser("~/nubes/chroma_db")
|
|
|
|
ALL_DIRS = [
|
|
"~/nubes/data/elsa_jsonl_wi_en", # 244k WI XML
|
|
"~/nubes/data/elsa_mdb_jsonl", # 191k rldal
|
|
"~/nubes/data/elsa_mdb_jsonl/ipsvrap", # 2.5M parts
|
|
"~/nubes/data/elsa_jsonl_htm_en", # 7k hs2+www
|
|
"~/nubes/data/elsa_jsonl", # 24k HTM
|
|
"~/nubes/data/elsa_mdb_jsonl/dbsvrfi", # car ref
|
|
"~/nubes/data/elsa_mdb_jsonl/dbsvrfz", # PR codes
|
|
]
|
|
|
|
def embed(texts):
|
|
for attempt in range(5):
|
|
try:
|
|
data = json.dumps({"content": texts}).encode()
|
|
req = urllib.request.Request(EF_URL, data=data,
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
body = json.loads(resp.read())
|
|
if isinstance(body, list):
|
|
return [item["embedding"][0] for item in body]
|
|
return body.get("embedding", [])
|
|
except Exception as e:
|
|
print(f"\n [embed retry {attempt+1}/5] {e}")
|
|
time.sleep(15)
|
|
raise Exception("Embedding failed after 5 retries")
|
|
|
|
def main():
|
|
# Удалить старую БД и создать новую
|
|
import shutil
|
|
if os.path.exists(DB_DIR):
|
|
shutil.rmtree(DB_DIR)
|
|
print("Старая ChromaDB удалена")
|
|
|
|
client = PersistentClient(path=DB_DIR)
|
|
collection = client.create_collection("elsa_docs")
|
|
|
|
# Собираем все JSONL
|
|
all_files = []
|
|
for d in ALL_DIRS:
|
|
path = os.path.expanduser(d)
|
|
files = sorted(glob.glob(os.path.join(path, "*.jsonl")))
|
|
all_files.extend(files)
|
|
print(f"{d}: {len(files)} files")
|
|
|
|
print(f"\nВсего {len(all_files)} JSONL файлов")
|
|
|
|
total = 0
|
|
errors = 0
|
|
|
|
for fpath in all_files:
|
|
docs, ids = [], []
|
|
prefix = os.path.basename(os.path.dirname(fpath)) + "_"
|
|
|
|
try:
|
|
with open(fpath) as fh:
|
|
for line in fh:
|
|
try:
|
|
d = json.loads(line)
|
|
title = d.get("title", "") or d.get("full_title", "")
|
|
text = d.get("text", "")
|
|
doc = (title + "\n" + text).strip()[:1000]
|
|
if len(doc) < 20:
|
|
continue
|
|
docs.append(doc)
|
|
ids.append(f"{prefix}{os.path.basename(fpath)}_{len(docs)}")
|
|
except:
|
|
continue
|
|
except Exception as e:
|
|
print(f" [SKIP {os.path.basename(fpath)}] read error: {e}")
|
|
errors += 1
|
|
continue
|
|
|
|
if not docs:
|
|
print(f" {os.path.basename(fpath)}: empty")
|
|
continue
|
|
|
|
print(f" {os.path.basename(fpath)}: {len(docs)} docs", end="", flush=True)
|
|
|
|
# Бачами по 50
|
|
for j in range(0, len(docs), 50):
|
|
batch_docs = docs[j:j+50]
|
|
batch_ids = ids[j:j+50]
|
|
|
|
for retry in range(5):
|
|
try:
|
|
embs = embed(batch_docs)
|
|
collection.add(ids=batch_ids, documents=batch_docs,
|
|
embeddings=embs, metadatas=[{"source": prefix.strip("_")}] * len(batch_docs))
|
|
total += len(batch_docs)
|
|
print(".", end="", flush=True)
|
|
break
|
|
except Exception as e:
|
|
if retry < 4:
|
|
print(f"R{retry}", end="", flush=True)
|
|
time.sleep(20)
|
|
else:
|
|
print(f"X", end="", flush=True)
|
|
errors += len(batch_docs)
|
|
|
|
time.sleep(0.2)
|
|
|
|
print(f" total={collection.count()}")
|
|
|
|
print(f"\n✅ Done. Total: {collection.count()} docs. Errors: {errors}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|