65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
#!/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())
|