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/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..4a08581 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,59 @@ +name: Docs + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: docs-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + build: + name: Build docs + runs-on: ubuntu-latest + permissions: + contents: read + 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 + permissions: + pages: write + id-token: write + 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/README.md b/README.md index bd6fb03..6bfe873 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, 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. 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..f300e7f --- /dev/null +++ b/docs/getting-started/installation.md @@ -0,0 +1,54 @@ +# 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. 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) +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..e362588 --- /dev/null +++ b/docs/getting-started/quickstart.md @@ -0,0 +1,54 @@ +# 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, seed=0) +idx.add_batch(list(range(10_000)), corpus) +``` + +`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 + +```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..10aa77f --- /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) +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..4efbbaf --- /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 | 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) + +| `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..eaf7115 --- /dev/null +++ b/docs/user-guide/filter-search.md @@ -0,0 +1,32 @@ +# Filtered search + +`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)} +hits = idx.search(query, k=5, filter_ids=filter_set) +``` + +## Performance + +- **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 + 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..398dd62 --- /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) +query = np.random.randn(384).astype(np.float32) + +idx = IVFPQSnapIndex( + dim=384, + nlist=512, # 4 * sqrt(N) + M=16, + K=256, + 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..0a0f000 --- /dev/null +++ b/docs/user-guide/pq.md @@ -0,0 +1,53 @@ +# 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) +query = np.random.randn(384).astype(np.float32) + +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) + +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..6d99251 --- /dev/null +++ b/docs/user-guide/residual.md @@ -0,0 +1,37 @@ +# 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) +query = np.random.randn(384).astype(np.float32) + +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) +``` + +## 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..392560d --- /dev/null +++ b/docs/user-guide/save-load.md @@ -0,0 +1,46 @@ +# 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; 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. diff --git a/docs/user-guide/snap-index.md b/docs/user-guide/snap-index.md new file mode 100644 index 0000000..ad9ce62 --- /dev/null +++ b/docs/user-guide/snap-index.md @@ -0,0 +1,63 @@ +# 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, 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: + +| `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/examples/filter_search.py b/examples/filter_search.py new file mode 100644 index 0000000..4eaf8c8 --- /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, 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..769bd2a --- /dev/null +++ b/examples/ivf_pq.py @@ -0,0 +1,55 @@ +"""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, + 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..f275a3d --- /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, 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..f927586 --- /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 +measure top-1 recall 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, 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..c4fbb64 --- /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, 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..6aee3e9 --- /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, 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() 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*"]