test build

backend/scripts/index_documents.py

@@ -0,0 +1,119 @@
+"""
+Index pre-parsed documents (JSON files from parse_pdf.py / parse_wikipedia.py)
+into PostgreSQL, Elasticsearch, FAISS, and MinHash LSH.
+
+Usage:
+    python scripts/index_documents.py <dir_with_json_files> [--limit N]
+"""
+import sys
+import os
+import json
+import argparse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from app.config import settings
+from app.models.db import SessionLocal, Document, Fingerprint, engine, Base
+from app.core.fingerprint import WinNowing, MinHashLSHIndex
+from app.core.embeddings import EmbeddingIndex
+from app.core.fulltext import FulltextSearch
+
+
+Base.metadata.create_all(bind=engine)
+
+
+def load_json_docs(directory: str, limit: int | None = None):
+    docs = []
+    for path in sorted(Path(directory).glob("*.json")):
+        try:
+            with open(path, encoding="utf-8") as f:
+                doc = json.load(f)
+            if doc.get("text") and doc.get("title"):
+                docs.append(doc)
+        except Exception as e:
+            print(f"Skip {path.name}: {e}")
+        if limit and len(docs) >= limit:
+            break
+    return docs
+
+
+def index_batch(docs: list, batch_size: int = 32):
+    db = SessionLocal()
+    winnowing = WinNowing()
+    minhash_lsh = MinHashLSHIndex()
+    embeddings = EmbeddingIndex()
+    fulltext = FulltextSearch()
+
+    print(f"Indexing {len(docs)} documents …")
+
+    embed_batch: list[tuple[str, str]] = []
+
+    for i, doc in enumerate(docs):
+        source = doc.get("source", "unknown")
+        ext_id = doc.get("id", "")
+        title = doc["title"]
+        text = doc["text"]
+        url = doc.get("url", "")
+        lang = doc.get("lang", "ru")
+
+        # 1. PostgreSQL
+        db_doc = Document(source=source, ext_id=ext_id, title=title, url=url, lang=lang)
+        db.add(db_doc)
+        db.flush()
+
+        # 2. Fingerprints → PostgreSQL
+        fp = winnowing.fingerprint(text)
+        for pos, h in enumerate(list(fp)[:500]):
+            db.add(Fingerprint(doc_id=db_doc.id, hash=format(h, '016x'), position=pos))
+
+        # 3. MinHash LSH (in-memory; call save separately if needed)
+        try:
+            minhash_lsh.add_document(str(db_doc.id), text)
+        except ValueError:
+            pass
+
+        # 4. Elasticsearch
+        fulltext.index_document(str(db_doc.id), title, text, source, lang, url)
+
+        # 5. Accumulate for batch embedding
+        embed_batch.append((str(db_doc.id), text))
+
+        if len(embed_batch) >= batch_size:
+            embeddings.batch_add(embed_batch, batch_size=batch_size)
+            embed_batch.clear()
+
+        if (i + 1) % 100 == 0:
+            db.commit()
+            print(f"  {i + 1}/{len(docs)}")
+
+    # Flush remaining embeddings
+    if embed_batch:
+        embeddings.batch_add(embed_batch, batch_size=batch_size)
+
+    db.commit()
+    db.close()
+
+    embeddings.save_index()
+    print("Done. FAISS index saved.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("directory", help="Directory with parsed JSON files")
+    parser.add_argument("--limit", type=int, default=None, help="Max documents to index")
+    parser.add_argument("--batch-size", type=int, default=32)
+    args = parser.parse_args()
+
+    docs = load_json_docs(args.directory, args.limit)
+    print(f"Loaded {len(docs)} documents from {args.directory}")
+
+    if not docs:
+        print("No documents found, exiting.")
+        sys.exit(0)
+
+    index_batch(docs, batch_size=args.batch_size)
+
+
+if __name__ == "__main__":
+    main()