From 5dff8194b7c221196138a18d2b81b354678a6b90 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 11:43:49 +0200 Subject: [PATCH 1/5] docs: add executable examples + CI smoke test Six stand-alone example scripts that exercise the public API end to end: quickstart, pq_index, ivf_pq (with fp16 rerank), filter_search, save_load, and streaming_ingest. Each example is tiny (synthetic data, a few seconds) and runs as part of the CI matrix on every OS + Python combination, so a broken public API shows up before release. --- .github/workflows/ci.yml | 14 +++++++++ examples/filter_search.py | 45 +++++++++++++++++++++++++++++ examples/ivf_pq.py | 56 ++++++++++++++++++++++++++++++++++++ examples/pq_index.py | 37 ++++++++++++++++++++++++ examples/quickstart.py | 35 ++++++++++++++++++++++ examples/save_load.py | 45 +++++++++++++++++++++++++++++ examples/streaming_ingest.py | 37 ++++++++++++++++++++++++ 7 files changed, 269 insertions(+) create mode 100644 examples/filter_search.py create mode 100644 examples/ivf_pq.py create mode 100644 examples/pq_index.py create mode 100644 examples/quickstart.py create mode 100644 examples/save_load.py create mode 100644 examples/streaming_ingest.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a375b4c..4be8a96 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -65,6 +65,20 @@ jobs: - name: Run tests run: pytest -q --cov=snapvec --cov-report=term-missing + - name: Run examples smoke test + shell: bash + run: | + set -e + for ex in examples/quickstart.py \ + examples/pq_index.py \ + examples/ivf_pq.py \ + examples/filter_search.py \ + examples/save_load.py \ + examples/streaming_ingest.py; do + echo "=== $ex ===" + python "$ex" + done + - name: Upload coverage (ubuntu + py3.12 only) if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12' uses: actions/upload-artifact@v4 diff --git a/examples/filter_search.py b/examples/filter_search.py new file mode 100644 index 0000000..03abed4 --- /dev/null +++ b/examples/filter_search.py @@ -0,0 +1,45 @@ +"""Filtered search: restrict results to a subset of ids. + +Every index type accepts ``filter_ids=`` to limit the candidates +returned. On IVFPQSnapIndex the filter is cluster-aware: probes skip +clusters that contain no matching id. + +Run with: python examples/filter_search.py +""" +from __future__ import annotations + +import numpy as np + +from snapvec import SnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, n_corpus = 64, 500 + + corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) + ids = [f"doc-{i:04d}" for i in range(n_corpus)] + + idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx.add_batch(ids, corpus) + + query = rng.standard_normal(dim).astype(np.float32) + + unrestricted = idx.search(query, k=5) + restricted_set = {f"doc-{i:04d}" for i in range(100)} + restricted = idx.search(query, k=5, filter_ids=restricted_set) + + print("unrestricted top-5:") + for doc_id, score in unrestricted: + print(f" {doc_id} score={score:+.4f}") + + print("\nrestricted to doc-0000..doc-0099:") + for doc_id, score in restricted: + print(f" {doc_id} score={score:+.4f}") + + assert all(h[0] in restricted_set for h in restricted) + print("\nall restricted hits are in the filter set: OK") + + +if __name__ == "__main__": + main() diff --git a/examples/ivf_pq.py b/examples/ivf_pq.py new file mode 100644 index 0000000..a90984d --- /dev/null +++ b/examples/ivf_pq.py @@ -0,0 +1,56 @@ +"""IVFPQSnapIndex with float16 rerank. + +IVF-PQ partitions the corpus into ``nlist`` coarse clusters and only +scans ``nprobe`` of them per query (sub-linear search). With +``keep_full_precision=True`` + ``rerank_candidates``, the top candidates +are rescored against the stored float16 vectors, which recovers almost +all the recall lost to PQ compression. + +Run with: python examples/ivf_pq.py +""" +from __future__ import annotations + +import numpy as np + +from snapvec import IVFPQSnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, n_corpus, n_queries = 128, 5000, 20 + + corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) + queries = rng.standard_normal((n_queries, dim)).astype(np.float32) + + idx = IVFPQSnapIndex( + dim=dim, + nlist=64, + M=16, + K=256, + normalized=True, + keep_full_precision=True, + seed=0, + ) + idx.fit(corpus[:2500]) + idx.add_batch(list(range(n_corpus)), corpus) + + print(f"IVFPQSnapIndex: n={len(idx)}, nlist=64, M=16") + + truth = (queries @ corpus.T).argmax(axis=1) + + # Baseline: IVF-PQ only. + pq_hits = [idx.search(q, k=1, nprobe=8)[0][0] for q in queries] + pq_recall = sum(h == t for h, t in zip(pq_hits, truth)) / n_queries + + # With rerank: candidates pre-selected by PQ, rescored in fp16. + rerank_hits = [ + idx.search(q, k=1, nprobe=8, rerank_candidates=50)[0][0] for q in queries + ] + rerank_recall = sum(h == t for h, t in zip(rerank_hits, truth)) / n_queries + + print(f"top-1 recall (PQ only): {pq_recall:.2f}") + print(f"top-1 recall (PQ + fp16 rerank): {rerank_recall:.2f}") + + +if __name__ == "__main__": + main() diff --git a/examples/pq_index.py b/examples/pq_index.py new file mode 100644 index 0000000..9a7d98e --- /dev/null +++ b/examples/pq_index.py @@ -0,0 +1,37 @@ +"""PQSnapIndex example. + +PQSnapIndex learns per-subspace k-means codebooks from a training sample, +then encodes each vector as M bytes. Much higher recall than scalar +quantization at the same bytes/vec, but requires a one-time ``fit`` call. + +Run with: python examples/pq_index.py +""" +from __future__ import annotations + +import numpy as np + +from snapvec import PQSnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, n_corpus, n_queries = 128, 2000, 10 + + corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) + queries = rng.standard_normal((n_queries, dim)).astype(np.float32) + + idx = PQSnapIndex(dim=dim, M=16, K=256, normalized=True, seed=0) + idx.fit(corpus[:1000]) # train codebooks on first half + idx.add_batch(list(range(n_corpus)), corpus) + + print(f"PQSnapIndex: n={len(idx)}, dim={dim}, M=16, K=256") + print(f"bytes/vec: {idx.M} (vs {dim * 4} for float32)") + + truth = (queries @ corpus.T).argmax(axis=1) + hits = [idx.search(q, k=1)[0][0] for q in queries] + recall_at_1 = sum(h == t for h, t in zip(hits, truth)) / n_queries + print(f"top-1 recall over {n_queries} queries: {recall_at_1:.2f}") + + +if __name__ == "__main__": + main() diff --git a/examples/quickstart.py b/examples/quickstart.py new file mode 100644 index 0000000..69dc259 --- /dev/null +++ b/examples/quickstart.py @@ -0,0 +1,35 @@ +"""Minimal SnapIndex example. + +Build a 4-bit scalar-quantized index over 1,000 random vectors and +recover the exact top-1 for a handful of queries without any training. + +Run with: python examples/quickstart.py +""" +from __future__ import annotations + +import numpy as np + +from snapvec import SnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, n_corpus, n_queries = 128, 1000, 5 + + corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) + queries = rng.standard_normal((n_queries, dim)).astype(np.float32) + + idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx.add_batch(list(range(n_corpus)), corpus) + + print(f"SnapIndex: n={len(idx)}, dim={dim}, bits=4") + + truth = (queries @ corpus.T).argmax(axis=1) + hits = [idx.search(q, k=1)[0][0] for q in queries] + + recall_at_1 = sum(h == t for h, t in zip(hits, truth)) / n_queries + print(f"top-1 recall over {n_queries} queries: {recall_at_1:.2f}") + + +if __name__ == "__main__": + main() diff --git a/examples/save_load.py b/examples/save_load.py new file mode 100644 index 0000000..75a4c52 --- /dev/null +++ b/examples/save_load.py @@ -0,0 +1,45 @@ +"""Save and load an index round-trip. + +Every index type has ``.save(path)`` / ``.load(path)``. Writes are +atomic (write to ``.tmp`` then rename) and include a CRC32 trailer +so transport or disk corruption is caught at load time. + +Run with: python examples/save_load.py +""" +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np + +from snapvec import SnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, n_corpus = 128, 500 + + corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) + + idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx.add_batch(list(range(n_corpus)), corpus) + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "demo.snpv" + idx.save(path) + size_kb = path.stat().st_size / 1024 + print(f"wrote {path.name} ({size_kb:.1f} KB)") + + loaded = SnapIndex.load(path) + print(f"loaded index: n={len(loaded)}, dim={loaded.dim}, bits={loaded.bits}") + + query = corpus[42] + hit_before = idx.search(query, k=1)[0] + hit_after = loaded.search(query, k=1)[0] + assert hit_before == hit_after, "round-trip mismatch" + print(f"round-trip top-1 identical: id={hit_after[0]} score={hit_after[1]:+.4f}") + + +if __name__ == "__main__": + main() diff --git a/examples/streaming_ingest.py b/examples/streaming_ingest.py new file mode 100644 index 0000000..c7228ea --- /dev/null +++ b/examples/streaming_ingest.py @@ -0,0 +1,37 @@ +"""Streaming ingest: grow an index with many small batches. + +SnapIndex is training-free so you can start searching after the first +``add_batch`` call and keep appending. PQ-based indices need an initial +``fit(sample)`` before adding vectors. + +Run with: python examples/streaming_ingest.py +""" +from __future__ import annotations + +import numpy as np + +from snapvec import SnapIndex + + +def main() -> None: + rng = np.random.default_rng(0) + dim, batch_size, n_batches = 64, 100, 10 + + idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + + for batch in range(n_batches): + vecs = rng.standard_normal((batch_size, dim)).astype(np.float32) + base = batch * batch_size + ids = list(range(base, base + batch_size)) + idx.add_batch(ids, vecs) + print(f"after batch {batch + 1:2d}: n={len(idx):5d}") + + assert len(idx) == batch_size * n_batches + + query = rng.standard_normal(dim).astype(np.float32) + hits = idx.search(query, k=3) + print(f"\ntop-3 for a fresh query: {[h[0] for h in hits]}") + + +if __name__ == "__main__": + main() From d3eba456a838a640feca320faacf09511b98c550 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 11:44:00 +0200 Subject: [PATCH 2/5] docs: add MkDocs site with Material theme + API reference - mkdocs.yml: Material theme with dark/light toggle, mkdocstrings for auto-generated API reference, navigation tree with getting-started, user-guide, architecture, benchmarks, API, changelog. - docs/: 16 pages. User guide has one page per index type plus 'choosing an index' decision tree, save/load, and filtered search. Architecture explains RHT, Lloyd-Max, PQ, IVF, and fp16 rerank. Benchmarks page carries the FIQA and sqlite-vec comparison numbers. API reference is five short pages that mkdocstrings expands from the existing docstrings -- zero extra maintenance burden. - docs/changelog.md uses mkdocs-material's snippet include so the canonical CHANGELOG.md stays at the repo root. - docs/blog/ (pre-existing drafts) is excluded from the built site. - pyproject.toml: new [docs] optional extra pulls mkdocs-material and mkdocstrings. Regular dev install unaffected. - .github/workflows/docs.yml: strict build on every push and PR, deploy to GitHub Pages on main. PR builds validate without deploying. - .gitignore: exclude the site/ build output. --- .github/workflows/docs.yml | 56 ++++++++++++++ .gitignore | 5 +- docs/api/helpers.md | 16 ++++ docs/api/ivf-pq.md | 3 + docs/api/pq.md | 3 + docs/api/residual.md | 3 + docs/api/snap-index.md | 3 + docs/architecture.md | 105 +++++++++++++++++++++++++++ docs/benchmarks.md | 59 +++++++++++++++ docs/changelog.md | 1 + docs/getting-started/installation.md | 52 +++++++++++++ docs/getting-started/quickstart.md | 51 +++++++++++++ docs/index.md | 50 +++++++++++++ docs/user-guide/choosing-an-index.md | 37 ++++++++++ docs/user-guide/filter-search.md | 29 ++++++++ docs/user-guide/ivf-pq.md | 76 +++++++++++++++++++ docs/user-guide/pq.md | 52 +++++++++++++ docs/user-guide/residual.md | 36 +++++++++ docs/user-guide/save-load.md | 39 ++++++++++ docs/user-guide/snap-index.md | 55 ++++++++++++++ mkdocs.yml | 91 +++++++++++++++++++++++ pyproject.toml | 5 ++ 22 files changed, 826 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/docs.yml create mode 100644 docs/api/helpers.md create mode 100644 docs/api/ivf-pq.md create mode 100644 docs/api/pq.md create mode 100644 docs/api/residual.md create mode 100644 docs/api/snap-index.md create mode 100644 docs/architecture.md create mode 100644 docs/benchmarks.md create mode 100644 docs/changelog.md create mode 100644 docs/getting-started/installation.md create mode 100644 docs/getting-started/quickstart.md create mode 100644 docs/index.md create mode 100644 docs/user-guide/choosing-an-index.md create mode 100644 docs/user-guide/filter-search.md create mode 100644 docs/user-guide/ivf-pq.md create mode 100644 docs/user-guide/pq.md create mode 100644 docs/user-guide/residual.md create mode 100644 docs/user-guide/save-load.md create mode 100644 docs/user-guide/snap-index.md create mode 100644 mkdocs.yml diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..263ff51 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,56 @@ +name: Docs + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: docs-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + pages: write + id-token: write + +jobs: + build: + name: Build docs + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - name: Install package + docs dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[docs]" + + - name: Build site (strict) + run: mkdocs build --strict + + - name: Upload Pages artifact + if: github.ref == 'refs/heads/main' + uses: actions/upload-pages-artifact@v3 + with: + path: site + + deploy: + name: Deploy to GitHub Pages + needs: build + if: github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.gitignore b/.gitignore index 87b7fb0..3627086 100644 --- a/.gitignore +++ b/.gitignore @@ -14,9 +14,12 @@ dist/ .pytest_cache/ .ruff_cache/ -# experiment caches and harness state — large binary, not source of truth +# experiment caches and harness state -- large binary, not source of truth experiments/.cache_*.npy .claude/scheduled_tasks.lock *.so snapvec/_fast.c build/ + +# mkdocs build output +site/ diff --git a/docs/api/helpers.md b/docs/api/helpers.md new file mode 100644 index 0000000..9faf609 --- /dev/null +++ b/docs/api/helpers.md @@ -0,0 +1,16 @@ +# Helpers + +Lower-level building blocks exposed for advanced use. Most users should +reach for one of the index classes. + +## `rht` + +::: snapvec.rht + +## `padded_dim` + +::: snapvec.padded_dim + +## `get_codebook` + +::: snapvec.get_codebook diff --git a/docs/api/ivf-pq.md b/docs/api/ivf-pq.md new file mode 100644 index 0000000..901b1d8 --- /dev/null +++ b/docs/api/ivf-pq.md @@ -0,0 +1,3 @@ +# IVFPQSnapIndex + +::: snapvec.IVFPQSnapIndex diff --git a/docs/api/pq.md b/docs/api/pq.md new file mode 100644 index 0000000..81469c1 --- /dev/null +++ b/docs/api/pq.md @@ -0,0 +1,3 @@ +# PQSnapIndex + +::: snapvec.PQSnapIndex diff --git a/docs/api/residual.md b/docs/api/residual.md new file mode 100644 index 0000000..1f7021c --- /dev/null +++ b/docs/api/residual.md @@ -0,0 +1,3 @@ +# ResidualSnapIndex + +::: snapvec.ResidualSnapIndex diff --git a/docs/api/snap-index.md b/docs/api/snap-index.md new file mode 100644 index 0000000..02cadb6 --- /dev/null +++ b/docs/api/snap-index.md @@ -0,0 +1,3 @@ +# SnapIndex + +::: snapvec.SnapIndex diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..d6f778a --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,105 @@ +# Architecture + +How `snapvec` compresses vectors without destroying inner-product recall. + +## Motivation + +Dense embeddings from modern encoders (384-dim BGE, 768-dim E5, 1024-dim+ +OpenAI, etc.) are expensive to store and scan at scale. At float32, +1M vectors of dim=768 take 3 GB. Float16 halves that but still scales +linearly with dim and only halves latency because memory bandwidth is +the bottleneck. + +`snapvec` targets the compression / recall frontier: give up a small +amount of recall (typically 1-5 percentage points) for 6-96x less disk +and RAM, and run search kernels that fit in L2 cache. + +## Pipeline (SnapIndex) + +1. **Normalize** (optional): unit-length every vector. Turns inner + product into cosine similarity and removes one nuisance dimension. +2. **Randomized Hadamard Transform (RHT)**: multiply by a signed random + Hadamard matrix. Decorrelates coordinates so their marginal + distribution approaches iid Gaussian, even when the input is not. +3. **Lloyd-Max scalar quantization**: quantize each coordinate to `bits` + levels using the Lloyd-Max boundaries for a unit Gaussian. This is + the optimal scalar quantizer for a Gaussian source under MSE. +4. **Bit-pack** the quantized codes so storage is exactly `bits * pdim / 8` + bytes per vector (`pdim` is the Hadamard-padded dim). + +At query time, the query goes through the same normalize + RHT + quantize +path, then a vectorized inner-product sum over the unpacked codes scores +the whole corpus. + +## Why RHT? + +Scalar quantization is optimal for a Gaussian source. Real embeddings are +not Gaussian along each coordinate. RHT is an isometry that makes +coordinates look Gaussian without changing inner products, so Lloyd-Max +becomes near-optimal. It also distributes signal energy across +coordinates, so no single quantization error dominates. + +The RHT is implemented as a reshape-view + in-place xor across `log2(pdim)` +levels, so it costs `O(pdim * log(pdim))` time but `O(log(pdim))` Python +dispatch overhead. + +## Lloyd-Max codebooks + +Pre-computed per bit-depth: + +- 2-bit: 4 reconstruction levels at Gaussian-optimal boundaries +- 3-bit: 8 levels +- 4-bit: 16 levels + +Stored as constants in [`snapvec/_codebooks.py`](https://github.com/stffns/snapvec/blob/main/snapvec/_codebooks.py). + +## Product Quantization (PQSnapIndex) + +For higher recall at the same bytes/vec, `PQSnapIndex` replaces scalar +quantization with product quantization: + +1. Split each vector into `M` contiguous sub-vectors. +2. For each subspace, learn `K=256` k-means centroids over a training + sample. +3. Encode each vector as `M` bytes (one centroid index per subspace). + +Query scoring uses an ADC (asymmetric distance computation) lookup table: +pre-compute `M * K` partial inner products against the query, then sum +`M` table lookups per corpus vector. The Cython+OpenMP kernel +(`snapvec._fast.adc_colmajor`) runs this in parallel with scores as a +stripe of ~4-8 MB that fits in L2. + +## Inverted File (IVFPQSnapIndex) + +For sub-linear search, layer an inverted file (IVF) on top of PQ: + +1. Learn `nlist` coarse centroids from a training sample. +2. Assign each corpus vector to its nearest coarse centroid. +3. Store vectors contiguous per cluster, with an `offsets` array for + random access. +4. At query time, rank coarse centroids by query-centroid distance and + visit only the top `nprobe` clusters. + +Each cluster stores PQ codes of the **residual** (vector minus its coarse +centroid), which is easier to quantize than the raw vector. + +### Float16 rerank + +PQ has a recall ceiling set by codebook granularity. To break it, +`IVFPQSnapIndex` with `keep_full_precision=True` stores each vector in +float16 as well. At query time, the top `rerank_candidates` returned by +the PQ pass are rescored against their stored fp16 vectors, and the +final top-k comes from the rerank scores. + +This recovers 4-7 percentage points of recall at the cost of one +`(rerank_candidates, dim) @ (dim,)` matmul per query, typically under +1 ms even at large `nprobe`. + +## References + +- Zandieh et al. (2025). *TurboQuant: Online Vector Quantization with + Near-optimal Distortion Rate.* [arXiv:2504.19874](https://arxiv.org/abs/2504.19874) +- Jegou et al. (2011). *Product Quantization for Nearest Neighbor Search.* + IEEE TPAMI. +- Andre et al. (2015). *Cache Locality is Not Enough: High-performance + Nearest Neighbor Search with Product Quantization Fast Scan.* diff --git a/docs/benchmarks.md b/docs/benchmarks.md new file mode 100644 index 0000000..23b4d0a --- /dev/null +++ b/docs/benchmarks.md @@ -0,0 +1,59 @@ +# Benchmarks + +All numbers measured on the same hardware with a fixed random seed. Raw +scripts live in [`experiments/`](https://github.com/stffns/snapvec/tree/main/experiments); +a reproducible CI-runnable suite is planned. + +## Headline: IVFPQSnapIndex on FIQA + +BEIR FIQA, N = 57,638, dim = 384 (BGE-small), nlist = 512, M = 192, K = 256. +`keep_full_precision=True`, `rerank_candidates=100`. + +| `nprobe` | Configuration | recall@10 | Latency (us/query) | +|---------|---------------|-----------|--------------------| +| 8 | PQ only | 0.85 | 180 | +| 32 | PQ only | 0.92 | 340 | +| 64 | PQ + fp16 rerank | **0.977** | **441** | +| 256 | PQ + fp16 rerank | **0.998** | **1021** | + +5.8x speedup vs v0.6 at identical recall, past the PQ-only 0.929 ceiling. + +## snapvec vs sqlite-vec (measured) + +Same corpus, same hardware. `snapvec` at `nprobe=64` + rerank. + +| N | sqlite-vec | snapvec | Speedup | Recall tradeoff | +|---|------------|---------|---------|-----------------| +| 10k | 2.3 ms | 0.44 ms | 5x | 0.997 | +| 57k | 15.1 ms | 0.44 ms | 34x | 0.977 | +| 100k | 23.8 ms | 1.04 ms | 23x | 0.994 | +| 500k | ~110 ms | 0.9 ms | 125x | ~0.97 | +| 1M | brute-force infeasible | 1.1 ms | -- | -- | + +Disk footprint is 2-8x smaller across the range. + +## Compression ratios + +For BGE-small (dim=384, float32 baseline = 1536 B/vec): + +| Index | Config | B/vec | Compression | +|-------|--------|-------|-------------| +| `SnapIndex` | bits=2 | 132 | 11.6x | +| `SnapIndex` | bits=3 | 196 | 7.8x | +| `SnapIndex` | bits=4 | 260 | 5.9x | +| `PQSnapIndex` | M=16 | 16 | 96x | +| `PQSnapIndex` | M=32 | 32 | 48x | +| `IVFPQSnapIndex` | M=192 + fp16 rerank | ~960 | 1.6x (rerank cache dominates) | +| `IVFPQSnapIndex` | M=192, no rerank | ~192 | 8x | + +## Reproduction + +```bash +pip install -e ".[dev]" +python experiments/bench_v090_fiqa.py # FIQA recall / latency +python experiments/bench_sqlite_vec_baseline.py # sqlite-vec comparison +``` + +The `experiments/` folder is WIP; expect rough edges. A first-class +`bench/` suite that runs in CI and emits machine-readable results is +tracked on the roadmap. diff --git a/docs/changelog.md b/docs/changelog.md new file mode 100644 index 0000000..786b75d --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1 @@ +--8<-- "CHANGELOG.md" diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md new file mode 100644 index 0000000..e3056a6 --- /dev/null +++ b/docs/getting-started/installation.md @@ -0,0 +1,52 @@ +# Installation + +## From PyPI + +```bash +pip install snapvec +``` + +Pre-compiled wheels are published for CPython 3.10-3.13 on Linux (x86_64, +aarch64), macOS (x86_64 and arm64, macOS 13+), and Windows (AMD64). NumPy +>= 1.24 is the only runtime dependency. + +## From source + +```bash +git clone https://github.com/stffns/snapvec.git +cd snapvec +pip install -e ".[dev]" +``` + +On macOS you need OpenMP from Homebrew before building: + +```bash +brew install libomp +``` + +Without `libomp`, the Cython kernels compile in serial mode. The library +still works but the parallel search paths do not. + +## Input dtype + +`snapvec` expects `np.float32` inputs everywhere. Passing `np.float64` +(the NumPy default for `np.array([...])`) is an error because silently +downcasting would be a surprise on the hot path. + +```python +# Models that return float32 directly (most modern embeddings) +vecs = model.encode(texts) # already float32 + +# Arrays from np.array([...]) default to float64 -- cast explicitly +vecs = np.array(vecs_list, dtype=np.float32) + +# Loading from disk -- respect the original dtype or force float32 +vecs = np.load("embeddings.npy").astype(np.float32, copy=False) +``` + +## Verify install + +```python +import snapvec +print(snapvec.__version__) +``` diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md new file mode 100644 index 0000000..5a5a0b5 --- /dev/null +++ b/docs/getting-started/quickstart.md @@ -0,0 +1,51 @@ +# Quickstart + +Five-minute tour. Run [`examples/quickstart.py`](https://github.com/stffns/snapvec/blob/main/examples/quickstart.py) +end-to-end for a working script. + +## 1. Build an index + +`SnapIndex` is the simplest index: training-free scalar quantization on +top of the randomized Hadamard transform. No `fit` call needed. + +```python +import numpy as np +from snapvec import SnapIndex + +rng = np.random.default_rng(0) +corpus = rng.standard_normal((10_000, 384)).astype(np.float32) + +idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx.add_batch(list(range(10_000)), corpus) +``` + +`ids` can be any hashable (int, str, UUID); they are round-tripped through +save/load. + +## 2. Query + +```python +query = rng.standard_normal(384).astype(np.float32) +hits = idx.search(query, k=10) + +for doc_id, score in hits: + print(doc_id, score) +``` + +`search` returns `list[tuple[id, float]]` sorted by descending score. + +## 3. Persist + +```python +idx.save("my.snpv") +loaded = SnapIndex.load("my.snpv") +``` + +Writes are atomic (temp file + rename) and CRC32-checksummed. + +## Next steps + +- [Choosing an index](../user-guide/choosing-an-index.md) -- decision tree + across the four index types. +- [User guide](../user-guide/snap-index.md) -- deep dives per index. +- [Architecture](../architecture.md) -- how the compression works. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..6d436a8 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,50 @@ +# snapvec + +**Fast compressed approximate nearest-neighbor search. NumPy + Cython compiled kernels.** + +`snapvec` ships four index types for embedding vector search, each targeting +a different point on the accuracy / storage / latency frontier: + +| Index | Training | Compression | Recall | Use when | +|-------|----------|-------------|--------|----------| +| [`SnapIndex`](user-guide/snap-index.md) | none | 6-12x | 0.92+ | Any distribution, no corpus sample | +| [`ResidualSnapIndex`](user-guide/residual.md) | none | 4-8x | 0.96 | Higher recall, still training-free | +| [`PQSnapIndex`](user-guide/pq.md) | one-off `fit` | 24-96x | 0.95 | Modern LLM embeddings, aggressive compression | +| [`IVFPQSnapIndex`](user-guide/ivf-pq.md) | one-off `fit` | 24-96x | 0.98 | Sub-linear search at scale (N > 100k) | + +All four file formats (`.snpv` / `.snpq` / `.snpr` / `.snpi`) carry a CRC32 +trailer -- silent disk or transport corruption is caught at `load()` time +instead of returning wrong results. + +## Install + +```bash +pip install snapvec +``` + +## Quickstart + +```python +import numpy as np +from snapvec import SnapIndex + +idx = SnapIndex(dim=384, bits=4, normalized=True) +idx.add_batch(list(range(10_000)), np.random.randn(10_000, 384).astype(np.float32)) + +results = idx.search(np.random.randn(384).astype(np.float32), k=10) +for doc_id, score in results: + print(doc_id, score) +``` + +See [Quickstart](getting-started/quickstart.md) for the end-to-end tour and +[Choosing an index](user-guide/choosing-an-index.md) for the decision tree. + +## Context + +`snapvec` was developed as the quantization layer for +[vstash](https://github.com/stffns/vstash), a local-first hybrid retrieval +system, to extend it to corpora beyond the float32 memory budget while +preserving its dependency-minimal design. It stands alone as a quantization +library, but the design constraints (NumPy-only base install, predictable +latency, reproducible index files) come from vstash's local-first +requirements. diff --git a/docs/user-guide/choosing-an-index.md b/docs/user-guide/choosing-an-index.md new file mode 100644 index 0000000..9dd1eaa --- /dev/null +++ b/docs/user-guide/choosing-an-index.md @@ -0,0 +1,37 @@ +# Choosing an index + +| If you... | Pick | Why | +|-----------|------|-----| +| Just want a dependency-minimal ANN with no training | **`SnapIndex`** | Works on any distribution, no `fit` call. 4-bit default is a good balance. | +| Need recall above 0.95 without a training pass | **`ResidualSnapIndex`** | Two-stage scalar quantization + optional rerank, still training-free. | +| Have a corpus sample and want aggressive compression | **`PQSnapIndex`** | Learned codebooks reach 16-32 B/vec with recall far above scalar at matched bytes. | +| Need sub-linear search above ~100k vectors | **`IVFPQSnapIndex`** | Partitioned search visits `nprobe / nlist` of the corpus. With `rerank_candidates` it breaks the PQ recall ceiling. | + +## Sizing rules of thumb + +| Parameter | Default | Guidance | +|-----------|---------|----------| +| `bits` (SnapIndex) | 4 | 4 for ~95% recall at 6x compression. 3 for a sweet spot around 7.8x. 2 only for huge corpora where scale beats precision. | +| `M` (PQ) | `dim // 4` | Higher M = higher recall, more disk. `M=16` is a common starting point for `dim=384` (BGE-small). | +| `K` (PQ) | 256 | Fixed at 256 (one byte per sub-index). | +| `nlist` (IVF) | `4 * sqrt(N)` | E.g. N=57k -> nlist=512, N=1M -> nlist=4096. | +| `nprobe` (IVF) | `nlist // 16` | Trades recall for latency; tune per query. | +| `rerank_candidates` (IVF) | `None` | Start at 100. Raises recall toward float32 ceiling. | + +## Bits vs recall (SnapIndex) + +| `bits` | Compression vs float32 | Recall@10 on real embeddings | +|--------|------------------------|-------------------------------| +| 2 | 11.6x | ~0.83 (synthetic), higher on clustered real data | +| 3 | 7.8x | ~0.92 | +| 4 | 5.9x | ~0.95 | + +## When does IVF-PQ pay off? + +- **N < 10k**: `SnapIndex` or `PQSnapIndex` full-scan is fine. +- **N = 50k-100k**: `PQSnapIndex` + scan or `IVFPQSnapIndex` -- both feasible. +- **N >= 100k**: `IVFPQSnapIndex` is the clear winner. +- **N >= 500k**: `IVFPQSnapIndex` is effectively required -- float32 + brute-force starts hitting RAM and latency walls. + +See the [benchmarks](../benchmarks.md) page for measured numbers. diff --git a/docs/user-guide/filter-search.md b/docs/user-guide/filter-search.md new file mode 100644 index 0000000..635d2f3 --- /dev/null +++ b/docs/user-guide/filter-search.md @@ -0,0 +1,29 @@ +# Filtered search + +Every index accepts `filter_ids=` to restrict results to a subset +of ids. + +```python +filter_set = {f"doc-{i:04d}" for i in range(100)} +hits = idx.search(query, k=5, filter_ids=filter_set) +``` + +## Performance + +- **SnapIndex / PQSnapIndex / ResidualSnapIndex**: the filter is applied + after scoring, so it does not save scoring work. Useful when the + filter set is small relative to corpus. +- **IVFPQSnapIndex** (cluster-aware): probe ranking is restricted to + clusters that contain at least one filter row, so sparse filters skip + clusters entirely. Rerank candidates are also drawn from the filtered + subset, not the unfiltered probe output. + +## Edge cases + +- Unknown ids in `filter_ids` are silently dropped. +- An entirely-unknown filter returns `[]`. +- A very sparse filter may require a larger `nprobe` on IVF-PQ to surface + `k` hits. + +See [`examples/filter_search.py`](https://github.com/stffns/snapvec/blob/main/examples/filter_search.py) +for a runnable example. diff --git a/docs/user-guide/ivf-pq.md b/docs/user-guide/ivf-pq.md new file mode 100644 index 0000000..54f91c8 --- /dev/null +++ b/docs/user-guide/ivf-pq.md @@ -0,0 +1,76 @@ +# IVFPQSnapIndex + +Sub-linear search at scale, with optional float16 rerank. Inverted-file +coarse partition on top of residual PQ: each query visits only +`nprobe / nlist` of the corpus. With `keep_full_precision=True` + +`rerank_candidates`, the top candidates are re-scored against the stored +float16 vectors, which recovers almost all the recall lost to PQ. + +## Headline numbers + +On BGE-small / FIQA (N = 57,638, dim = 384): + +- recall@10 = **0.977 at 441 us / query** +- recall@10 = **0.998 at 1021 us / query** + +5.8x faster than v0.6 at identical recall, past the PQ-only 0.929 ceiling. + +## When to use + +- N >= 100k -- below that, full-scan `PQSnapIndex` is comparable and simpler. +- Latency budget in the sub-millisecond range. +- Recall target >= 0.97. + +## Basic usage + +```python +import numpy as np +from snapvec import IVFPQSnapIndex + +corpus = np.random.randn(100_000, 384).astype(np.float32) + +idx = IVFPQSnapIndex( + dim=384, + nlist=512, # 4 * sqrt(N) + M=16, + K=256, + normalized=True, + keep_full_precision=True, + seed=0, +) +idx.fit(corpus[:20_000]) +idx.add_batch(list(range(100_000)), corpus) + +hits = idx.search(query, k=10, nprobe=32, rerank_candidates=100) +``` + +## Sizing + +| Parameter | Guidance | +|-----------|----------| +| `nlist` | `4 * sqrt(N)`; clamp between 32 and 65536 | +| `nprobe` | Start at `nlist // 16`; sweep to tune recall vs latency | +| Training set size | `>= 30 * nlist` rows (FAISS rule of thumb) | +| `rerank_candidates` | `None` for PQ-only; `100` for strong recall lift | +| `keep_full_precision` | `True` to enable `rerank_candidates` | + +## Operating points (FIQA, BGE-small, N=57k) + +| `nprobe` | `rerank_candidates` | recall@10 | latency | +|---------|---------------------|-----------|---------| +| 8 | None | 0.85 | 180 us | +| 32 | None | 0.92 | 340 us | +| 64 | 100 | 0.977 | 441 us | +| 256 | 200 | 0.998 | 1021 us | + +See [benchmarks](../benchmarks.md) for the full sweep and reproduction +instructions. + +## File format + +On-disk extension: `.snpi`. Magic `SNPI`, v4 as of v0.9.0 (adds +float16 rerank cache). + +## API + +See [`IVFPQSnapIndex` API reference](../api/ivf-pq.md). diff --git a/docs/user-guide/pq.md b/docs/user-guide/pq.md new file mode 100644 index 0000000..17089b5 --- /dev/null +++ b/docs/user-guide/pq.md @@ -0,0 +1,52 @@ +# PQSnapIndex + +Train-once product quantization. Learns per-subspace k-means codebooks +that adapt to the corpus distribution. Delivers 15-18 percentage points +higher recall@10 than `SnapIndex` at matched bytes/vec on modern LLM +embeddings, and opens ultra-compressed modes (16 / 32 / 64 B/vec) that +scalar quantization cannot reach. + +## When to use + +- You have access to a representative corpus sample (~10-50k rows). +- Recall target is above 0.95. +- You want aggressive compression (16-32 B/vec). + +## Basic usage + +```python +import numpy as np +from snapvec import PQSnapIndex + +corpus = np.random.randn(50_000, 384).astype(np.float32) + +idx = PQSnapIndex(dim=384, M=16, K=256, normalized=True, seed=0) +idx.fit(corpus[:10_000]) # train codebooks on a sample +idx.add_batch(list(range(50_000)), corpus) + +hits = idx.search(query, k=10) +``` + +## Choosing `M` + +- `M` must divide the effective dim (or `pdim` if `use_rht=True`). +- Storage = `M` bytes/vec when `normalized=True` (plus 4 bytes otherwise + for the stored norm). +- For `dim=384` (BGE-small): `M=16` gives 16 B/vec (24x compression vs + float32); `M=32` gives 32 B/vec and ~2pp more recall. + +## Why `use_rht=False` by default + +`PQSnapIndex` learns codebooks directly on the embedding space, so the +RHT step (which decorrelates coordinates for scalar quantization) usually +hurts more than it helps: the k-means fit is already capturing +subspace-local structure. Pass `use_rht=True` only if your embeddings are +very non-gaussian and you see recall drop on held-out queries. + +## File format + +On-disk extension: `.snpq`. Magic `SNPQ`, v1 as of v0.9.0. + +## API + +See [`PQSnapIndex` API reference](../api/pq.md). diff --git a/docs/user-guide/residual.md b/docs/user-guide/residual.md new file mode 100644 index 0000000..f88a603 --- /dev/null +++ b/docs/user-guide/residual.md @@ -0,0 +1,36 @@ +# ResidualSnapIndex + +Training-free, two-stage scalar quantization. Cascades a coarse Lloyd-Max +quantizer and a second-stage residual quantizer to reach operating points +that `SnapIndex` alone cannot (5-7 bits/coord with recall up to 0.96). + +Exposes a **coarse-pass + rerank** search mode that converges to full- +reconstruction recall at `O(rerank_M)` candidates instead of `O(N)`; +`rerank_M = 100` already saturates on the tested corpora. + +## When to use + +- You need recall above 0.95 without running a corpus `fit`. +- You prefer scalar quantization to PQ for simplicity / portability. + +## Basic usage + +```python +import numpy as np +from snapvec import ResidualSnapIndex + +corpus = np.random.randn(10_000, 384).astype(np.float32) + +idx = ResidualSnapIndex(dim=384, b1=3, b2=3, normalized=True, seed=0) +idx.add_batch(list(range(10_000)), corpus) + +hits = idx.search(query, k=10, rerank_M=100) +``` + +## File format + +On-disk extension: `.snpr`. CRC32-checksummed. + +## API + +See [`ResidualSnapIndex` API reference](../api/residual.md). diff --git a/docs/user-guide/save-load.md b/docs/user-guide/save-load.md new file mode 100644 index 0000000..3601c1f --- /dev/null +++ b/docs/user-guide/save-load.md @@ -0,0 +1,39 @@ +# Save and load + +Every index type has `.save(path)` / `.load(path)`. + +```python +idx.save("my.snpv") # SnapIndex +loaded = SnapIndex.load("my.snpv") +``` + +## Guarantees + +- **Atomic writes**: the save target is written to `.tmp` first, + then renamed. A crash mid-write leaves the old file intact (or absent). +- **CRC32 trailer** (since v0.7): an 8-byte checksum is appended. `load()` + verifies it; silent corruption raises `ValueError`. +- **Forward-compatible reads**: older files load in newer versions + wherever the format is documented as backward-compatible; new files + never claim to be readable by older versions. + +## File extensions per index + +| Index | Extension | Magic | +|-------|-----------|-------| +| `SnapIndex` | `.snpv` | `SNPV` | +| `PQSnapIndex` | `.snpq` | `SNPQ` | +| `ResidualSnapIndex` | `.snpr` | `SNPR` | +| `IVFPQSnapIndex` | `.snpi` | `SNPI` | + +Paths can use any extension -- the magic header determines the format. +The canonical extensions are a convention for your tooling. + +## IDs + +`ids` can be any hashable Python value. They are serialized as strings +in the file and round-tripped through `load`. If you pass integers, +they come back as integers; strings come back as strings. + +See [`examples/save_load.py`](https://github.com/stffns/snapvec/blob/main/examples/save_load.py) +for a runnable example. diff --git a/docs/user-guide/snap-index.md b/docs/user-guide/snap-index.md new file mode 100644 index 0000000..3a8cc92 --- /dev/null +++ b/docs/user-guide/snap-index.md @@ -0,0 +1,55 @@ +# SnapIndex + +Training-free scalar-quantized index. Implements +[TurboQuant](https://arxiv.org/abs/2504.19874): randomized Hadamard +transform followed by Lloyd-Max scalar quantization. Works out of the +box on any vector distribution, no calibration or corpus sample required. + +## When to use + +- You don't want to run a one-off `fit` step. +- `dim * N` fits comfortably in RAM even at float32. +- Recall target is 0.92-0.95 at 6-12x compression. + +## Basic usage + +```python +import numpy as np +from snapvec import SnapIndex + +corpus = np.random.randn(10_000, 384).astype(np.float32) +idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx.add_batch(list(range(10_000)), corpus) + +query = np.random.randn(384).astype(np.float32) +hits = idx.search(query, k=10) +``` + +## Bits guidance + +Pick 4-bit unless you have a specific reason: + +| `bits` | Compression | Recall@10 on real embeddings | Notes | +|--------|-------------|------------------------------|-------| +| 2 | 11.6x | ~0.83 | Only for aggressive compression | +| 3 | 7.8x | ~0.92 | Middle ground; tightly packed since v0.3 | +| 4 | 5.9x | ~0.95 | Default, recommended | + +## Unbiased-estimator mode + +For use cases that need unbiased inner-product estimates (KV-cache, +attention), pass `use_prod=True`. Applies the QJL correction at the cost +of roughly 2x search latency: + +```python +idx = SnapIndex(dim=dim, bits=3, use_prod=True) +``` + +## File format + +On-disk extension: `.snpv`. CRC32-checksummed, atomic writes (temp file ++ rename). + +## API + +See [`SnapIndex` API reference](../api/snap-index.md). diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..4314760 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,91 @@ +site_name: snapvec +site_description: Fast compressed approximate nearest-neighbor search. NumPy + Cython. +site_url: https://stffns.github.io/snapvec/ +repo_url: https://github.com/stffns/snapvec +repo_name: stffns/snapvec +edit_uri: edit/main/docs/ + +theme: + name: material + palette: + - media: "(prefers-color-scheme: light)" + scheme: default + primary: indigo + accent: indigo + toggle: + icon: material/brightness-7 + name: Switch to dark mode + - media: "(prefers-color-scheme: dark)" + scheme: slate + primary: indigo + accent: indigo + toggle: + icon: material/brightness-4 + name: Switch to light mode + features: + - navigation.instant + - navigation.tracking + - navigation.sections + - navigation.indexes + - navigation.top + - content.code.copy + - content.code.annotate + - search.suggest + - search.highlight + +plugins: + - search + - mkdocstrings: + handlers: + python: + options: + show_source: false + show_root_heading: true + show_root_full_path: false + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + docstring_style: numpy + members_order: source + filters: ["!^_"] + +markdown_extensions: + - admonition + - pymdownx.details + - pymdownx.superfences + - pymdownx.tabbed: + alternate_style: true + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.inlinehilite + - pymdownx.snippets + - toc: + permalink: true + - tables + - attr_list + +exclude_docs: | + blog/ + +nav: + - Home: index.md + - Getting started: + - Installation: getting-started/installation.md + - Quickstart: getting-started/quickstart.md + - User guide: + - Choosing an index: user-guide/choosing-an-index.md + - SnapIndex: user-guide/snap-index.md + - PQSnapIndex: user-guide/pq.md + - IVFPQSnapIndex: user-guide/ivf-pq.md + - ResidualSnapIndex: user-guide/residual.md + - Save and load: user-guide/save-load.md + - Filtered search: user-guide/filter-search.md + - Architecture: architecture.md + - Benchmarks: benchmarks.md + - API reference: + - SnapIndex: api/snap-index.md + - PQSnapIndex: api/pq.md + - IVFPQSnapIndex: api/ivf-pq.md + - ResidualSnapIndex: api/residual.md + - Helpers: api/helpers.md + - Changelog: changelog.md diff --git a/pyproject.toml b/pyproject.toml index 2595834..d5846cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,11 @@ dev = [ "ruff", "Cython>=3.0", ] +docs = [ + "mkdocs>=1.6", + "mkdocs-material>=9.5", + "mkdocstrings[python]>=0.27", +] [tool.setuptools.packages.find] include = ["snapvec*"] From 03a1a08691a33b5ea62c7024c98c79d9fdf738dd Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 11:44:06 +0200 Subject: [PATCH 3/5] docs: trim README to hero + quickstart + docs links Reduce README from 673 lines to 91. Keep the hero, the decision-tree summary table, install + quickstart, and the headline benchmark. Push everything else (sizing, algorithm, full API reference, roadmap) to the MkDocs site where it lives under structured navigation. Link the mkdocs.io site from the badges row and from a dedicated Docs section so the README is primarily an entry point, not a reference manual. Preserves the Context / Contributing / License sections at the bottom. --- README.md | 688 +++++------------------------------------------------- 1 file changed, 53 insertions(+), 635 deletions(-) diff --git a/README.md b/README.md index bd6fb03..149e27d 100644 --- a/README.md +++ b/README.md @@ -3,671 +3,89 @@ [![PyPI version](https://img.shields.io/pypi/v/snapvec.svg)](https://pypi.org/project/snapvec/) [![Python versions](https://img.shields.io/pypi/pyversions/snapvec.svg)](https://pypi.org/project/snapvec/) [![CI](https://github.com/stffns/snapvec/actions/workflows/ci.yml/badge.svg)](https://github.com/stffns/snapvec/actions/workflows/ci.yml) +[![Docs](https://github.com/stffns/snapvec/actions/workflows/docs.yml/badge.svg)](https://stffns.github.io/snapvec/) [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT) [![Downloads](https://static.pepy.tech/badge/snapvec/month)](https://pepy.tech/project/snapvec) -**Fast compressed approximate nearest-neighbor search. NumPy + Cython compiled kernels.** +**Fast compressed approximate nearest-neighbor search. NumPy + Cython compiled kernels.** -`snapvec` ships **four index types** for embedding vector search, each targeting a different point on the accuracy / storage / latency frontier: +Four index types for embedding vector search, each targeting a different +point on the accuracy / storage / latency frontier: -- **`SnapIndex`** — *training-free scalar*. Implements [TurboQuant](https://arxiv.org/abs/2504.19874) (randomized Hadamard transform + Lloyd-Max scalar quantization). Works out-of-the-box on any vector distribution, no calibration or corpus sample required. **~6× / ~8× / ~12× compression at 4 / 3 / 2 bits** with **>0.92 recall@10** on real embeddings. -- **`ResidualSnapIndex`** — *training-free, two-stage scalar*. Cascades coarse + residual Lloyd-Max quantization to reach operating points `SnapIndex` alone cannot (5 / 6 / 7 bits/coord with recall up to 0.96), and offers a coarse-pass + rerank search mode that converges to full-reconstruction recall at O(`rerank_M`) candidates instead of O(N) (e.g. `rerank_M = 100` already saturates on the tested corpora). -- **`PQSnapIndex`** — *train-once product quantization*. Learns per-subspace k-means codebooks; delivers **+15–18 pp recall@10** over `SnapIndex` at matched bytes/vec on modern LLM embeddings, and opens ultra-compressed modes (16 / 32 / 64 B/vec) that scalar quantization cannot reach. Cost is one offline `fit(sample)` call. -- **`IVFPQSnapIndex`** — *sub-linear search at scale*, with an optional **float16 rerank pass** that breaks the PQ recall ceiling. Inverted-file coarse partition on top of residual PQ, visiting only `nprobe / nlist` of the corpus per query. With `keep_full_precision=True` + `rerank_candidates=100`, IVF-PQ reaches **recall@10 = 0.977 at 441 us/query** on BGE-small / FIQA (N=57K) -- 5.8x faster than v0.6 at identical recall, past the PQ-only 0.929 ceiling. +| Index | Training | Compression | Recall | Use when | +|-------|----------|-------------|--------|----------| +| `SnapIndex` | none | 6-12x | 0.92+ | Any distribution, no corpus sample | +| `ResidualSnapIndex` | none | 4-8x | 0.96 | Higher recall, still training-free | +| `PQSnapIndex` | one-off `fit` | 24-96x | 0.95 | Modern LLM embeddings, aggressive compression | +| `IVFPQSnapIndex` | one-off `fit` | 24-96x | 0.98 | Sub-linear search at scale (N > 100k) | -All four file formats (`.snpv` / `.snpq` / `.snpr` / `.snpi`) carry a CRC32 trailer from v0.7 on -- silent disk or transport corruption is caught at `load()` time instead of returning wrong results. +Headline number: **recall@10 = 0.977 at 441 us/query** on BEIR FIQA +(N = 57,638, BGE-small), 25-125x faster than sqlite-vec at comparable +recall. -### Context +## Install -snapvec was developed as the quantization layer for [vstash](https://github.com/stffns/vstash), a local-first hybrid retrieval system, to extend it to corpora beyond the float32 memory budget while preserving its dependency-minimal design. It stands alone as a quantization library, but the design constraints (NumPy-only base install, no SIMD intrinsics required, predictable latency) come from vstash's local-first requirements. - -``` +```bash pip install snapvec ``` -### Which `bits` should I use? - -| If you are… | Pick | Why | -|-------------|------|-----| -| Building RAG / semantic search (the default) | **`bits=4`** | ~95% recall@10 on real embeddings, 5.9× smaller than float32 | -| Squeezing massive corpora where scale beats precision | **`bits=2`** | 11.6× smaller; recall@10 ≈ 0.83 on synthetic, higher on clustered real data | -| Willing to pay a bit of accuracy for ~25% more compression vs 4-bit | **`bits=3`** | 7.8× smaller, now tightly packed — a genuine middle ground as of v0.3.0 | -| Need unbiased inner-product estimates (KV-cache, attention) | **`bits=3` or `4` + `use_prod=True`** | QJL correction at the cost of ~2× search latency | - -### Which Index should I use? - -| If you need… | Use | Why | -|---|---|---| -| **Maximum accuracy on LLM embeddings** | `PQSnapIndex` | Exploits the embedding model's natural cluster structure. 128 B/vec with PQ matches 256 B/vec with scalar Lloyd-Max on BGE-small. | -| **Zero setup / plug-and-play** | `SnapIndex` | No `.fit()` needed. RHT Gaussianizes any distribution, so fixed Lloyd-Max codebooks work out of the box. | -| **Massive scale (N ≫ 10⁶)** | `PQSnapIndex` | Enables 16 B and 32 B per-vector modes that are not physically reachable with scalar quantization (`SnapIndex` floors at ~128 B/vec for 384-dim embeddings). | -| **Sub-linear search latency (N ≳ 10⁴)** | `IVFPQSnapIndex` | Inverted-file partition + residual PQ. Visits only `nprobe / nlist` of the corpus per query. With rerank: **0.977 recall at 441us** on FIQA (N=57K) -- 3.5x faster than PQ full-scan at higher recall. | -| **Unbiased inner-product estimates** (KV-cache, attention) | `SnapIndex(use_prod=True)` | QJL residual correction; see the TurboQuant_prod section below. | +On macOS you also need `brew install libomp` to build from source; the +wheels on PyPI bundle it. -On BGE-small / SciFact (3-seed mean, disjoint train/eval split): - -| Mode | Bytes/vec | Recall@10 | Δ vs scalar baseline | -|---|---|---|---| -| `SnapIndex(bits=2)` | 128 + 4 | 0.686 | — | -| `PQSnapIndex(M=128, K=256, normalized=True)` | 128 | **0.871** | **+18.5 pp** | -| `SnapIndex(bits=3)` | 192 + 4 | 0.781 | — | -| `PQSnapIndex(M=192, K=256, normalized=True)` | 192 | **0.937** | **+15.6 pp** | -| `SnapIndex(bits=4)` | 256 + 4 | 0.870 | — | -| `PQSnapIndex(M=256, K=256, use_rht=True)` | 256 | **0.948** | **+7.8 pp** | - -Full sweeps in `experiments/bench_pq_scaleup_validation.py` (scalar vs PQ, K ∈ {16, 64, 256}, 3 seeds), `experiments/bench_ivf_pq_contiguous.py` (IVF-PQ layout), and `experiments/bench_v090_fiqa.py` (definitive FIQA recall + latency). - ---- - -## Quick start +## Quickstart ```python import numpy as np from snapvec import SnapIndex -# Build index -idx = SnapIndex(dim=384, bits=4) # 4-bit, ~6x compression -# ⚠ pass float32 arrays — see "Input dtype" below -embeddings = np.asarray(embeddings, dtype=np.float32) -idx.add_batch(ids=list(range(N)), vectors=embeddings) - -# Query -results = idx.search(query_vector.astype(np.float32), k=10) # [(id, score), ...] - -# Persist -idx.save("my_index.snpv") -idx2 = SnapIndex.load("my_index.snpv") # atomic save, v1/v2/v3 compatible -``` - -### Quick start — `PQSnapIndex` (higher recall, one-off `fit`) - -```python -from snapvec import PQSnapIndex - -# Build: choose M subspaces dividing dim (or pdim if use_rht=True) and K ≤ 256 centroids. -# Storage = M bytes/vec when normalized=True (+4 bytes otherwise). -idx = PQSnapIndex(dim=384, M=192, K=256, normalized=True) # 192 B/vec - -# Train once on a representative sample (≥ K, ~10–50 k is plenty) -idx.fit(training_vectors.astype(np.float32)) - -# Index and query exactly like SnapIndex -idx.add_batch(ids=list(range(N)), vectors=corpus.astype(np.float32)) -results = idx.search(query.astype(np.float32), k=10) - -# Persist to its own format (atomic save, .snpq magic SNPQ v1) -idx.save("my_index.snpq") -idx2 = PQSnapIndex.load("my_index.snpq") -``` - -Pick `PQSnapIndex` when you can afford the one-off `fit` step and want the recall lift documented above; pick `SnapIndex` when you need a truly training-free, drop-in index. - -### Quick start — `IVFPQSnapIndex` (sub-linear search at scale) - -```python -from snapvec import IVFPQSnapIndex - -# Rule-of-thumb: nlist ≈ 4·√N (e.g. nlist = 256 for N = 20 000). -idx = IVFPQSnapIndex(dim=384, nlist=256, M=192, K=256, normalized=True) - -# fit() trains both the coarse centroids and the residual PQ codebook -idx.fit(training_vectors.astype(np.float32)) -idx.add_batch(ids=list(range(N)), vectors=corpus.astype(np.float32)) - -# nprobe trades recall for latency. Defaults to max(1, nlist // 16). -results = idx.search(query.astype(np.float32), k=10, nprobe=16) - -idx.save("my_index.snpi") -idx2 = IVFPQSnapIndex.load("my_index.snpi") -``` - -#### Float32 rerank for recall above the PQ ceiling *(v0.6)* - -The IVF-PQ pass has a recall ceiling set by the per-subspace codebook capacity — at `M = 192, K = 256` that ceiling is **0.929** on BGE-small / FIQA, regardless of `nprobe`. Past that, no amount of probing can recover vectors the PQ decoder couldn't distinguish. - -Opt in to a float32 rerank that recovers the missing recall: - -```python -idx = IVFPQSnapIndex( - dim=384, nlist=512, M=192, K=256, normalized=True, - keep_full_precision=True, # +dim × 4 bytes / vec (cached originals) -) -idx.fit(training_vectors) -idx.add_batch(ids, corpus) - -# IVF-PQ picks top-100 candidates, then float32 reranks → top-10. -hits = idx.search(query, k=10, nprobe=32, rerank_candidates=100) -``` - -On BGE-small / FIQA (N=57,638, `nlist=512`, `M=192`, `K=256`, 500 queries, v0.9.0): - -#### Operating points - -| Profile | `nprobe` | Recall@10 | us/query | Use case | -|---------|---:|---:|---:|---| -| **Fastest** | 16 | 0.883 | 304 | Latency-critical, recall <90% acceptable | -| **Interactive** | 32 | 0.944 | 369 | RAG chat, autocomplete (default recommendation) | -| **High precision** | 64 | 0.977 | 441 | Apps where recall matters | -| **Near-exact** | 128 | 0.992 | 635 | When you need to be sure | - -All four sub-millisecond at N=57K with dim=384. The sweet spot is **nprobe=64**: the biggest recall jump is from nprobe=32 to nprobe=64 (+3.3pp). After that, marginal gains decay: nprobe=128 adds only +1.5pp for 44% more latency. - -#### Full sweep - -| `nprobe` | PQ recall | PQ us | + rerank recall | + rerank us | -|---:|---:|---:|---:|---:| -| 4 | 0.661 | 142 | 0.674 | 166 | -| 8 | 0.771 | 195 | 0.795 | 223 | -| 16 | 0.845 | 276 | **0.883** | 304 | -| 32 | 0.891 | 332 | **0.944** | 369 | -| 64 | 0.914 | 419 | **0.977** | 441 | -| 128 | 0.924 | 617 | **0.992** | 635 | -| 256 | 0.929 | 1004 | **0.998** | 1021 | - -Rerank cost is a single `(rerank_candidates, dim) @ (dim,)` matmul (~38k ops at N=100 candidates). `rerank_candidates=100` already saturates the lift -- going to 200 doesn't move the number. +rng = np.random.default_rng(0) +corpus = rng.standard_normal((10_000, 384)).astype(np.float32) -#### Why IVF-PQ + rerank beats PQ full-scan +idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx.add_batch(list(range(10_000)), corpus) -| Mode | Recall@10 | us/query | -|---|---:|---:| -| PQSnapIndex full scan (no IVF, no rerank) | 0.915 | 1,529 | -| IVFPQSnapIndex nprobe=64 + rerank(100) | **0.977** | **441** | +query = rng.standard_normal(384).astype(np.float32) +for doc_id, score in idx.search(query, k=10): + print(doc_id, score) -More recall, 3.5x faster. The rerank pass breaks the accuracy/speed tradeoff: IVF-PQ + rerank is strictly superior to full-scan PQ -- not a compromise, an improvement on both axes. - -Trade-off: the float16 rerank cache adds `dim x 2 bytes` per vector. At `d=384` that's +768 B/vec on top of the `M=192` codes. Opt-in; default `keep_full_precision=False` preserves the codes-only storage footprint. - -#### Reproducibility - -The recall numbers above (0.977 at nprobe=64, 0.992 at nprobe=128) reproduce the v0.6.0 measurements exactly, across three releases with non-trivial algorithmic changes (column-major layout, batched matmul LUT, Cython compiled kernels). Same config, same data, same number -- the evaluation pipeline is deterministic and the rerank path is reproducible. - -#### Performance vs v0.6.0 - -v0.9.0 adds Cython+OpenMP compiled ADC kernels. Latency comparison on the same FIQA corpus: - -| `nprobe` + rerank | v0.6.0 | v0.9.0 | Speedup | -|---|---:|---:|---:| -| 32 | 1,380 us | **369 us** | **3.7x** | -| 64 | 2,550 us | **441 us** | **5.8x** | -| 128 | 4,590 us | **635 us** | **7.2x** | - -snapvec reaches recall 0.98 at sub-ms latency on BEIR FiQA (BGE-small, N=57K) with a NumPy-only runtime dependency. This is in the same recall range as libraries that depend on SIMD popcount intrinsics (RaBitQ) or GPU-accelerated codebooks (FAISS-GPU), achieved here with a single compiled wheel -- trading some asymptotic optimality for distribution simplicity. - -The Cython+OpenMP kernel matches Numba-parallel performance without the 143MB LLVM runtime, so Numba was evaluated and dropped from the dependency graph. snapvec's only runtime dependency is numpy. - -### Sizing `nlist` and the training set - -Two simple rules from FAISS that snapvec follows literally: - -- **`nlist ≈ 4·√N`** is a good starting point. For N = 20 k that is 256; for N = 1 M that is 4 096. -- **`n_train ≥ 30 · nlist`** for the coarse k-means to be stable. Fewer training samples leave many clusters empty or under-trained, and recall stops responding to `nprobe` (we saw this on the v0.5 N=1M baseline: 50 k train at nlist=4 096 → ratio of 12, recall pinned at 0.731 across all `nprobe`). `IVFPQSnapIndex.fit()` emits a `UserWarning` when this ratio is violated. - -| N | suggested `nlist` | minimum `n_train` | -|---:|---:|---:| -| 10 000 | 256 | 7 680 | -| 100 000 | 1 024 | 30 720 | -| 1 000 000 | 4 096 | 122 880 | -| 10 000 000 | 16 384 | 491 520 | - -### Why `PQSnapIndex` defaults to `use_rht=False` - -The Randomized Hadamard Transform (RHT) is essential for *scalar* quantization: it Gaussianizes each coordinate so that a single fixed Lloyd-Max codebook is optimal regardless of the input distribution. That is the entire point of `SnapIndex`. - -For Product Quantization the same rotation is actively harmful: - -- **Structure preservation.** PQ learns a codebook per subspace by finding clusters inside that subspace. RHT spreads every coordinate's variance uniformly across the whole vector, which increases entropy per subspace and destroys the clustered geometry k-means is trying to exploit. -- **Empirical cost.** Turning the RHT off and running k-means on the raw embedding subspaces buys **+15.6 pp recall@10 at 192 B/vec** on BGE-small / SciFact (0.937 vs 0.781 for `SnapIndex(bits=3)`, and 0.781 vs 0.776 for PQ-with-RHT at the same storage). Validated across K ∈ {16, 64, 256} with 3-seed mean and std ≤ 0.01. - -`use_rht=True` remains available for the rare cases where the embedding distribution is genuinely isotropic already, or for sanity-checking against the scalar pipeline. See `experiments/bench_pq_no_rht_ablation.py` for the controlled A/B. - -### Input dtype: pass `np.float32` - -`add_batch`, `add`, and `search` accept any array-like input but cast -internally to `float32`. **If you pass `float64` (the NumPy default), -a full-size temporary copy is allocated during the cast** — for -`add_batch(1M × 384)` that's a transient 1.5 GB allocation on top of -your input array, gone only after quantization completes. - -To avoid this: - -```python -# Models that return float32 directly (most modern embeddings) -embeddings = model.encode(texts) # already float32 → no copy -idx.add_batch(ids, embeddings) - -# Arrays from np.array([...]) default to float64 — cast explicitly -embeddings = np.asarray(my_list, dtype=np.float32) # cast once upfront -idx.add_batch(ids, embeddings) - -# Loading from disk — respect the original dtype or force float32 -embeddings = np.load("vecs.npy").astype(np.float32, copy=False) -``` - -The cast is a correctness convenience, not a design choice — the whole -pipeline (normalize, RHT, quantize) operates in `float32`. - ---- - -## Technical background - -### The problem: embedding vectors are expensive - -Modern embedding models produce float32 vectors of dimension `d ∈ {384, 768, 1536}`. -Storing N vectors requires `4·N·d` bytes; brute-force search costs `O(N·d)` per query. -For N = 1M, d = 384: **1.5 GB RAM**, with inner products dominating inference time. - -Product Quantization (PQ) splits vectors into M sub-vectors and quantizes each -independently. It is effective but requires training a K-means codebook per dataset. -Random Binary Quantization (RaBitQ, 1-bit) is fast but coarse. - -**TurboQuant** achieves near-optimal distortion at b bits per coordinate -**without training codebooks**, by first rotating the space with a -randomized Hadamard transform to make coordinates approximately -Gaussian, then quantizing each coordinate independently with the -optimal scalar quantizer for N(0,1). - -> 📄 **Paper:** Zandieh, Daliri, Hadian, Mirrokni (2025). -> *TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate.* -> ICLR 2026 — [arXiv:2504.19874](https://arxiv.org/abs/2504.19874) -> -> `snapvec` is an independent implementation of the algorithm described -> in that paper, with Cython-compiled hot paths for production latency. -> If you use snapvec in academic work, please cite the TurboQuant paper. - ---- - -### Algorithm - -#### Step 1 — Normalize - -Given a raw embedding `v ∈ ℝᵈ`, compute the unit vector `v̂ = v / ‖v‖` and store -`‖v‖` separately (float32, 4 bytes per vector). - -#### Step 2 — Randomized Hadamard Transform (RHT) - -Pad `v̂` to the next power of 2 (`d' = 2^⌈log₂ d⌉`), then apply: - -``` -x = (1/√d') · H · D · v̂ +idx.save("my.snpv") ``` -where: -- `D = diag(σ₁, …, σ_d')` — diagonal matrix of i.i.d. ±1 random signs (seed-deterministic) -- `H` — unnormalized Walsh-Hadamard matrix (butterfly pattern) +Runnable end-to-end scripts for every index live in +[`examples/`](examples/). -By the Johnson-Lindenstrauss lemma, each coordinate `xᵢ ≈ N(0, 1/d')`. -After rescaling `x̃ = x · √d'`, the coordinates are approximately `N(0,1)` -regardless of the original distribution of `v`. +## Documentation -**Complexity:** O(d log d) — no matrix multiplication, no codebook training. +Full docs: **** -**Implementation note (v0.3.0):** the butterfly is fully vectorised via a -reshape view — each level becomes one pair of NumPy ops on the whole -array instead of a Python `for` loop over `n / (2h)` slices. The diff is -tiny and the speedup is ~24× on a single query at `padded_dim = 512`: +- [Installation](https://stffns.github.io/snapvec/getting-started/installation/) +- [Choosing an index](https://stffns.github.io/snapvec/user-guide/choosing-an-index/) +- [Architecture](https://stffns.github.io/snapvec/architecture/) (RHT, Lloyd-Max, PQ, IVF) +- [Benchmarks](https://stffns.github.io/snapvec/benchmarks/) +- [API reference](https://stffns.github.io/snapvec/api/snap-index/) -```python -# Before: Python loop over n/(2h) slices per level → O(d) dispatches -while h < n: - for i in range(0, n, h * 2): - a = x[..., i : i + h].copy() - b = x[..., i + h : i + 2 * h] - x[..., i : i + h] = a + b - x[..., i + h : i + 2 * h] = a - b - h *= 2 - -# After: reshape view → one pair of ops per level → O(log d) dispatches -while h < n: - view = x.reshape(*x.shape[:-1], n // (2 * h), 2, h) - a = view[..., 0, :].copy() - view[..., 0, :] += view[..., 1, :] # view[0] ← a + b - view[..., 1, :] = a - view[..., 1, :] # view[1] ← a - b - h *= 2 -``` +## Context -`reshape` returns a view only when the array is C-contiguous, so writes -to `view` propagate back to `x` — no extra allocation per level beyond -the single `a.copy()`. In the RHT pipeline we enforce this explicitly -(via `np.ascontiguousarray` at the call site and a defensive assert -inside `_fwht_inplace`); `.astype` alone does not guarantee C-order -when the input is Fortran-contiguous. - -#### Step 3 — Lloyd-Max scalar quantization - -The optimal scalar quantizer for N(0,1) at b bits partitions ℝ into 2^b intervals -and assigns each the conditional mean as reconstruction value. -These boundaries and centroids are precomputed and hardcoded in `snapvec._codebooks` -(no scipy required at runtime): - -| bits | levels | distortion (MSE) | bytes/coord | -|------|--------|------------------|-------------| -| 2 | 4 | 0.1175 | 0.25 | -| 3 | 8 | 0.0311 | 0.375 | -| 4 | 16 | 0.0077 | 0.50 | - -The quantized vector is stored as a `uint8` index matrix, tightly -bit-packed to `b/8` bytes per coordinate both on disk and in RAM. For 3-bit -this means packing 8 indices into 3 bytes (24 bits); for 2-bit and 4-bit -the packing is byte-aligned (4 per byte, 2 per byte respectively). - -#### Step 4 — Approximate inner product - -At search time the query `q` is rotated (not quantized) and the approximate -cosine similarity is computed as: - -``` -score(q, v) = (1/d') · Σᵢ centroid[idx_qᵢ] · centroid[idx_vᵢ] -``` - -This is a single float16 matrix–vector product against the cached centroid expansions. - ---- - -### TurboQuant_prod: unbiased estimator with QJL correction - -The MSE quantizer introduces a small systematic downward bias. The `use_prod=True` mode -corrects this using a **Quantized Johnson-Lindenstrauss (QJL)** residual: - -**Build time (per stored vector):** - -1. Quantize at `(b-1)` bits MSE, compute residual `r = x̃ - x̃_MSE` -2. Store `sign(S·r)` as a 1-bit vector (int8 ±1 in practice), - where `S ∈ ℝ^(d'×d')` is a fixed random Gaussian matrix -3. Store `‖r‖ / √d'` (one float32 per vector) - -**Query time (correction term):** - -``` -correctionᵢ = √(π/2) / d' · ‖rᵢ‖ · dot(S·q̂, sign(S·rᵢ)) -final_scoreᵢ = mse_scoreᵢ + correctionᵢ -``` - -This follows from Lemma 4 of Zandieh et al. (2025): -`E[sign(S·r)] = √(2/π) · S·r / ‖S·r‖`, giving an unbiased estimate of `⟨r, q̂⟩`. - -**When to use `use_prod=True`:** -- When you need accurate inner product magnitudes (KV-cache, attention approximation) -- **Not** recommended for pure ranking/NNS — the added QJL variance degrades recall@k - relative to MSE-only at equal total bits - ---- - -### Compression ratios - -Measured on `d = 384` (BGE-small); the RHT pads to the next power of 2 so -`padded_dim = 512`. All per-vector byte counts below include the padded -dimensions — they are the numbers you actually pay in RAM and on disk. -The trailing `+ 4` is the float32 norm; in `normalized=True` mode it is -still persisted (as 1.0) to keep the on-disk layout stable across modes. - -| Backend | Bytes/vec (disk) | Bytes/vec (RAM, idle) | Disk ratio | RAM ratio | -|----------------|------------------|-----------------------|------------|-----------| -| float32 | 1 536 | 1 536 | 1.0× | 1.0× | -| 4-bit snapvec | 256 + 4 | 256 + 4 | **5.9×** | **5.9×** | -| 3-bit snapvec | 192 + 4 | 192 + 4 | **7.8×** | **7.8×** | -| 2-bit snapvec | 128 + 4 | 128 + 4 | **11.6×** | **11.6×** | -| int8 (naïve) | 384 + 4 | 384 + 4 | 4.0× | 4.0× | - -**RAM = disk:** indices are tightly bit-packed in both. The same byte -layout is used everywhere, so `save` / `load` copy bytes directly without -an intermediate unpack/repack step. This halves indices RAM vs the -pre-v0.2 behaviour (which stored uint8 in RAM and only packed on disk). - -Bit-packing scheme: -- **4-bit** (`0.5 bytes/coord`): 2 indices per byte (byte-aligned). -- **3-bit** (`0.375 bytes/coord`): 8 indices → 3 bytes, cross-byte tight - packing (v3 file format; v1/v2 files are read transparently via the - legacy byte-aligned decoder). -- **2-bit** (`0.25 bytes/coord`): 4 indices per byte (byte-aligned). - -Unpacking happens when the `float16` centroid cache is built (cached -full-scan path — once, off the hot matmul) or per-chunk/per-query in -`chunk_size` / `filter_ids` modes. - -**Search cache:** a lazy `float16` centroid expansion of shape -`(N, padded_dim)` is materialised on first query for fast matmul (~5 ms -at N = 100 k). It is evicted on writes and can be avoided entirely via -`chunk_size` for memory-constrained deployments. - -#### Real-world footprint: 1 M vectors at d = 768 - -Typical for BGE-base, E5-base, `nomic-embed-text-v1`, and other 768-dim -models. The RHT pads to 1024. - -| Backend | Idle RAM | + cache (float16) | Warm peak | -|----------------|----------|--------------------|-----------| -| float32 | **2.86 GiB** | — | 2.86 GiB | -| int8 (naïve) | 0.72 GiB | — | 0.72 GiB | -| 4-bit snapvec | **0.48 GiB** | +1.91 GiB | 2.39 GiB | -| 3-bit snapvec | **0.36 GiB** | +1.91 GiB | 2.27 GiB | -| 2-bit snapvec | **0.24 GiB** | +1.91 GiB | 2.15 GiB | - -Numbers use binary units (1 GiB = 2³⁰ bytes); e.g. float32 is -`1M × 768 × 4 B = 2.86 GiB`. - -The cache is materialised only during active search and is evicted on -any write. With `chunk_size` set, warm peak drops to roughly -`idle RAM + chunk_size × padded_dim × 2 B` (the per-chunk float16 -scratch) — e.g. `chunk_size=10_000` at `padded_dim=1024` adds ~20 MiB, -at the cost of ~10× query latency. This is the usual memory/latency -trade-off, exposed as a first-class flag. - -For a single-server RAG index at this scale, 4-bit snapvec idles at -**half a GiB** where float32 idles at ~3 GiB. - ---- - -### Recall benchmarks - -Measured on synthetic unit-sphere vectors (`d=384`, `N=10 000`, 100 queries). -**Baseline: exact cosine float32 brute-force.** - -| bits | recall@1 | recall@10 | recall@50 | -|------|----------|-----------|-----------| -| 2 | 0.72 | 0.83 | 0.91 | -| 3 | 0.81 | 0.91 | 0.96 | -| 4 | 0.86 | 0.93 | 0.95 | - -Recall improves with clustered (real-world) data. On BGE-small-en embeddings -from mixed document corpora, 4-bit achieves **recall@10 ≈ 0.95**. - -> **Note on published results:** The TurboQuant paper (Zandieh et al., 2025) reports -> recall up to 0.99, measured against HNSW graph navigation (not brute-force float32), -> on GloVe `d=200` data, using recall@1 with large `k_probe`. These conditions differ -> from the above; both results are correct under their respective definitions. - ---- - -### File format (`.snpv`) - -``` -Offset Size Field -────────────────────────────────────────────────── -0 4 B magic: "SNPV" -4 4 B version: uint32 (1, 2, or 3) -8 4 B dim: uint32 — original embedding dimension -12 4 B bits: uint32 — total bits (2, 3, or 4) -16 4 B seed: uint32 — rotation seed -20 4 B n: uint32 — number of stored vectors -24 4 B flags: uint32 — bit-0: use_prod, bit-1: normalized [v2/v3 only] -────────────────────────────────────────────────── -28 4 B packed_len: uint32 -32 * indices: bit-packed uint8 MSE indices - n×4 B norms: float32 per-vector original norms -[prod only] - n×d' B qjl_signs: int8 sign(S·r) per vector - n×4 B rnorms: float32 ‖r‖/√d per vector -────────────────────────────────────────────────── - n×(2+L) ids: uint16-length-prefixed UTF-8 strings -``` - -Saves are **atomic** on POSIX: writes to `.snpv.tmp` then `os.replace()`. -Backward compatible: v1 (mse-only, pre-flags), v2 (flags + byte-aligned -3-bit), and v3 (tight 3-bit packing) files all load correctly — the -reader dispatches the 3-bit decoder on version. - ---- - -## API reference - -### `SnapIndex(dim, bits=4, seed=0, use_prod=False, chunk_size=None, normalized=False)` - -| Parameter | Type | Default | Description | -|---------------|------------|---------|-------------| -| `dim` | int | — | Embedding dimension | -| `bits` | int | 4 | Bits per coordinate: 2, 3, or 4 | -| `seed` | int | 0 | Rotation seed — must be consistent across build and query | -| `use_prod` | bool | False | Enable QJL unbiased estimator (requires bits ≥ 3) | -| `chunk_size` | int \| None | None | Stream search in chunks without the float16 cache (for N > 500k) | -| `normalized` | bool | False | Skip norm computation: trust that input vectors are unit-length | - -### Methods - -```python -idx.add(id, vector) # Add one vector -idx.add_batch(ids, vectors) # Add N vectors (~50x faster than loop) -idx.delete(id) -> bool # Remove by id, O(1) lookup -idx.search(query, k=10, filter_ids=None) # [(id, score), ...] descending -idx.save(path) # Atomic binary save to .snpv -SnapIndex.load(path) # Load from .snpv file -idx.stats() -> dict # Compression / memory diagnostics -len(idx) # Number of stored vectors -repr(idx) # SnapIndex(dim=384, bits=4, mode=mse, n=1000) -``` - -> **All vector inputs should be `np.float32`.** Passing `float64` triggers -> a full-size temporary cast inside `add_batch` / `search` (see "Input -> dtype" in Quick start). Most embedding models already return `float32`; -> only ad-hoc arrays built from Python lists via `np.array(...)` default -> to `float64`. - ---- - -## Relation to TurboQuant / PolarQuant - -`snapvec` implements the core compression pipeline from: - -> Zandieh, A., Daliri, M., Hadian, A., & Mirrokni, V. (2025). -> **TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate.** -> *ICLR 2026.* [arXiv:2504.19874](https://arxiv.org/abs/2504.19874) - -The same algorithm was published concurrently as "PolarQuant" at AISTATS 2026. -Both names were already taken on PyPI; `snapvec` = **Hada**mard + Lloyd-**Max**, -named after its two core operations. - -> **Note on `PQSnapIndex`.** While `SnapIndex` follows the TurboQuant paper -> strictly (RHT + Lloyd-Max), `PQSnapIndex` deliberately departs from that -> pipeline: it skips the RHT and learns per-subspace k-means codebooks on the -> raw embedding space, which lets it exploit the natural clustered structure -> of modern LLM embeddings. The trade-off is that a training-free guarantee -> is replaced with a one-off `fit(sample)` call — in exchange for a strictly -> better recall / storage Pareto frontier on the retrieval workloads we've -> measured. See the decision table at the top of this README. - -Key contributions of this implementation over the reference: - -- **No scipy** -- codebooks hardcoded, numpy is the only runtime dependency. Cython kernels compile at install time -- **Batch WHT** — single O(n·d·log d) call for bulk inserts (~50x faster than loop) -- **Float16 cache** — centroid expansions in half precision, ~2x faster matmul -- **Packed RAM indices** — 2/4-bit indices stored bit-packed in RAM (2× less memory), - unpacked lazily when the float16 cache is built — zero impact on warm query latency -- **Pre-filtering** — `filter_ids` restricts search to a subset with O(|filter| · d) cost -- **Chunked streaming search** — `chunk_size` avoids the float16 cache for large N -- **O(1) delete** — `_id_to_pos` dict + position compaction -- **Atomic saves** — `.snpv.tmp` → `os.replace()` pattern -- **Versioned format** — v1/v2/v3 all loadable, forward-compatible flags field, - legacy 3-bit decoder kept around for pre-v0.3 indices - ---- - -## Roadmap - -### Pure-Python improvements (no new dependencies) - -- **Quantized norms on disk** -- store per-vector norms as uint8 with a - (min, max) header. Saves 3 bytes/vec on disk with <0.1% precision loss. -- **Typed ID storage** -- persist integer IDs as fixed-width uint32/uint64 - instead of UTF-8 strings when all IDs are numeric. -- **Positional-ID fast path** -- when IDs are `0..N-1`, skip `_ids` list - and `_id_to_pos` dict entirely; position is the ID. -- **API parity** -- `filter_ids` for PQSnapIndex and ResidualSnapIndex - (currently only SnapIndex and IVFPQSnapIndex support it). - -### Future directions - -- **Batch RHT + quantization** -- Cython FWHT + `searchsorted` for - `add_batch`. Target: `add_batch(100k, d=384)` 2.5s -> ~200ms. - Lowest priority -- indexing is typically amortized offline. -- **OpenMP tuning** -- the current parallel threshold (n >= 2000) and - schedule (static) were chosen empirically on Apple Silicon. Linux - and many-core x86 may benefit from dynamic scheduling or different - thresholds. - -_Previously planned and now shipped:_ -_Tight 3-bit packing (v0.3.0). Vectorized FWHT (v0.3.0). Compiled -ADC kernels (v0.9.0 via Cython+OpenMP -- originally scoped as Rust -Phase 1-2, delivered with equivalent performance and simpler build)._ - -### Non-goals - -- **Rust / native extensions beyond Cython**. The v0.9.0 Cython+OpenMP - kernels match the performance targets originally scoped for a Rust - accelerator, at the cost of one compiled wheel build. Adding a second - native backend would duplicate maintenance without clear user benefit. -- **Trained codebooks (PQ / OPQ / RaBitQ-trained)** for SnapIndex. Keeps - the "no training required" guarantee; the data-agnostic Lloyd-Max - tables make SnapIndex safe to use without a representative sample. - (PQSnapIndex and IVFPQSnapIndex do use trained codebooks by design.) -- **Graph indices (HNSW / NSG)**. Different trade-off space; snapvec - targets flat and inverted-file indices where compression and - predictable latency matter more than sub-linear graph traversal. -- **GPU acceleration**. The compiled ADC path is compute-bound on CPU - at current scales; GPU would help only at very large N where - network/transfer cost already dominates. - ---- - -## Changelog - -See [`CHANGELOG.md`](./CHANGELOG.md) for the per-release history. Recent -highlights: - -- **v0.9.0** -- Cython+OpenMP compiled ADC kernels. 5.8x faster at recall 0.977 on FIQA. Zero extra runtime dependencies. -- **v0.8.0** -- Thread-safe search via `freeze()`, `filter_ids` for IVFPQSnapIndex. -- **v0.7.0** -- CRC32 trailers on all file formats, fp16 rerank cache (half storage). -- **v0.6.0** -- Float32 rerank pass breaks the PQ recall ceiling (0.929 -> 0.994). -- **v0.5.0** -- IVFPQSnapIndex with sub-linear search, 12x faster build at N=1M. -- **v0.3.0** -- Tight 3-bit packing (7.8x compression), vectorised FWHT. - -## Installation - -```bash -pip install snapvec -``` - -**Requirements:** Python >= 3.10, NumPy >= 1.24. No other runtime dependencies. -Pre-built wheels will be available for common platforms (macOS, Linux x86_64, -Windows) via `pip install`; building from source requires a C compiler for the -Cython extension. - -For development: - -```bash -git clone https://github.com/stffns/snapvec -cd snapvec -pip install -e ".[dev]" -python setup.py build_ext --inplace # compile Cython kernels -pytest tests/ -v -``` +`snapvec` was developed as the quantization layer for +[vstash](https://github.com/stffns/vstash), a local-first hybrid retrieval +system, to extend it to corpora beyond the float32 memory budget while +preserving its dependency-minimal design. It stands alone as a +quantization library, but the design constraints (NumPy-only base +install, predictable latency, reproducible index files) come from +vstash's local-first requirements. -On macOS, `brew install libomp` enables OpenMP parallel scoring. -On Linux, OpenMP support is typically available via gcc out of the box. +## Contributing ---- +See [CONTRIBUTING.md](CONTRIBUTING.md) for dev setup, the test matrix, +and the release process. Bugs and feature requests go to +[issues](https://github.com/stffns/snapvec/issues); questions and +usage help to [discussions](https://github.com/stffns/snapvec/discussions). ## License -MIT © 2025 Jayson Steffens. +MIT (c) 2025 Jayson Steffens. -The TurboQuant algorithm is described in [arXiv:2504.19874](https://arxiv.org/abs/2504.19874) -by Zandieh et al. (Google Research / ICLR 2026). This package is an independent implementation. +The TurboQuant algorithm is described in +[arXiv:2504.19874](https://arxiv.org/abs/2504.19874) by Zandieh et al. +(Google Research / ICLR 2026). This package is an independent +implementation. From f0f26cf548373bae1d706b905654a174fb8ee020 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 11:51:19 +0200 Subject: [PATCH 4/5] docs: address PR #49 bot review feedback All 5 inline comments were valid. - installation.md: float64 is silently cast by np.asarray, not raised as an error. Warn about the hidden allocation instead. - choosing-an-index.md: rename 'Default' column to 'Required?'; clarify that M (PQ) and nlist (IVF) are required ctor args, and give recommended starting values in the guidance column. - filter-search.md: only SnapIndex and IVFPQSnapIndex expose filter_ids today; PQSnapIndex and ResidualSnapIndex do not. - filter-search.md: SnapIndex applies the filter before the matmul, so it *does* save scoring work on sparse filters; updated to describe the actual O(|filter_ids| * dim) cost. - save-load.md: _decode_id parses numeric-looking strings back to int/float. Document the exact behaviour and the 'prefix to keep it a string' workaround. --- docs/getting-started/installation.md | 8 +++++--- docs/user-guide/choosing-an-index.md | 16 ++++++++-------- docs/user-guide/filter-search.md | 13 ++++++++----- docs/user-guide/save-load.md | 13 ++++++++++--- 4 files changed, 31 insertions(+), 19 deletions(-) diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index e3056a6..f300e7f 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -29,9 +29,11 @@ still works but the parallel search paths do not. ## Input dtype -`snapvec` expects `np.float32` inputs everywhere. Passing `np.float64` -(the NumPy default for `np.array([...])`) is an error because silently -downcasting would be a surprise on the hot path. +`snapvec` expects `np.float32` inputs everywhere. Internally, `add_batch` +and `search` pass your arrays through `np.asarray(..., dtype=np.float32)`, +so a `float64` input will be cast silently. That cast allocates a full +copy of the batch on every call, which is easy to miss in a hot path. +Prefer `float32` at the source. ```python # Models that return float32 directly (most modern embeddings) diff --git a/docs/user-guide/choosing-an-index.md b/docs/user-guide/choosing-an-index.md index 9dd1eaa..4efbbaf 100644 --- a/docs/user-guide/choosing-an-index.md +++ b/docs/user-guide/choosing-an-index.md @@ -9,14 +9,14 @@ ## Sizing rules of thumb -| Parameter | Default | Guidance | -|-----------|---------|----------| -| `bits` (SnapIndex) | 4 | 4 for ~95% recall at 6x compression. 3 for a sweet spot around 7.8x. 2 only for huge corpora where scale beats precision. | -| `M` (PQ) | `dim // 4` | Higher M = higher recall, more disk. `M=16` is a common starting point for `dim=384` (BGE-small). | -| `K` (PQ) | 256 | Fixed at 256 (one byte per sub-index). | -| `nlist` (IVF) | `4 * sqrt(N)` | E.g. N=57k -> nlist=512, N=1M -> nlist=4096. | -| `nprobe` (IVF) | `nlist // 16` | Trades recall for latency; tune per query. | -| `rerank_candidates` (IVF) | `None` | Start at 100. Raises recall toward float32 ceiling. | +| Parameter | Required? | Guidance | +|-----------|-----------|----------| +| `bits` (SnapIndex) | no, defaults to 4 | 4 for ~95% recall at 6x compression. 3 for a sweet spot around 7.8x. 2 only for huge corpora where scale beats precision. | +| `M` (PQ) | **yes** | Higher M = higher recall, more disk. Starting point `dim // 4` (e.g. `M=96` for `dim=384`); many users ship with `M=16-32` for aggressive compression. | +| `K` (PQ) | no, defaults to 256 | Leave at 256 (one byte per sub-index). | +| `nlist` (IVF) | **yes** | Target `4 * sqrt(N)`. E.g. N=57k -> nlist=512, N=1M -> nlist=4096. | +| `nprobe` (IVF) | no, defaults to `nlist // 16` | Trades recall for latency; tune per query. | +| `rerank_candidates` (IVF) | no, defaults to `None` | Pass `100` to rerank the PQ candidates with the stored fp16 vectors. Raises recall toward the float32 ceiling. Requires `keep_full_precision=True` at construction. | ## Bits vs recall (SnapIndex) diff --git a/docs/user-guide/filter-search.md b/docs/user-guide/filter-search.md index 635d2f3..eaf7115 100644 --- a/docs/user-guide/filter-search.md +++ b/docs/user-guide/filter-search.md @@ -1,7 +1,9 @@ # Filtered search -Every index accepts `filter_ids=` to restrict results to a subset -of ids. +`SnapIndex` and `IVFPQSnapIndex` accept `filter_ids=` to restrict +results to a subset of ids. `PQSnapIndex` and `ResidualSnapIndex` do +not yet support this argument; filter the returned list in Python if +you need it. ```python filter_set = {f"doc-{i:04d}" for i in range(100)} @@ -10,9 +12,10 @@ hits = idx.search(query, k=5, filter_ids=filter_set) ## Performance -- **SnapIndex / PQSnapIndex / ResidualSnapIndex**: the filter is applied - after scoring, so it does not save scoring work. Useful when the - filter set is small relative to corpus. +- **SnapIndex**: the filter is resolved to a sorted row-index slice + **before** the inner-product matmul, so a sparse filter actively + reduces scoring work (cost ~ `O(|filter_ids| * dim)` instead of + `O(N * dim)`). - **IVFPQSnapIndex** (cluster-aware): probe ranking is restricted to clusters that contain at least one filter row, so sparse filters skip clusters entirely. Rerank candidates are also drawn from the filtered diff --git a/docs/user-guide/save-load.md b/docs/user-guide/save-load.md index 3601c1f..392560d 100644 --- a/docs/user-guide/save-load.md +++ b/docs/user-guide/save-load.md @@ -31,9 +31,16 @@ The canonical extensions are a convention for your tooling. ## IDs -`ids` can be any hashable Python value. They are serialized as strings -in the file and round-tripped through `load`. If you pass integers, -they come back as integers; strings come back as strings. +`ids` can be any hashable Python value. They are serialized as strings +in the file; on `load()`, `snapvec` tries to decode each string as +`int`, then `float`, and falls back to the raw string. That means: + +- `42` -> `42` (int round-trip). +- `3.14` -> `3.14` (float round-trip). +- `"abc"` -> `"abc"` (string round-trip). +- `"123"` -> `123` (the string form of a number is loaded back as `int`, + not `str`). If you need to preserve a numeric-looking string verbatim, + prefix it (for example, `"id-123"`) before ingesting. See [`examples/save_load.py`](https://github.com/stffns/snapvec/blob/main/examples/save_load.py) for a runnable example. From 5974ac8822a678d9ef2fc79b7c9c9d10db406167 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 11:54:06 +0200 Subject: [PATCH 5/5] docs: address Copilot review on PR #49 Thirteen of 16 comments were about the same mismatch: examples and docs passed normalized=True with raw random-normal vectors. normalized=True means 'caller guarantees unit-length inputs' and skips internal normalization, so the snippets were silently computing uncalibrated scores instead of cosine similarity. Drop normalized=True everywhere the surrounding code produces non-unit vectors; the default (False) already does the right thing. Additional fixes: - examples/quickstart.py docstring: SnapIndex is approximate. Reword from 'recover the exact top-1' to 'measure top-1 recall'. - docs/getting-started/quickstart.md: expand the 'ids can be any hashable' line to call out that only numeric-looking values round-trip to their original type; everything else comes back as str. - docs/user-guide/snap-index.md: add an admonition explaining when to use normalized=True so the flag does not get cargo-culted. - docs/user-guide/pq.md / ivf-pq.md / residual.md: define the query variable in each snippet so they are self-contained. - .github/workflows/docs.yml: scope pages:write / id-token:write to the deploy job only; workflow-level permissions applied to PR builds too (principle of least privilege). --- .github/workflows/docs.yml | 7 +++++-- README.md | 2 +- docs/getting-started/quickstart.md | 9 ++++++--- docs/index.md | 2 +- docs/user-guide/ivf-pq.md | 2 +- docs/user-guide/pq.md | 3 ++- docs/user-guide/residual.md | 3 ++- docs/user-guide/snap-index.md | 10 +++++++++- examples/filter_search.py | 2 +- examples/ivf_pq.py | 1 - examples/pq_index.py | 2 +- examples/quickstart.py | 4 ++-- examples/save_load.py | 2 +- examples/streaming_ingest.py | 2 +- 14 files changed, 33 insertions(+), 18 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 263ff51..4a08581 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -13,13 +13,13 @@ concurrency: permissions: contents: read - pages: write - id-token: write jobs: build: name: Build docs runs-on: ubuntu-latest + permissions: + contents: read steps: - uses: actions/checkout@v4 @@ -47,6 +47,9 @@ jobs: needs: build if: github.ref == 'refs/heads/main' runs-on: ubuntu-latest + permissions: + pages: write + id-token: write environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} diff --git a/README.md b/README.md index 149e27d..6bfe873 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ from snapvec import SnapIndex rng = np.random.default_rng(0) corpus = rng.standard_normal((10_000, 384)).astype(np.float32) -idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx = SnapIndex(dim=384, bits=4, seed=0) idx.add_batch(list(range(10_000)), corpus) query = rng.standard_normal(384).astype(np.float32) diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index 5a5a0b5..e362588 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -15,12 +15,15 @@ from snapvec import SnapIndex rng = np.random.default_rng(0) corpus = rng.standard_normal((10_000, 384)).astype(np.float32) -idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx = SnapIndex(dim=384, bits=4, seed=0) idx.add_batch(list(range(10_000)), corpus) ``` -`ids` can be any hashable (int, str, UUID); they are round-tripped through -save/load. +`ids` can be any hashable. They are serialized as strings by `save()`, +so only numeric-looking values (`int`, `float`) round-trip to their +original type; other values (UUIDs, tuples, arbitrary objects) come +back as their `str()` form. See [Save and load](../user-guide/save-load.md) +for the exact behavior. ## 2. Query diff --git a/docs/index.md b/docs/index.md index 6d436a8..10aa77f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -28,7 +28,7 @@ pip install snapvec import numpy as np from snapvec import SnapIndex -idx = SnapIndex(dim=384, bits=4, normalized=True) +idx = SnapIndex(dim=384, bits=4) idx.add_batch(list(range(10_000)), np.random.randn(10_000, 384).astype(np.float32)) results = idx.search(np.random.randn(384).astype(np.float32), k=10) diff --git a/docs/user-guide/ivf-pq.md b/docs/user-guide/ivf-pq.md index 54f91c8..398dd62 100644 --- a/docs/user-guide/ivf-pq.md +++ b/docs/user-guide/ivf-pq.md @@ -28,13 +28,13 @@ import numpy as np from snapvec import IVFPQSnapIndex corpus = np.random.randn(100_000, 384).astype(np.float32) +query = np.random.randn(384).astype(np.float32) idx = IVFPQSnapIndex( dim=384, nlist=512, # 4 * sqrt(N) M=16, K=256, - normalized=True, keep_full_precision=True, seed=0, ) diff --git a/docs/user-guide/pq.md b/docs/user-guide/pq.md index 17089b5..0a0f000 100644 --- a/docs/user-guide/pq.md +++ b/docs/user-guide/pq.md @@ -19,8 +19,9 @@ import numpy as np from snapvec import PQSnapIndex corpus = np.random.randn(50_000, 384).astype(np.float32) +query = np.random.randn(384).astype(np.float32) -idx = PQSnapIndex(dim=384, M=16, K=256, normalized=True, seed=0) +idx = PQSnapIndex(dim=384, M=16, K=256, seed=0) idx.fit(corpus[:10_000]) # train codebooks on a sample idx.add_batch(list(range(50_000)), corpus) diff --git a/docs/user-guide/residual.md b/docs/user-guide/residual.md index f88a603..6d99251 100644 --- a/docs/user-guide/residual.md +++ b/docs/user-guide/residual.md @@ -20,8 +20,9 @@ import numpy as np from snapvec import ResidualSnapIndex corpus = np.random.randn(10_000, 384).astype(np.float32) +query = np.random.randn(384).astype(np.float32) -idx = ResidualSnapIndex(dim=384, b1=3, b2=3, normalized=True, seed=0) +idx = ResidualSnapIndex(dim=384, b1=3, b2=3, seed=0) idx.add_batch(list(range(10_000)), corpus) hits = idx.search(query, k=10, rerank_M=100) diff --git a/docs/user-guide/snap-index.md b/docs/user-guide/snap-index.md index 3a8cc92..ad9ce62 100644 --- a/docs/user-guide/snap-index.md +++ b/docs/user-guide/snap-index.md @@ -18,13 +18,21 @@ import numpy as np from snapvec import SnapIndex corpus = np.random.randn(10_000, 384).astype(np.float32) -idx = SnapIndex(dim=384, bits=4, normalized=True, seed=0) +idx = SnapIndex(dim=384, bits=4, seed=0) idx.add_batch(list(range(10_000)), corpus) query = np.random.randn(384).astype(np.float32) hits = idx.search(query, k=10) ``` +!!! tip "`normalized` is an optimization, not a default" + If your embeddings are already unit-length (for example, cosine-space + outputs from most modern sentence encoders), pass `normalized=True` + to skip the internal L2 normalization step. With raw vectors (like + the example above), leave it at the default `False`. Passing + `normalized=True` on non-unit inputs silently skips normalization + and scores will not match cosine similarity. + ## Bits guidance Pick 4-bit unless you have a specific reason: diff --git a/examples/filter_search.py b/examples/filter_search.py index 03abed4..4eaf8c8 100644 --- a/examples/filter_search.py +++ b/examples/filter_search.py @@ -20,7 +20,7 @@ def main() -> None: corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) ids = [f"doc-{i:04d}" for i in range(n_corpus)] - idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx = SnapIndex(dim=dim, bits=4, seed=0) idx.add_batch(ids, corpus) query = rng.standard_normal(dim).astype(np.float32) diff --git a/examples/ivf_pq.py b/examples/ivf_pq.py index a90984d..769bd2a 100644 --- a/examples/ivf_pq.py +++ b/examples/ivf_pq.py @@ -27,7 +27,6 @@ def main() -> None: nlist=64, M=16, K=256, - normalized=True, keep_full_precision=True, seed=0, ) diff --git a/examples/pq_index.py b/examples/pq_index.py index 9a7d98e..f275a3d 100644 --- a/examples/pq_index.py +++ b/examples/pq_index.py @@ -20,7 +20,7 @@ def main() -> None: corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) queries = rng.standard_normal((n_queries, dim)).astype(np.float32) - idx = PQSnapIndex(dim=dim, M=16, K=256, normalized=True, seed=0) + idx = PQSnapIndex(dim=dim, M=16, K=256, seed=0) idx.fit(corpus[:1000]) # train codebooks on first half idx.add_batch(list(range(n_corpus)), corpus) diff --git a/examples/quickstart.py b/examples/quickstart.py index 69dc259..f927586 100644 --- a/examples/quickstart.py +++ b/examples/quickstart.py @@ -1,7 +1,7 @@ """Minimal SnapIndex example. Build a 4-bit scalar-quantized index over 1,000 random vectors and -recover the exact top-1 for a handful of queries without any training. +measure top-1 recall for a handful of queries without any training. Run with: python examples/quickstart.py """ @@ -19,7 +19,7 @@ def main() -> None: corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) queries = rng.standard_normal((n_queries, dim)).astype(np.float32) - idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx = SnapIndex(dim=dim, bits=4, seed=0) idx.add_batch(list(range(n_corpus)), corpus) print(f"SnapIndex: n={len(idx)}, dim={dim}, bits=4") diff --git a/examples/save_load.py b/examples/save_load.py index 75a4c52..c4fbb64 100644 --- a/examples/save_load.py +++ b/examples/save_load.py @@ -22,7 +22,7 @@ def main() -> None: corpus = rng.standard_normal((n_corpus, dim)).astype(np.float32) - idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx = SnapIndex(dim=dim, bits=4, seed=0) idx.add_batch(list(range(n_corpus)), corpus) with tempfile.TemporaryDirectory() as tmp: diff --git a/examples/streaming_ingest.py b/examples/streaming_ingest.py index c7228ea..6aee3e9 100644 --- a/examples/streaming_ingest.py +++ b/examples/streaming_ingest.py @@ -17,7 +17,7 @@ def main() -> None: rng = np.random.default_rng(0) dim, batch_size, n_batches = 64, 100, 10 - idx = SnapIndex(dim=dim, bits=4, normalized=True, seed=0) + idx = SnapIndex(dim=dim, bits=4, seed=0) for batch in range(n_batches): vecs = rng.standard_normal((batch_size, dim)).astype(np.float32)