# 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`](https://huggingface.co/datasets/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

```bash
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

```bash
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

```bash
# 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:

```bash
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](https://huggingface.co/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.