From bf475d5e1e470a2e298a7a202f9f14e9c6a07211 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sun, 20 Sep 2026 05:49:06 -0700 Subject: [PATCH] Add arxiv-complete retrieval layer (semantic + keyword) --- retrieval/.gitignore | 11 ++ retrieval/README.md | 140 ++++++++++++++++++++++ retrieval/build_index.py | 64 ++++++++++ retrieval/data.py | 136 ++++++++++++++++++++++ retrieval/embed.py | 208 +++++++++++++++++++++++++++++++++ retrieval/index.py | 233 +++++++++++++++++++++++++++++++++++++ retrieval/query.py | 207 ++++++++++++++++++++++++++++++++ retrieval/requirements.txt | 18 +++ retrieval/search.py | 152 ++++++++++++++++++++++++ 9 files changed, 1169 insertions(+) create mode 100644 retrieval/.gitignore create mode 100644 retrieval/README.md create mode 100644 retrieval/build_index.py create mode 100644 retrieval/data.py create mode 100644 retrieval/embed.py create mode 100644 retrieval/index.py create mode 100644 retrieval/query.py create mode 100644 retrieval/requirements.txt create mode 100644 retrieval/search.py diff --git a/retrieval/.gitignore b/retrieval/.gitignore new file mode 100644 index 0000000..8aa40e8 --- /dev/null +++ b/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 diff --git a/retrieval/README.md b/retrieval/README.md new file mode 100644 index 0000000..28ca65e --- /dev/null +++ b/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` (FAISS), `.meta.pkl` (parallel + arrays: id, year, category, url, title, abstract), and `.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. diff --git a/retrieval/build_index.py b/retrieval/build_index.py new file mode 100644 index 0000000..a73f7ab --- /dev/null +++ b/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()) diff --git a/retrieval/data.py b/retrieval/data.py new file mode 100644 index 0000000..602d1ef --- /dev/null +++ b/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]) diff --git a/retrieval/embed.py b/retrieval/embed.py new file mode 100644 index 0000000..2404849 --- /dev/null +++ b/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() diff --git a/retrieval/index.py b/retrieval/index.py new file mode 100644 index 0000000..77c7d49 --- /dev/null +++ b/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: + + * ``.index`` -- the FAISS flat L2 index (float32 vectors) + * ``.meta.pkl`` -- parallel arrays: paper_id, year, category, url, + title, abstract (aligned to the FAISS rows) + * ``.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 diff --git a/retrieval/query.py b/retrieval/query.py new file mode 100644 index 0000000..47e1e95 --- /dev/null +++ b/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] diff --git a/retrieval/requirements.txt b/retrieval/requirements.txt new file mode 100644 index 0000000..6b39865 --- /dev/null +++ b/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 diff --git a/retrieval/search.py b/retrieval/search.py new file mode 100644 index 0000000..9504a19 --- /dev/null +++ b/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())