Add arxiv-complete retrieval layer (semantic + keyword)

2026-09-20 05:49:06 -07:00
parent aab8e35fe9
commit bf475d5e1e
9 changed files with 1169 additions and 0 deletions
@@ -0,0 +1,140 @@
+# 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.