86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Полная индексация через sentence-transformers (CPU, стабильно)."""
|
|
import json, os, sys, time, glob, shutil
|
|
from chromadb import PersistentClient
|
|
from chromadb.utils import embedding_functions
|
|
|
|
ALL_DIRS = [
|
|
"~/nubes/data/elsa_jsonl_wi_en",
|
|
"~/nubes/data/elsa_mdb_jsonl",
|
|
"~/nubes/data/elsa_mdb_jsonl/ipsvrap",
|
|
"~/nubes/data/elsa_jsonl_htm_en",
|
|
"~/nubes/data/elsa_jsonl",
|
|
"~/nubes/data/elsa_mdb_jsonl/dbsvrfi",
|
|
"~/nubes/data/elsa_mdb_jsonl/dbsvrfz",
|
|
]
|
|
|
|
DB_DIR = os.path.expanduser("~/nubes/chroma_db")
|
|
MODEL = "all-MiniLM-L6-v2" # 90 MB, CPU
|
|
|
|
ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL)
|
|
|
|
if os.path.exists(DB_DIR):
|
|
shutil.rmtree(DB_DIR)
|
|
print("Old chroma_db deleted")
|
|
|
|
client = PersistentClient(path=DB_DIR)
|
|
collection = client.create_collection("elsa_docs", embedding_function=ef)
|
|
|
|
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"\nTotal: {len(all_files)} JSONL files")
|
|
|
|
total = 0
|
|
for fpath in all_files:
|
|
docs, ids, metas = [], [], []
|
|
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)}")
|
|
metas.append({"source": prefix.strip("_")})
|
|
except:
|
|
continue
|
|
except:
|
|
continue
|
|
|
|
if not docs:
|
|
continue
|
|
|
|
print(f"{os.path.basename(fpath)}: {len(docs)} docs", end="", flush=True)
|
|
|
|
for j in range(0, len(docs), 500):
|
|
batch = docs[j:j+500]
|
|
batch_ids = ids[j:j+500]
|
|
for retry in range(3):
|
|
try:
|
|
collection.add(ids=batch_ids, documents=batch,
|
|
metadatas=metas[j:j+500])
|
|
total += len(batch)
|
|
print(".", end="", flush=True)
|
|
break
|
|
except Exception as e:
|
|
if retry < 2:
|
|
print(f"R{retry}", end="", flush=True)
|
|
time.sleep(5)
|
|
else:
|
|
print(f"X", end="", flush=True)
|
|
|
|
print(f" total={collection.count()}")
|
|
|
|
print(f"\n✅ Done. Total: {collection.count()} docs")
|