126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Ingest JSONL into ChromaDB — uses llama.cpp /embedding endpoint."""
|
|
import json, os, sys, glob, urllib.request, time
|
|
import chromadb
|
|
|
|
LLAMA_URL = "http://localhost:8081/embedding"
|
|
|
|
def llama_embed(texts):
|
|
"""Get embeddings from nomic-embed server with retry."""
|
|
data = json.dumps({"content": texts}).encode()
|
|
req = urllib.request.Request(LLAMA_URL, data=data,
|
|
headers={"Content-Type": "application/json"})
|
|
for attempt in range(3):
|
|
try:
|
|
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:
|
|
if attempt < 2:
|
|
print(f"[retry {attempt+1}] {e}")
|
|
time.sleep(5)
|
|
else:
|
|
raise
|
|
|
|
|
|
class LlamaEmbedFn:
|
|
"""Custom embedding function using llama.cpp."""
|
|
def __call__(self, input):
|
|
all_embs = []
|
|
for i in range(0, len(input), 32):
|
|
batch = input[i:i+32]
|
|
embs = llama_embed(batch)
|
|
if isinstance(embs[0], float):
|
|
all_embs.append(embs)
|
|
else:
|
|
all_embs.extend(embs)
|
|
return all_embs
|
|
|
|
|
|
def main():
|
|
jsonl_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser("~/nubes/data/elsa_jsonl_wi_en")
|
|
db_dir = sys.argv[2] if len(sys.argv) > 2 else os.path.expanduser("~/nubes/chroma_db")
|
|
collection_name = sys.argv[3] if len(sys.argv) > 3 else "elsa_docs"
|
|
|
|
ef = LlamaEmbedFn()
|
|
|
|
client = chromadb.PersistentClient(path=db_dir)
|
|
|
|
try:
|
|
client.delete_collection(collection_name)
|
|
except:
|
|
pass
|
|
|
|
collection = client.create_collection(
|
|
name=collection_name,
|
|
embedding_function=ef,
|
|
metadata={"hnsw:space": "cosine"}
|
|
)
|
|
|
|
jsonl_files = sorted(glob.glob(os.path.join(jsonl_dir, "*.jsonl")))
|
|
print(f"Found {len(jsonl_files)} JSONL files")
|
|
print(f"Using llama.cpp at {LLAMA_URL}")
|
|
|
|
total = 0
|
|
for jf in jsonl_files:
|
|
ids, docs, metas = [], [], []
|
|
|
|
with open(jf, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
try:
|
|
d = json.loads(line)
|
|
title = d.get("title", "").strip()
|
|
text = d.get("text", "").strip()
|
|
ft = d.get("full_title", "").strip()
|
|
|
|
if ft and text:
|
|
doc = f"{ft}\n{text}"
|
|
elif title and text:
|
|
doc = f"{title}\n{text}"
|
|
else:
|
|
doc = text or title or ft
|
|
|
|
if not doc or len(doc) < 20:
|
|
continue
|
|
|
|
# Split long docs into chunks of ~1000 chars
|
|
chunks = [doc[i:i+1000] for i in range(0, len(doc), 1000)]
|
|
for chunk in chunks:
|
|
if len(chunk) < 20:
|
|
continue
|
|
ids.append(str(total))
|
|
docs.append(chunk)
|
|
metas.append({
|
|
"source": d.get("source", ""),
|
|
"file": d.get("file", ""),
|
|
"title": title or ft or doc[:100],
|
|
})
|
|
total += 1
|
|
except:
|
|
continue
|
|
|
|
if ids:
|
|
# Send in smaller batches of 500
|
|
for i in range(0, len(ids), 500):
|
|
batch_ids = ids[i:i+500]
|
|
batch_docs = docs[i:i+500]
|
|
batch_metas = metas[i:i+500]
|
|
for attempt in range(3):
|
|
try:
|
|
collection.add(ids=batch_ids, documents=batch_docs, metadatas=batch_metas)
|
|
break
|
|
except Exception as e:
|
|
if attempt < 2:
|
|
print(f" [retry {os.path.basename(jf)} batch {attempt+1}] {e}")
|
|
time.sleep(10)
|
|
else:
|
|
print(f" [SKIP batch] {e}")
|
|
print(f" {os.path.basename(jf)}: +{len(ids)} (total {total})")
|
|
|
|
print(f"\nDone. {total} documents → {db_dir}/{collection_name}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|