141 lines
5.9 KiB
Markdown
141 lines
5.9 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) and a dependency-free **local hashing** fallback (offline, no API key). |
|
|
| `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 fallback
|
|
|
|
This machine (Strix Halo, ~104 GiB unified RAM) has **OOM-killed local models
|
|
before**, so we deliberately do *not* import torch or run a large local
|
|
embedding model here.
|
|
|
|
**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, …).
|
|
|
|
**Offline fallback** — with no key configured, the code automatically uses
|
|
`HashingEmbedder`, a deterministic, dependency-free embedding that projects a
|
|
bag of character n-grams into a fixed-size vector via the hashing trick. It needs
|
|
no weights, no GPU, no internet and no API key, so the whole ANN pipeline is
|
|
fully exercisable for demos and CI. (It is not semantically meaningful — synonyms
|
|
are not mapped together — but it shares enough lexical overlap to cluster near
|
|
the right topics.)
|
|
|
|
### Swapping in a real local model (if you have the RAM)
|
|
|
|
Add a class in `embed.py`:
|
|
|
|
```python
|
|
class LocalEmbedder(Embedder):
|
|
DIM = 768
|
|
def __init__(self):
|
|
from sentence_transformers import SentenceTransformer
|
|
self._model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
def embed(self, texts):
|
|
return self._model.encode(texts, normalize_embeddings=True)
|
|
```
|
|
|
|
and pass it via `embedder=LocalEmbedder()` to the index/query builders. The
|
|
rest of the pipeline is agnostic to where the vectors come from.
|
|
|
|
## 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.
|