diff --git a/retrieval/README.md b/retrieval/README.md index 28ca65e..18b8c20 100644 --- a/retrieval/README.md +++ b/retrieval/README.md @@ -13,7 +13,7 @@ because keyword-only search misses synonyms. | 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). | +| `embed.py` | Pluggable embedding providers: an OpenAI-compatible **remote** provider (production), a **local semantic** sentence-transformers model (default local path), and a dependency-free **local hashing** fallback (offline, no deps). | | `index.py` | Builds and persists a FAISS flat L2 index + metadata sidecar (paper_id → vector, title, year, category, URL, abstract). Incremental adds. | | `query.py` | The interface. Embeds the query, runs ANN search, ranks, applies filters. Falls back to keyword search if no index is available. | | `search.py` | CLI tying it together. | @@ -72,11 +72,11 @@ python search.py --query "speculative decoding" --keyword-only python search.py --query "speculative decoding" --survey-only --category cs.LG ``` -## Embeddings: remote (production) vs local fallback +## Embeddings: remote (production) vs local semantic vs hashing -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. +This machine (Strix Halo, ~104 GiB unified RAM) has **OOM-killed large local +models before**, so we deliberately avoid loading big ones. Three options, all +behind the same `Embedder` interface: **Production path** — set these env vars and the CLI uses a remote OpenAI-compatible embeddings endpoint: @@ -90,30 +90,21 @@ 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.) +**Local semantic (default)** — with no API key configured, the code +automatically uses `LocalEmbedder`, a small CPU +[sentence-transformers](https://huggingface.co/sentence-transformers) model +(`paraphrase-MiniLM-L3-v2`, ~90 MB). This produces *genuinely semantic* vectors +(synonyms and paraphrases map close together) at no API cost, and fits comfortably +in RAM. It's slower than keyword hashing — embedding the full 3.15M corpus takes +~1-2 hours on this box, so for a demo index a single `--category` subset instead. +Override the model with `EMBEDDING_LOCAL_MODEL` (e.g. `all-MiniLM-L6-v2` for +higher quality at ~3x the time). -### 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. +**Dependency-free fallback** — pass `--embedder hashing` (or `--keyword-only`) +for `HashingEmbedder`, a deterministic embedding that projects a bag of character +n-grams into a fixed-size vector via the hashing trick. No weights, no GPU, no +internet, no API key — handy for CI or a smoke test, but not semantically +meaningful (synonyms aren't mapped together). ## How the index works diff --git a/retrieval/requirements.txt b/retrieval/requirements.txt index 6b39865..d2d5609 100644 --- a/retrieval/requirements.txt +++ b/retrieval/requirements.txt @@ -13,6 +13,10 @@ numpy polars faiss-cpu +# Needed for the local semantic embedder (LocalEmbedder, the default local path). +# Uses a small CPU sentence-transformers model (~90 MB, no GPU needed). +sentence-transformers + # Only needed for the remote (production) embedding provider. -# Comment out if you only use the local HashingEmbedder fallback. +# Comment out if you only use the local embedders. requests