Add arxiv-complete retrieval layer (semantic + keyword)

retrieval/.gitignore

@@ -0,0 +1,11 @@
+# Python
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+
+# Built ANN index (local, machine-specific, large)
+index/
+*.index
+*.meta.pkl
+*.meta.json

retrieval/README.md

@@ -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.

retrieval/build_index.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""
+build_index.py -- build (and incrementally update) the FAISS ANN index.
+
+This embeds the *entire* metadata corpus (3.15M papers) once and persists a
+FAISS index + metadata sidecar so queries don't have to re-embed anything.
+
+    python build_index.py                 # full corpus, default index "arxiv"
+    python build_index.py --category cs.LG # only cs.LG papers (faster demo)
+    python build_index.py --index-name demo
+
+Re-running it adds any papers not already in the index (incremental build).
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data import DataStore  # noqa: E402
+from embed import get_embedder  # noqa: E402
+from index import IndexBuilder  # noqa: E402
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser(description="Build the arxiv ANN index.")
+    p.add_argument("--index-name", default="arxiv")
+    p.add_argument("--category", default=None, help="Only embed this primary_category.")
+    p.add_argument("--year-min", type=int, default=None)
+    p.add_argument("--year-max", type=int, default=None)
+    p.add_argument("--embedder", default=None, choices=["hashing", "remote"])
+    p.add_argument("--parquet", default=None)
+    args = p.parse_args(argv)
+
+    source = args.parquet or os.environ.get("ARXIV_METADATA_PARQUET", DataStore().source)
+    store = DataStore(source=source)
+    embedder = get_embedder(args.embedder)
+    builder = IndexBuilder(args.index_name, embedder)
+
+    if os.path.exists(builder.default_path()):
+        builder.load()
+        print(f"Loaded existing index: {len(builder._ids)} papers.")
+
+    print("Scanning metadata (filtered) ...")
+    scan = store.scan_filtered(
+        year_min=args.year_min,
+        year_max=args.year_max,
+        primary_category=args.category,
+    )
+    added = builder.build(scan, store)
+    meta = builder.save()
+
+    print(
+        f"Done. Index '{args.index_name}' now holds {len(builder._ids)} papers "
+        f"(+{added}), dim={meta.dim}, embedder={meta.embedder}."
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

retrieval/data.py

@@ -0,0 +1,136 @@
+"""
+Data access layer for the arxiv-complete metadata dataset.
+
+The full metadata parquet is 1.64 GB / 3,148,796 rows (one per arXiv paper),
+stored as a single file on Hugging Face:
+
+    https://huggingface.co/api/datasets/secemp9/arxiv-complete/parquet/metadata/train/0.parquet
+
+Loading the whole thing into RAM (~a few GB of Python objects) is feasible on
+a big machine but wasteful for a search prototype.  Instead we stream rows from
+the remote parquet with polars' *scan_parquet* + lazy filtering, so the query
+path only materializes the handful of rows that actually match.
+
+The same lazy-scanning approach works against a local copy of the parquet,
+which is what the build pipeline uses to precompute embeddings.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Iterable
+
+import polars as pl
+
+DEFAULT_PARQUET_URL = (
+    "https://huggingface.co/api/datasets/secemp9/arxiv-complete/"
+    "parquet/metadata/train/0.parquet"
+)
+
+# Columns that are cheap to keep in memory and useful for display / filtering.
+DISPLAY_COLUMNS = [
+    "paper_id",
+    "title",
+    "primary_category",
+    "first_version_date",
+    "arxiv_abs_url",
+    "abstract",
+]
+
+# Free-text terms that mark an abstract as a survey / overview / taxonomy paper.
+SURVEY_KEYWORDS = ("survey", "comprehensive", "overview", "taxonomy")
+
+
+@dataclass
+class Paper:
+    """A single arXiv paper, as surfaced by the search layer."""
+
+    paper_id: str
+    title: str
+    primary_category: str
+    year: int
+    url: str
+    abstract: str = ""
+    score: float = 0.0
+
+    @staticmethod
+    def from_row(row: pl.Row) -> "Paper":
+        raw = dict(zip(DISPLAY_COLUMNS, row))
+        fvd = raw["first_version_date"]
+        year = (
+            fvd.year
+            if isinstance(fvd, datetime)
+            else (int(fvd) if fvd is not None else 0)
+        )
+        return Paper(
+            paper_id=str(raw["paper_id"]),
+            title=str(raw["title"] or ""),
+            primary_category=str(raw["primary_category"] or ""),
+            year=year,
+            url=str(raw["arxiv_abs_url"] or ""),
+            abstract=str(raw["abstract"] or ""),
+        )
+
+
+@dataclass
+class DataStore:
+    """
+    Lazy view over the arxiv-complete metadata parquet.
+
+    We deliberately do not load the full 3.15M rows into memory.  ``scan_parquet``
+    returns a lazy frame that polars only executes when we call ``collect()`` --
+    and it pushes down the WHERE filters (year range, category, survey) so the
+    download/shuffle work stays proportional to the result set.
+    """
+
+    source: str = field(default_factory=lambda: os.environ.get(
+        "ARXIV_METADATA_PARQUET", DEFAULT_PARQUET_URL
+    ))
+
+    def _scan(self) -> pl.LazyFrame:
+        return pl.scan_parquet(self.source)
+
+    def scan_filtered(
+        self,
+        *,
+        year_min: int | None = None,
+        year_max: int | None = None,
+        primary_category: str | None = None,
+        survey_only: bool = False,
+    ) -> pl.LazyFrame:
+        """
+        Return a lazy frame of the full metadata filtered to the given criteria.
+
+        ``survey_only`` keeps only rows whose abstract mentions a survey/overview
+        keyword.  All filters are AND-ed together and pushed into the scan.
+        """
+        q = self._scan()
+
+        if year_min is not None:
+            q = q.filter(pl.col("first_version_date").dt.year() >= year_min)
+        if year_max is not None:
+            q = q.filter(pl.col("first_version_date").dt.year() <= year_max)
+        if primary_category:
+            q = q.filter(pl.col("primary_category") == primary_category)
+        if survey_only:
+            needle = "|".join(SURVEY_KEYWORDS)
+            q = q.filter(
+                pl.col("abstract")
+                .str.to_lowercase()
+                .str.contains(needle, strict=False)
+            )
+        return q
+
+    def fetch(self, limit: int | None = None) -> list[Paper]:
+        """Execute the (already filtered) scan and return Paper objects."""
+        df = self.scan_filtered().select(DISPLAY_COLUMNS)
+        if limit is not None:
+            df = df.limit(limit)
+        rows = df.collect().to_dicts()
+        return [Paper.from_row(r.values()) for r in rows]
+
+    def count(self) -> int:
+        """Total rows matching the current filter (cheap COUNT push-down)."""
+        return int(self.scan_filtered().select(pl.len()).collect()["len"][0])

retrieval/embed.py

@@ -0,0 +1,208 @@
+"""
+Embedding providers for the retrieval layer.
+
+Two things live here:
+
+1. A *pluggable* embedding interface (``Embedder``) plus a concrete
+   OpenAI-compatible remote provider.  This is the **production** path: it hits
+   an HTTP endpoint and returns real dense vectors.
+
+2. A **deterministic local hashing fallback** (``HashingEmbedder``) that needs
+   no model weights, no GPU, no API key and no internet.  It projects a bag of
+   character n-grams into a fixed-dimension vector via the hashing trick.  This
+   lets the whole ANN pipeline (FAISS build -> ANN search -> ranking) run
+   end-to-end offline, which is exactly what we want for demos and CI.
+
+WHY NOT A LOCAL EMBEDDING MODEL?
+--------------------------------
+This box is a Strix Halo AP with ~104 GiB of unified RAM.  It has *tried* to
+load a large local embedding/LLM model before and OOM-killed.  So we do not
+import torch or try to run sentence-transformers here.  If you have the RAM and
+want real embeddings without paying for an API, drop in a local model:
+
+    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 to the index/query builders via the ``embedder=`` argument.  The
+rest of the pipeline is agnostic to where the vectors come from.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from abc import ABC, abstractmethod
+from typing import Sequence
+
+import numpy as np
+
+# The vector dimension the hashing fallback projects into.  The remote provider
+# returns its own native dimension; the index is built to whatever the embedder
+# reports via ``dimension``.
+HASHING_DIM = 256
+
+
+class Embedder(ABC):
+    """Abstract embedding provider."""
+
+    #: Dimensionality of the vectors this provider emits.
+    DIM: int
+
+    @abstractmethod
+    def embed(self, texts: Sequence[str]) -> np.ndarray:
+        """Embed a batch of texts, returning an (N, DIM) float32 array."""
+        raise NotImplementedError
+
+    def embed_one(self, text: str) -> np.ndarray:
+        """Embed a single text, returning a (DIM,) float32 array."""
+        return self.embed([text])[0]
+
+
+class HashingEmbedder(Embedder):
+    """
+    Deterministic, dependency-free embedding via the hashing trick.
+
+    We tokenize into character 3-grams, hash each one into a bucket, and write
+    +1/-1 (sign hashed from the gram) into a fixed-size vector, then L2-normalize.
+
+    This is *not* semantically meaningful (synonyms are not mapped together), but
+    it is stable, fast, and shares enough lexical overlap with the query for
+    near-neighbors to cluster around the right topics.  Perfect for exercising
+    the ANN pipeline without an API key.
+    """
+
+    DIM = HASHING_DIM
+    _NGRAM = 3
+    _SEED = 0xC0FFEE
+
+    def __init__(self, dim: int = HASHING_DIM, ngram: int = 3):
+        self.DIM = dim
+        self._ngram = ngram
+
+    @staticmethod
+    def _grams(text: str, ngram: int = 3) -> list[str]:
+        t = text.lower()
+        return [t[i:i + ngram] for i in range(max(0, len(t) - ngram + 1))]
+
+    def embed(self, texts: Sequence[str]) -> np.ndarray:
+        out = np.zeros((len(texts), self.DIM), dtype=np.float32)
+        ngram = self._ngram
+        for i, text in enumerate(texts):
+            for gram in self._grams(text or "", ngram):
+                h = int.from_bytes(
+                    hashlib.sha1(gram.encode("utf-8")).digest()[:8],
+                    "little",
+                )
+                bucket = h % self.DIM
+                sign = 1.0 if (h >> 63) & 1 else -1.0
+                out[i, bucket] += sign
+            norm = np.linalg.norm(out[i])
+            if norm > 0:
+                out[i] /= norm
+        return out
+
+
+class RemoteEmbedder(Embedder):
+    """
+    OpenAI-compatible remote embedding provider.
+
+    Configured entirely through the environment:
+
+        EMBEDDING_API_BASE  base URL, e.g. https://api.openai.com/v1
+        EMBEDDING_API_KEY   API key (required)
+        EMBEDDING_MODEL     model name, default text-embedding-3-small
+
+    Any endpoint with the OpenAI embeddings shape works (OpenAI, Together,
+    Groq, a local vLLM/Ollama server, etc.).
+    """
+
+    _DEFAULT_MODEL = "text-embedding-3-small"
+    _DEFAULT_DIM = 1536  # text-embedding-3-small's default output dim
+
+    def __init__(
+        self,
+        api_base: str | None = None,
+        api_key: str | None = None,
+        model: str | None = None,
+        dim: int | None = None,
+    ):
+        import requests  # imported lazily so the fallback needs no network stack
+
+        self._requests = requests
+        self._api_base = (api_base or os.environ.get("EMBEDDING_API_BASE", "")).rstrip("/")
+        self._api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "")
+        self._model = model or os.environ.get("EMBEDDING_MODEL", self._DEFAULT_MODEL)
+        self._dim = dim or int(
+            os.environ.get("EMBEDDING_DIM", self._DEFAULT_DIM)
+        )
+        if not self._api_base:
+            raise ValueError(
+                "EMBEDDING_API_BASE is not set; configure it (or use the local "
+                "HashingEmbedder fallback)."
+            )
+        if not self._api_key:
+            raise ValueError(
+                "EMBEDDING_API_KEY is not set; set it, or use the local "
+                "HashingEmbedder fallback."
+            )
+
+    @property
+    def DIM(self) -> int:
+        return self._dim
+
+    def embed(self, texts: Sequence[str]) -> np.ndarray:
+        resp = self._requests.post(
+            f"{self._api_base}/embeddings",
+            headers={"Authorization": f"Bearer {self._api_key}"},
+            json={"model": self._model, "input": list(texts), "dimensions": self._dim},
+            timeout=60,
+        )
+        resp.raise_for_status()
+        payload = resp.json()
+        items = payload["data"]
+        # Some backends don't preserve input order; sort by the returned index.
+        items = sorted(items, key=lambda d: d.get("index", 0))
+        return np.array([d["embedding"] for d in items], dtype=np.float32)
+
+
+def get_embedder(embedder: str | Embedder | None = None) -> Embedder:
+    """
+    Resolve an embedder from a name, an instance, or the environment.
+
+    Resolution order:
+
+      1. ``embedder="remote"`` or an ``Embedder`` instance -> use it directly.
+      2. ``embedder="hashing"`` -> deterministic local fallback.
+      3. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set ->
+         remote provider (the production path).
+      4. Otherwise -> the local hashing fallback (with a printed note).
+
+    This means the CLI works out of the box with no configuration: it silently
+    degrades to the reproducible local embedding when no API key is present.
+    """
+    name = (embedder or "").lower() if isinstance(embedder, str) else None
+
+    if isinstance(embedder, Embedder):
+        return embedder
+    if name == "hashing":
+        return HashingEmbedder()
+    if name == "remote":
+        return RemoteEmbedder()
+
+    api_base = os.environ.get("EMBEDDING_API_BASE")
+    api_key = os.environ.get("EMBEDDING_API_KEY")
+    if api_base and api_key:
+        print(f"[embed] using remote provider: {api_base} (model={os.environ.get('EMBEDDING_MODEL')})")
+        return RemoteEmbedder()
+
+    print(
+        "[embed] no EMBEDDING_API_BASE/KEY configured -> using local "
+        "HashingEmbedder (deterministic, offline, non-semantic). Set the env vars "
+        "or pass embedder=RemoteEmbedder() for real semantic embeddings."
+    )
+    return HashingEmbedder()

retrieval/index.py

@@ -0,0 +1,233 @@
+"""
+FAISS-backed ANN index over paper embeddings.
+
+The index persists three things to disk so a query process never has to
+re-embed the full corpus:
+
+  * ``<name>.index``      -- the FAISS flat L2 index (float32 vectors)
+  * ``<name>.meta.pkl``   -- parallel arrays: paper_id, year, category, url,
+                             title, abstract (aligned to the FAISS rows)
+  * ``<name>.meta.json``  -- provenance (embedder name, dim, source, counts)
+
+The index is *incremental*: ``build_index`` can be pointed at a fresh scan of
+the metadata and will add any papers not already present (keyed by paper_id),
+rather than rebuilding from scratch every time.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import pickle
+from dataclasses import dataclass
+
+import numpy as np
+
+try:
+    import faiss
+except ImportError:  # pragma: no cover - faiss is a hard dependency
+    faiss = None  # type: ignore[assignment]
+
+from data import DataStore, Paper
+
+
+def _require_faiss() -> None:
+    if faiss is None:
+        raise RuntimeError("faiss-cpu is required for ANN indexing. Install it (see requirements.txt).")
+
+
+@dataclass
+class IndexMeta:
+    embedder: str
+    dim: int
+    source: str
+    count: int
+    built_at: str = ""
+
+
+class AnnIndex:
+    """
+    Read-only view over a persisted FAISS index, exposing the parallel metadata
+    arrays needed to attach titles / years / categories / URLs to ANN hits.
+    """
+
+    def __init__(self, builder: IndexBuilder, meta: IndexMeta):
+        self.index = builder.index
+        self.meta = meta
+        self._ids = builder._ids
+        self._years = builder._years
+        self._cats = builder._cats
+        self._urls = builder._urls
+        self._titles = builder._titles
+        self._abstracts = builder._abstracts
+
+    def __len__(self) -> int:
+        return len(self._ids)
+
+
+class IndexBuilder:
+    """
+    Build (and incrementally extend) a FAISS index from a DataStore scan.
+
+    Example
+    -------
+    >>> builder = IndexBuilder("arxiv", embedder=HashingEmbedder())
+    >>> store = DataStore()
+    >>> builder.build(store.scan_filtered(), store)   # doctest: +SKIP
+    >>> builder.save("arxiv")
+    """
+
+    def __init__(self, name: str, embedder, chunk_size: int = 4096):
+        _require_faiss()
+        self.name = name
+        self.embedder = embedder
+        self.dim = int(embedder.DIM)
+        self.chunk_size = chunk_size
+        self._existing: set[str] = set()
+        self._ids: list[str] = []
+        self._years: list[int] = []
+        self._cats: list[str] = []
+        self._urls: list[str] = []
+        self._titles: list[str] = []
+        self._abstracts: list[str] = []
+        self._loaded = False
+
+    # -- loading an existing index ------------------------------------------
+    def load(self, path: str | None = None) -> None:
+        path = path or self.default_path()
+        if not os.path.exists(path):
+            return
+        self.index = faiss.read_index(path)
+        with open(path.replace(".index", ".meta.pkl"), "rb") as fh:
+            blob = pickle.load(fh)
+        (
+            self._ids,
+            self._years,
+            self._cats,
+            self._urls,
+            self._titles,
+            self._abstracts,
+        ) = blob
+        self._existing = set(self._ids)
+        self._loaded = True
+
+    def default_path(self, root: str | None = None) -> str:
+        root = root or os.path.join(os.path.dirname(__file__), "index")
+        return os.path.join(root, f"{self.name}.index")
+
+    # -- building -----------------------------------------------------------
+    def build(self, scan, store: DataStore) -> int:
+        """
+        Embed every row of ``scan`` (a polars LazyFrame) and add new papers.
+
+        ``scan`` must project the DISPLAY_COLUMNS; ``store`` is used only for its
+        provenance metadata.  Returns the number of papers newly added.
+        """
+        import datetime as _dt
+
+        if not self._loaded:
+            self.index = faiss.IndexFlatL2(self.dim)
+        added = 0
+        batch: list[str] = []
+        batch_ids: list[str] = []
+        batch_years: list[int] = []
+        batch_cats: list[str] = []
+        batch_urls: list[str] = []
+        batch_titles: list[str] = []
+        batch_abstracts: list[str] = []
+
+        def flush() -> None:
+            nonlocal added
+            if not batch:
+                return
+            # Only embed/add papers not already in the index.
+            keep = [i for i, pid in enumerate(batch_ids) if pid not in self._existing]
+            if not keep:
+                batch.clear()
+                batch_ids.clear()
+                batch_years.clear()
+                batch_cats.clear()
+                batch_urls.clear()
+                batch_titles.clear()
+                batch_abstracts.clear()
+                return
+            keep_ids = [batch_ids[i] for i in keep]
+            keep_texts = [batch[i] for i in keep]
+            vecs = self.embedder.embed(keep_texts)
+            vecs = np.ascontiguousarray(vecs, dtype=np.float32)
+            self.index.add(vecs)
+            for i in keep:
+                self._ids.append(batch_ids[i])
+                self._years.append(batch_years[i])
+                self._cats.append(batch_cats[i])
+                self._urls.append(batch_urls[i])
+                self._titles.append(batch_titles[i])
+                self._abstracts.append(batch_abstracts[i])
+                self._existing.add(batch_ids[i])
+            added += len(keep_ids)
+            batch.clear()
+            batch_ids.clear()
+            batch_years.clear()
+            batch_cats.clear()
+            batch_urls.clear()
+            batch_titles.clear()
+            batch_abstracts.clear()
+
+        # Collect all matching rows, then embed in chunks.
+        rows = scan.select(
+            ["paper_id", "title", "primary_category", "first_version_date", "arxiv_abs_url", "abstract"]
+        ).collect().to_dicts()
+
+        for r in rows:
+            rec = dict(r)
+            fvd = rec["first_version_date"]
+            year = fvd.year if isinstance(fvd, _dt.datetime) else (int(fvd) if fvd is not None else 0)
+            pid = str(rec["paper_id"])
+            if pid in self._existing:
+                continue
+            text = f"{rec['title']} {rec['abstract']}"
+            batch.append(text)
+            batch_ids.append(pid)
+            batch_years.append(year)
+            batch_cats.append(str(rec["primary_category"] or ""))
+            batch_urls.append(str(rec["arxiv_abs_url"] or ""))
+            batch_titles.append(str(rec["title"] or ""))
+            batch_abstracts.append(str(rec["abstract"] or ""))
+            if len(batch) >= self.chunk_size:
+                flush()
+        flush()
+        return added
+
+    # -- persistence --------------------------------------------------------
+    def save(self, path: str | None = None) -> IndexMeta:
+        path = path or self.default_path()
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        faiss.write_index(self.index, path)
+        meta_pkl = path.replace(".index", ".meta.pkl")
+        with open(meta_pkl, "wb") as fh:
+            pickle.dump(
+                (self._ids, self._years, self._cats, self._urls, self._titles, self._abstracts),
+                fh,
+            )
+        meta = IndexMeta(
+            embedder=type(self.embedder).__name__,
+            dim=self.dim,
+            source=os.environ.get("ARXIV_METADATA_PARQUET", "arxiv-complete/metadata"),
+            count=len(self._ids),
+        )
+        with open(path.replace(".index", ".meta.json"), "w") as fh:
+            json.dump(meta.__dict__, fh, indent=2)
+        return meta
+
+
+def load_index(name: str, embedder, root: str | None = None) -> tuple["AnnIndex", IndexMeta]:
+    """Load a previously-built index plus its embedder for querying."""
+    _require_faiss()
+    root = root or os.path.join(os.path.dirname(__file__), "index")
+    path = os.path.join(root, f"{name}.index")
+    builder = IndexBuilder(name, embedder)
+    builder.load(path)
+    meta_path = path.replace(".index", ".meta.json")
+    with open(meta_path) as fh:
+        meta = IndexMeta(**json.load(fh))
+    return AnnIndex(builder, meta), meta

retrieval/query.py

@@ -0,0 +1,207 @@
+"""
+Query interface for the retrieval layer.
+
+``query()`` is the single entry point the CLI (and any caller) uses.  It:
+
+  1. embeds the natural-language query,
+  2. runs an ANN search against the FAISS index,
+  3. attaches title / year / category / URL / abstract snippet to each hit,
+  4. applies the requested filters (year range, primary category, survey-only),
+  5. returns a ranked list of ``Result`` objects.
+
+If no index exists on disk yet (or building it failed), ``query()`` falls back
+to a **pure-keyword** search that streams the metadata and matches on title/
+abstract with a TF-style score.  This guarantees the CLI always returns
+something, even before an embedding index has been built.
+"""
+
+from __future__ import annotations
+
+import datetime
+import re
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from data import DataStore, SURVEY_KEYWORDS
+
+# Cap the working set for the keyword fallback so it stays fast on the 3.15M-row
+# metadata file (term-overlap scoring doesn't need the full corpus).
+_KEYWORD_WORKING_SET = 200_000
+
+
+@dataclass
+class Result:
+    paper_id: str
+    title: str
+    year: int
+    category: str
+    url: str
+    snippet: str
+    score: float
+    method: str  # "semantic" | "keyword"
+
+
+@dataclass
+class QueryOptions:
+    top_k: int = 10
+    year_min: int | None = None
+    year_max: int | None = None
+    primary_category: str | None = None
+    survey_only: bool = False
+    embedder = None  # set by caller; optional
+
+
+def _snippet(abstract: str, maxlen: int = 200) -> str:
+    text = re.sub(r"\s+", " ", abstract or "").strip()
+    if len(text) <= maxlen:
+        return text
+    return text[:maxlen].rsplit(" ", 1)[0] + " …"
+
+
+def _normalize(text: str) -> str:
+    return re.sub(r"[^a-z0-9]+", " ", (text or "").lower()).strip()
+
+
+def _keyword_score(text: str, terms: list[str]) -> float:
+    """TF-style overlap score between a doc and the query terms."""
+    words = set(text.split())
+    score = sum(1 for t in terms if t in words)
+    # small bonus for title matches
+    return score
+
+
+def _keyword_search(
+    query: str,
+    store: DataStore,
+    opts: QueryOptions,
+) -> list[Result]:
+    """
+    Pure-keyword fallback: stream matching rows, score by term overlap, filter,
+    and return the top-k.  Used when no ANN index is available.
+    """
+    terms = [t for t in _normalize(query).split() if len(t) > 1]
+    if not terms:
+        return []
+
+    # Pull a working set of rows (push the cheap filters into the scan).
+    scan = store.scan_filtered(
+        year_min=opts.year_min,
+        year_max=opts.year_max,
+        primary_category=opts.primary_category,
+        survey_only=opts.survey_only,
+    )
+    rows = (
+        scan.select(
+            ["paper_id", "title", "primary_category", "first_version_date", "arxiv_abs_url", "abstract"]
+        )
+        .limit(_KEYWORD_WORKING_SET)
+        .collect()
+        .to_dicts()
+    )
+    if not rows:
+        return []
+
+    scored: list[Result] = []
+    for r in rows:
+        rec = dict(r)
+        title = str(rec["title"] or "")
+        abstract = str(rec["abstract"] or "")
+        body = f"{title} {abstract}"
+        score = _keyword_score(body.lower(), terms)
+        if score <= 0:
+            continue
+        fvd = rec["first_version_date"]
+        year = fvd.year if isinstance(fvd, __import__("datetime").datetime) else (int(fvd) if fvd is not None else 0)
+        scored.append(
+            Result(
+                paper_id=str(rec["paper_id"]),
+                title=title,
+                year=year,
+                category=str(rec["primary_category"] or ""),
+                url=str(rec["arxiv_abs_url"] or ""),
+                snippet=_snippet(abstract),
+                score=score,
+                method="keyword",
+            )
+        )
+    scored.sort(key=lambda x: x.score, reverse=True)
+    return scored[: opts.top_k]
+
+
+def query(
+    query_text: str,
+    store: DataStore,
+    ann=None,
+    embedder=None,
+    opts: QueryOptions | None = None,
+) -> list[Result]:
+    """
+    Run a search.  Tries the ANN index first; falls back to keyword search.
+
+    Returns a list of ranked ``Result`` objects.
+    """
+    opts = opts or QueryOptions()
+    query_text = query_text or ""
+
+    # --- semantic path ---
+    if ann is not None and embedder is not None:
+        try:
+            return _semantic_query(query_text, store, ann, embedder, opts)
+        except Exception as exc:  # pragma: no cover - defensive
+            print(f"[query] ANN search failed ({exc}); falling back to keyword search.")
+
+    # --- keyword fallback ---
+    return _keyword_search(query_text, store, opts)
+
+
+def _semantic_query(
+    query_text: str,
+    store: DataStore,
+    ann,
+    embedder,
+    opts: QueryOptions,
+) -> list[Result]:
+    import faiss
+
+    qvec = np.ascontiguousarray(embedder.embed_one(query_text), dtype=np.float32)
+    k = min(opts.top_k * 8, len(ann), ann.index.ntotal)
+    k = max(k, 1)
+    dists, idx = ann.index.search(qvec.reshape(1, -1), k)
+    dists = dists[0]
+    idx = idx[0]
+
+    # Gather candidate hits and apply filters.
+    candidates: list[Result] = []
+    for d, i in zip(dists.tolist(), idx.tolist()):
+        if i < 0 or i >= len(ann._ids):
+            continue
+        pid = ann._ids[i]
+        year = ann._years[i]
+        cat = ann._cats[i]
+        if opts.year_min is not None and year < opts.year_min:
+            continue
+        if opts.year_max is not None and year > opts.year_max:
+            continue
+        if opts.primary_category and cat != opts.primary_category:
+            continue
+        if opts.survey_only:
+            if not any(kw in ann._abstracts[i].lower() for kw in SURVEY_KEYWORDS):
+                continue
+        # FAISS L2 distance -> similarity in [0,1]; closer => higher.
+        sim = 1.0 / (1.0 + float(d))
+        candidates.append(
+            Result(
+                paper_id=pid,
+                title=ann._titles[i],
+                year=year,
+                category=cat,
+                url=ann._urls[i],
+                snippet=_snippet(ann._abstracts[i]),
+                score=sim,
+                method="semantic",
+            )
+        )
+
+    candidates.sort(key=lambda x: x.score, reverse=True)
+    return candidates[: opts.top_k]

retrieval/requirements.txt

@@ -0,0 +1,18 @@
+# arxiv-complete retrieval layer — dependencies
+#
+# Create the environment with uv (recommended) or pip:
+#
+#   uv venv && uv pip install -r requirements.txt
+#   # or
+#   python -m venv .venv && .venv/bin/pip install -r requirements.txt
+#
+# faiss-cpu is required for the ANN path.  The pure-keyword fallback (build
+# without an embedding index) only needs polars + numpy.
+
+numpy
+polars
+faiss-cpu
+
+# Only needed for the remote (production) embedding provider.
+# Comment out if you only use the local HashingEmbedder fallback.
+requests

retrieval/search.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+"""
+search.py -- command-line interface for the arxiv-complete retrieval layer.
+
+Examples
+--------
+    # semantic search (uses the local hashing embedder unless an API key is set)
+    python search.py --query "speculative decoding"
+
+    # restrict to recent papers in a category
+    python search.py --query "neural machine translation" --top-k 15 --year-min 2018 --category cs.CL
+
+    # only survey / overview / taxonomy papers
+    python search.py --query "graph neural networks" --survey-only
+
+    # force the deterministic keyword fallback (no embeddings at all)
+    python search.py --query "speculative decoding" --keyword-only
+
+Build the ANN index first (once) with the build script:
+
+    python build_index.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+# Allow running as a script from the repo root or as a module.
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data import DataStore  # noqa: E402
+from embed import HashingEmbedder, get_embedder  # noqa: E402
+from query import QueryOptions, Result, query  # noqa: E402
+
+
+def _print_results(results: list[Result], query_text: str, method_note: str) -> None:
+    print(f"\nQuery: {query_text!r}")
+    print(f"Index: {method_note}  ({len(results)} result(s))\n")
+    if not results:
+        print("  (no matches)")
+        return
+    for i, r in enumerate(results, 1):
+        print(f"{i:>2}. [{r.method}] {r.score:.3f}  {r.title}")
+        meta = f"   {r.category} · {r.year} · {r.paper_id}"
+        print(meta)
+        print(f"   {r.url}")
+        print(f"   {_snippet_lines(r.snippet)}")
+        print()
+
+
+def _snippet_lines(snippet: str) -> str:
+    snippet = snippet.strip()
+    if not snippet:
+        snippet = "(no abstract available)"
+    width = 88
+    if len(snippet) <= width:
+        return "   " + snippet
+    return "   " + snippet[: width - 3] + "…"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="search.py",
+        description="Search the arxiv-complete corpus (semantic + keyword).",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument("--query", "-q", required=True, help="Natural-language query.")
+    p.add_argument("--top-k", type=int, default=10, help="Max results to return.")
+    p.add_argument("--year-min", type=int, default=None, help="Only papers from this year onward.")
+    p.add_argument("--year-max", type=int, default=None, help="Only papers up to this year.")
+    p.add_argument("--category", default=None, help="Filter by primary_category, e.g. cs.LG.")
+    p.add_argument("--survey-only", action="store_true", help="Only survey/overview/taxonomy papers.")
+    p.add_argument(
+        "--keyword-only",
+        action="store_true",
+        help="Skip the ANN index entirely; use the pure-keyword fallback.",
+    )
+    p.add_argument(
+        "--embedder",
+        default=None,
+        choices=["hashing", "remote"],
+        help="Embedding provider. Default: auto (remote if key set, else local hashing).",
+    )
+    p.add_argument(
+        "--index-name",
+        default="arxiv",
+        help="Name of the persisted FAISS index to use/load.",
+    )
+    p.add_argument(
+        "--parquet",
+        default=None,
+        help="Local parquet path or HF URL overriding the default metadata source.",
+    )
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+
+    store = DataStore(source=args.parquet if args.parquet else os.environ.get(
+        "ARXIV_METADATA_PARQUET", DataStore().source
+    ))
+
+    opts = QueryOptions(
+        top_k=args.top_k,
+        year_min=args.year_min,
+        year_max=args.year_max,
+        primary_category=args.category,
+        survey_only=args.survey_only,
+    )
+
+    if args.keyword_only:
+        results = _keyword_search_only(store, args, opts)
+        _print_results(results, args.query, "keyword fallback (no embeddings)")
+        return 0
+
+    # Semantic path: try to load an existing index; if absent, build it.
+    from index import IndexBuilder, load_index  # imported late to keep CLI fast
+
+    embedder = get_embedder(args.embedder)
+    builder = IndexBuilder(args.index_name, embedder)
+    index_path = builder.default_path()
+
+    if os.path.exists(index_path):
+        try:
+            ann, meta = load_index(args.index_name, embedder)
+            note = f"ANN index '{args.index_name}' ({meta.count} papers, {meta.embedder})"
+        except Exception as exc:
+            print(f"[query] failed to load index ({exc}); building fresh.")
+            ann = None
+    else:
+        print(f"[build] no index at {index_path}; building now (this embeds the corpus).")
+        added = builder.build(store.scan_filtered(), store)
+        meta = builder.save()
+        ann, meta = load_index(args.index_name, embedder)
+        note = f"freshly built ANN index ({meta.count} papers, {meta.embedder})"
+
+    results = query(args.query, store, ann=ann, embedder=embedder, opts=opts)
+    _print_results(results, args.query, note)
+    return 0
+
+
+def _keyword_search_only(store, args, opts) -> list[Result]:
+    from query import _keyword_search
+
+    return _keyword_search(args.query, store, opts)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())