120 lines
3.4 KiB
Python
120 lines
3.4 KiB
Python
"""
|
|
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()
|