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,92 @@
+#!/usr/bin/env python3
+"""Ingest WI JSONL into ChromaDB — one file at a time, batches of 50."""
+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")
+# Default: ingest ALL JSONL dirs
+ALL_DIRS = [
+    os.path.expanduser("~/nubes/data/elsa_jsonl_wi_en"),        # 244k WI XML
+    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl"),            # 191k rldal
+    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/ipsvrap"),   # 2.5M parts
+    os.path.expanduser("~/nubes/data/elsa_jsonl_htm_en"),        # 7k hs2+www
+    os.path.expanduser("~/nubes/data/elsa_jsonl"),               # 24k HTM
+    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/dbsvrfi"),   # car ref
+    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/dbsvrfz"),   # PR codes
+]
+
+DATA_DIR = sys.argv[1] if len(sys.argv) > 1 else None
+COLLECTION_NAME = "elsa_docs"
+
+# Connect to existing ChromaDB
+client = PersistentClient(path=DB_DIR)
+try:
+    collection = client.get_or_create_collection(COLLECTION_NAME)
+except:
+    collection = client.create_collection(COLLECTION_NAME)
+
+existing = collection.count()
+print(f"Collection has {existing} docs already")
+
+# Find remaining files
+if DATA_DIR:
+    files = sorted(glob.glob(os.path.join(DATA_DIR, "*.jsonl")))
+else:
+    files = sorted(glob.glob(os.path.join(ALL_DIRS[0], "*.jsonl")))
+    for d in ALL_DIRS[1:]:
+        files.extend(sorted(glob.glob(os.path.join(d, "*.jsonl"))))
+
+print(f"Processing {len(files)} JSONL files")
+
+for fpath in files:
+    docs, ids = [], []
+    prefix = os.path.basename(os.path.dirname(fpath)) + "_"
+    with open(fpath) as fh:
+        for line in fh:
+            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)}")
+    
+    if not docs:
+        print(f"{os.path.basename(fpath)}: empty, skip")
+        continue
+    
+    print(f"{os.path.basename(fpath)}: {len(docs)} docs", end="")
+    
+    for j in range(0, len(docs), 50):
+        batch_docs = docs[j:j+50]
+        batch_ids = ids[j:j+50]
+        data = json.dumps({"content": batch_docs}).encode()
+        
+        for retry in range(3):
+            try:
+                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):
+                    embs = [item["embedding"][0] for item in body]
+                else:
+                    embs = body.get("embedding", [])
+                
+                collection.add(ids=batch_ids, documents=batch_docs, embeddings=embs)
+                print(".", end="", flush=True)
+                break
+            except Exception as e:
+                if retry < 2:
+                    print(f"R{retry}", end="", flush=True)
+                    time.sleep(10)
+                else:
+                    print(f"X({e})", end="", flush=True)
+        
+        time.sleep(0.3)
+    
+    print(f" total={collection.count()}")
+
+print(f"\nDone. Total: {collection.count()} documents")