93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
#!/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")
|