132 lines
6.0 KiB
Markdown
132 lines
6.0 KiB
Markdown
# 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.
|