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,115 @@
+#!/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()