aisys/retrieval/README.md

Retrieval layer for arxiv-complete

A small prototype that turns the full arXiv metadata population (3.15M papers, one parquet file, ~1.6 GB) into a searchable corpus — so you can answer "what is the existing research on topic X?" with real recall and immediacy, instead of the navigational links (blog posts, GitHub lists) a web search returns.

It layers semantic (embedding-based) retrieval on top of keyword search, because keyword-only search misses synonyms.

What it is

File	Role
data.py	Lazy loader for the arxiv-complete metadata parquet (polars scan_parquet). Streams rows so only the matching subset is materialized.
embed.py	Pluggable embedding providers: an OpenAI-compatible remote provider (production), a local semantic sentence-transformers model (default local path), and a dependency-free local hashing fallback (offline, no deps).
index.py	Builds and persists a FAISS flat L2 index + metadata sidecar (paper_id → vector, title, year, category, URL, abstract). Incremental adds.
query.py	The interface. Embeds the query, runs ANN search, ranks, applies filters. Falls back to keyword search if no index is available.
search.py	CLI tying it together.
build_index.py	Embeds the corpus and persists the index.

Data source

• Dataset: secemp9/arxiv-complete
• Config used: metadata — 3,148,796 rows, one per arXiv paper.
• Columns: paper_id, title, authors, abstract, categories, primary_category, submitter, license, doi, journal_ref, comments, report_no, msc_class, acm_class, proxy, n_versions, first_version_date, latest_version_date, oai_datestamp, oai_sets, arxiv_abs_url.
• Snapshot through 2026-08-26. License per paper is mixed (CC-BY, CC-BY-NC-ND, PD, …) and carried in the license column.
• Related configs exist too (versions, files, and content blobs like latex/source/paper_text/ps/pdf) — this prototype only consumes metadata.

Setup

cd aisys/retrieval
uv venv ../.venv && source ../.venv/bin/activate
uv pip install -r requirements.txt

Requires Python 3.10+. faiss-cpu is needed for the ANN path; the keyword fallback works with just polars + numpy.

Running

python search.py --query "speculative decoding"

Options

--query, -q          Natural-language query (required)
--top-k              Max results (default 10)
--year-min           Only papers from this year onward
--year-max           Only papers up to this year
--category           Filter by primary_category, e.g. cs.LG
--survey-only        Only survey/overview/taxonomy papers
--keyword-only       Skip embeddings; pure keyword fallback
--embedder           "hashing" | "remote" (default: auto)
--index-name         Persisted index name (default "arxiv")
--parquet            Override the metadata parquet source

Worked example

# Semantic search (builds the index on first run)
python search.py --query "speculative decoding" --top-k 5

# Keyword-only, no embeddings
python search.py --query "speculative decoding" --keyword-only

# Surveys only, in a category
python search.py --query "speculative decoding" --survey-only --category cs.LG

Embeddings: remote (production) vs local semantic vs hashing

This machine (Strix Halo, ~104 GiB unified RAM) has OOM-killed large local models before, so we deliberately avoid loading big ones. Three options, all behind the same Embedder interface:

Production path — set these env vars and the CLI uses a remote OpenAI-compatible embeddings endpoint:

export EMBEDDING_API_BASE=https://api.openai.com/v1
export EMBEDDING_API_KEY=sk-...
export EMBEDDING_MODEL=text-embedding-3-small

Any endpoint with the OpenAI embeddings shape works (OpenAI, Together, Groq, a local vLLM/Ollama server, …).

Local semantic (default) — with no API key configured, the code automatically uses LocalEmbedder, a small CPU sentence-transformers model (paraphrase-MiniLM-L3-v2, ~90 MB). This produces genuinely semantic vectors (synonyms and paraphrases map close together) at no API cost, and fits comfortably in RAM. It's slower than keyword hashing — embedding the full 3.15M corpus takes ~1-2 hours on this box, so for a demo index a single --category subset instead. Override the model with EMBEDDING_LOCAL_MODEL (e.g. all-MiniLM-L6-v2 for higher quality at ~3x the time).

Dependency-free fallback — pass --embedder hashing (or --keyword-only) for HashingEmbedder, a deterministic embedding that projects a bag of character n-grams into a fixed-size vector via the hashing trick. No weights, no GPU, no internet, no API key — handy for CI or a smoke test, but not semantically meaningful (synonyms aren't mapped together).

How the index works

1. build_index.py scans the metadata (optionally filtered by category/year), embeds title + abstract for each paper, and adds the vectors to a FAISS IndexFlatL2.
2. It persists <index-name>.index (FAISS), <index-name>.meta.pkl (parallel arrays: id, year, category, url, title, abstract), and <index-name>.meta.json (provenance).
3. search.py loads that index, embeds the query, runs ANN search, and attaches metadata to each hit. Re-running build_index.py adds only new papers (incremental, keyed by paper_id).

Notes / decisions

• The corpus is large (3.15M rows), so the data layer streams via lazy scan_parquet and pushes filters (year, category, survey) into the scan rather than loading everything into RAM.
• --survey-only matches abstracts containing survey, comprehensive, overview, or taxonomy. Combined with semantic search this isolates the handful of survey/overview papers about a topic (e.g. ~28 for "speculative decoding").
• The license column is loaded but not yet surfaced in results — a natural follow-up for any downstream licensing filter.