pybed

Fast, on-device static embeddings + semantic search in pure Python — a Python port of gobed / zbed.

The model is sentence-transformers/static-retrieval-mrl-en-v1, quantized to int8 with a per-token float32 scale and reduced to 512 dimensions (~16 MB, bundled in model/). It is a static embedder: there is no transformer forward pass. "Inference" is a token-row lookup, scale, and mean-pool — a handful of memory reads per token.

text ──tokenize(int16)──▶ token ids ──gather rows · scale · mean-pool──▶ 512-d vector ──quantize──▶ int8

Features

• Low-allocation BERT WordPiece tokenizer — int16 token ids (the 30522 vocab fits a signed 16-bit int), written into a caller-supplied pre-allocated buffer so the hot path allocates nothing for token ids.
• Static int8 model — loaded straight from safetensors; embedding is a vectorized numpy gather + mean-pool.
• Two search indexes over int8 vectors:
  • FlatIndex — exact cosine top-k (int8 matmul on CPU, or a custom CUDA kernel on GPU).
  • CagraIndex — a custom, dependency-free CAGRA-style graph index: a fixed-degree k-NN graph searched with greedy best-first traversal from multiple entry points.
• Custom CUDA kernels (pybed/cuda_kernels.py, via cupy RawKernel) for int8 cosine scoring and token embedding, with automatic numpy fallback if a usable GPU/cupy build is not present.

Install (uv)

uv venv                       # create .venv
uv pip install -e .           # CPU only (numpy is the only core dep)
uv pip install -e ".[gpu]"    # + cupy-cuda12x for the custom CUDA kernels
uv pip install -e ".[gpu,dev]"  # + pytest

uv run pytest                 # run the suite (GPU tests auto-skip without a GPU)
uv run pybed embed "hello"    # run the CLI

A committed uv.lock pins the full resolution for reproducible installs (uv sync).

Plain pip works too: pip install -e ".[gpu,dev]".

GPU / nvrtc requirement. The CUDA kernels need an nvrtc that knows your GPU's arch. On Blackwell (RTX 5090, sm_120) that means nvrtc ≥ 12.8 — cupy finds it from your system CUDA toolkit (/usr/local/cuda, here 12.9) as long as no older nvidia-cuda-nvrtc-cu12 wheel (e.g. 12.6) shadows it in the environment. A clean uv venv does the right thing. Check with:

from cupy.cuda import nvrtc; print(nvrtc.getSupportedArchs())  # must include 120

Usage

from pybed import EmbedModel, FlatIndex, CagraIndex

model = EmbedModel.from_dir("model")

# Embed
vec = model.embed("a fluffy cat purring")          # float32 [512]
q, scale, norm = model.embed_quantized("a cat")     # int8 [512], scale, l2-norm

# Index a corpus (int8 vectors) and search
import numpy as np
docs = ["a fluffy cat", "the stock market rallied", "a rocket to mars"]
emb = np.stack([model.embed_quantized(d)[0] for d in docs])

index = FlatIndex(emb)                # exact
# index = CagraIndex(emb, degree=32)  # graph ANN
qv, _, qn = model.embed_quantized("kitten meowing")
for r in index.search(qv, qn, top_k=5):
    print(r.score, docs[r.doc_idx])

CLI

pybed embed "semantic file search"
pybed search "opus audio" --path . --cagra --limit 10

Quality — NDCG@10

tests/test_ndcg.py runs the whole pipeline (tokenize → embed → quantize → search) end-to-end on a hardcoded 30-document, 6-topic corpus with 6 labelled queries:

index	mean NDCG@10	recall@10 vs exact
FlatIndex	0.998	1.000 (exact)
CagraIndex	0.998	1.000

pytest -s            # prints the NDCG numbers

Benchmark (RTX 5090 / CUDA 12.9)

$ uv run python bench.py --gpu
vocab=30522 dim=512 | cuda_usable=True
tokenize         :   17.0 us/op   (58,700 ops/s)
embed (tok+pool) :   33.2 us/op   (30,100 ops/s)
flat search (20,000 docs) [gpu]:   0.34 ms/query  (2,959 qps)   # vs 104.7 ms on CPU — ~300x
cagra search (5,000 docs)      :   2.25 ms/query  (445 qps)

The custom CUDA kernel turns the 20k-doc exact scan from 105 ms into 0.34 ms.

CUDA kernels

pybed/cuda_kernels.py ships hand-written CUDA (compiled at import via cupy RawModule) for:

• cosine_scores — one block per doc, cooperative int8 dot-product reduction (used by FlatIndex(use_gpu=True)),
• embed_tokens — one thread per output dim, gather + scale + mean-pool (used by EmbedModel.embed_gpu).

cuda_usable() runs a one-time probe-launch; if the kernel can't both compile and launch on the current device, every GPU call transparently falls back to numpy — so use_gpu=True / embed_gpu are always safe. Verified running on an RTX 5090 (sm_120) — tests/test_gpu.py asserts the kernels match the numpy reference (and auto-skips on CPU-only hosts).

Layout

pybed/
  tokenizer.py     WordPiece, int16, zero-alloc hot path
  model.py         safetensors load + static int8 embedding
  search.py        FlatIndex (exact) + CagraIndex (graph ANN)
  cuda_kernels.py  cupy RawKernel CUDA + probe/fallback
  cli.py           embed / search CLI
model/             modelint8_512dim.safetensors (~16 MB) + tokenizer.json
tests/             tokenizer, embed, NDCG@10, and GPU-kernel tests
bench.py           throughput benchmark (--gpu)
uv.lock            pinned, reproducible resolution

License

MIT