Add arxiv-complete retrieval layer (semantic + keyword)

retrieval/build_index.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""
+build_index.py -- build (and incrementally update) the FAISS ANN index.
+
+This embeds the *entire* metadata corpus (3.15M papers) once and persists a
+FAISS index + metadata sidecar so queries don't have to re-embed anything.
+
+    python build_index.py                 # full corpus, default index "arxiv"
+    python build_index.py --category cs.LG # only cs.LG papers (faster demo)
+    python build_index.py --index-name demo
+
+Re-running it adds any papers not already in the index (incremental build).
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data import DataStore  # noqa: E402
+from embed import get_embedder  # noqa: E402
+from index import IndexBuilder  # noqa: E402
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser(description="Build the arxiv ANN index.")
+    p.add_argument("--index-name", default="arxiv")
+    p.add_argument("--category", default=None, help="Only embed this primary_category.")
+    p.add_argument("--year-min", type=int, default=None)
+    p.add_argument("--year-max", type=int, default=None)
+    p.add_argument("--embedder", default=None, choices=["hashing", "remote"])
+    p.add_argument("--parquet", default=None)
+    args = p.parse_args(argv)
+
+    source = args.parquet or os.environ.get("ARXIV_METADATA_PARQUET", DataStore().source)
+    store = DataStore(source=source)
+    embedder = get_embedder(args.embedder)
+    builder = IndexBuilder(args.index_name, embedder)
+
+    if os.path.exists(builder.default_path()):
+        builder.load()
+        print(f"Loaded existing index: {len(builder._ids)} papers.")
+
+    print("Scanning metadata (filtered) ...")
+    scan = store.scan_filtered(
+        year_min=args.year_min,
+        year_max=args.year_max,
+        primary_category=args.category,
+    )
+    added = builder.build(scan, store)
+    meta = builder.save()
+
+    print(
+        f"Done. Index '{args.index_name}' now holds {len(builder._ids)} papers "
+        f"(+{added}), dim={meta.dim}, embedder={meta.embedder}."
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())