Add arxiv-complete retrieval layer (semantic + keyword)

retrieval/search.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+"""
+search.py -- command-line interface for the arxiv-complete retrieval layer.
+
+Examples
+--------
+    # semantic search (uses the local hashing embedder unless an API key is set)
+    python search.py --query "speculative decoding"
+
+    # restrict to recent papers in a category
+    python search.py --query "neural machine translation" --top-k 15 --year-min 2018 --category cs.CL
+
+    # only survey / overview / taxonomy papers
+    python search.py --query "graph neural networks" --survey-only
+
+    # force the deterministic keyword fallback (no embeddings at all)
+    python search.py --query "speculative decoding" --keyword-only
+
+Build the ANN index first (once) with the build script:
+
+    python build_index.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+# Allow running as a script from the repo root or as a module.
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data import DataStore  # noqa: E402
+from embed import HashingEmbedder, get_embedder  # noqa: E402
+from query import QueryOptions, Result, query  # noqa: E402
+
+
+def _print_results(results: list[Result], query_text: str, method_note: str) -> None:
+    print(f"\nQuery: {query_text!r}")
+    print(f"Index: {method_note}  ({len(results)} result(s))\n")
+    if not results:
+        print("  (no matches)")
+        return
+    for i, r in enumerate(results, 1):
+        print(f"{i:>2}. [{r.method}] {r.score:.3f}  {r.title}")
+        meta = f"   {r.category} · {r.year} · {r.paper_id}"
+        print(meta)
+        print(f"   {r.url}")
+        print(f"   {_snippet_lines(r.snippet)}")
+        print()
+
+
+def _snippet_lines(snippet: str) -> str:
+    snippet = snippet.strip()
+    if not snippet:
+        snippet = "(no abstract available)"
+    width = 88
+    if len(snippet) <= width:
+        return "   " + snippet
+    return "   " + snippet[: width - 3] + "…"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="search.py",
+        description="Search the arxiv-complete corpus (semantic + keyword).",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument("--query", "-q", required=True, help="Natural-language query.")
+    p.add_argument("--top-k", type=int, default=10, help="Max results to return.")
+    p.add_argument("--year-min", type=int, default=None, help="Only papers from this year onward.")
+    p.add_argument("--year-max", type=int, default=None, help="Only papers up to this year.")
+    p.add_argument("--category", default=None, help="Filter by primary_category, e.g. cs.LG.")
+    p.add_argument("--survey-only", action="store_true", help="Only survey/overview/taxonomy papers.")
+    p.add_argument(
+        "--keyword-only",
+        action="store_true",
+        help="Skip the ANN index entirely; use the pure-keyword fallback.",
+    )
+    p.add_argument(
+        "--embedder",
+        default=None,
+        choices=["hashing", "remote"],
+        help="Embedding provider. Default: auto (remote if key set, else local hashing).",
+    )
+    p.add_argument(
+        "--index-name",
+        default="arxiv",
+        help="Name of the persisted FAISS index to use/load.",
+    )
+    p.add_argument(
+        "--parquet",
+        default=None,
+        help="Local parquet path or HF URL overriding the default metadata source.",
+    )
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+
+    store = DataStore(source=args.parquet if args.parquet else os.environ.get(
+        "ARXIV_METADATA_PARQUET", DataStore().source
+    ))
+
+    opts = QueryOptions(
+        top_k=args.top_k,
+        year_min=args.year_min,
+        year_max=args.year_max,
+        primary_category=args.category,
+        survey_only=args.survey_only,
+    )
+
+    if args.keyword_only:
+        results = _keyword_search_only(store, args, opts)
+        _print_results(results, args.query, "keyword fallback (no embeddings)")
+        return 0
+
+    # Semantic path: try to load an existing index; if absent, build it.
+    from index import IndexBuilder, load_index  # imported late to keep CLI fast
+
+    embedder = get_embedder(args.embedder)
+    builder = IndexBuilder(args.index_name, embedder)
+    index_path = builder.default_path()
+
+    if os.path.exists(index_path):
+        try:
+            ann, meta = load_index(args.index_name, embedder)
+            note = f"ANN index '{args.index_name}' ({meta.count} papers, {meta.embedder})"
+        except Exception as exc:
+            print(f"[query] failed to load index ({exc}); building fresh.")
+            ann = None
+    else:
+        print(f"[build] no index at {index_path}; building now (this embeds the corpus).")
+        added = builder.build(store.scan_filtered(), store)
+        meta = builder.save()
+        ann, meta = load_index(args.index_name, embedder)
+        note = f"freshly built ANN index ({meta.count} papers, {meta.embedder})"
+
+    results = query(args.query, store, ann=ann, embedder=embedder, opts=opts)
+    _print_results(results, args.query, note)
+    return 0
+
+
+def _keyword_search_only(store, args, opts) -> list[Result]:
+    from query import _keyword_search
+
+    return _keyword_search(args.query, store, opts)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())