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