ElsaWin RAG pipeline: full parsing chain + ChromaDB ingest + docs

2026-06-09 06:58:57 +04:00
parent 89d2301fdf
commit 84ae6cfa9d
11 changed files with 1228 additions and 0 deletions
@@ -0,0 +1,85 @@
+#!/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")