rembed

README

rembed

Pure-Go embedding inference engine for BERT-family encoder models. Text in, L2-normalized embedding vectors out — no cgo, no ONNX Runtime, one static binary.

// Loads straight from the Hugging Face Hub (pure Go, cached locally) —
// no Python, no conversion step:
emb, err := rembed.Load("sentence-transformers/all-MiniLM-L6-v2")
vecs, err := emb.Embed(ctx, []string{"hello world"})
// vecs[0] is a []float32 of emb.Dim() (384 for MiniLM-L6-v2)

EmbedTokens returns per-token hidden states (ONNX Runtime's last_hidden_state) for rerankers and late-interaction retrieval, and a multi-text Embed call fans out across texts for near-linear batch throughput (bit-identical to one-at-a-time results).

Load accepts a Hub model id (downloaded into $REMBED_CACHE, default the user cache dir; HF_TOKEN honored), a git-cloned HF repo directory, or a converted model dir. Options: rembed.WithInt8() (weight-only quantization, ~4× less weight traffic, cosine ≥ 0.999 vs fp32), rembed.WithWorkers(n) (CPU cap for servers), and rembed.WithDim(d) (Matryoshka: truncate to d dims and re-normalize — EmbeddingGemma 768→512/256/128 — for cheaper storage and search; CLI -dim).

Status: the optimization ladder is complete — naive baseline to statistical parity with (and, with int8, consistently ahead of) ONNX Runtime on the reference laptop: ~45× → 0.89× across six rungs, every step measured against a golden ONNX reference within 1e-4 (int8: cosine ≥ 0.999). See DESIGN.md for the architecture and bench/RESULTS.md for the full measured ladder, including the failed experiments. Weight-only int8 is opt-in via rembed.WithInt8(); rembed.WithWorkers(n) caps per-call CPU for throughput-saturated servers.

Supported models

Seven architectures: BERT-family, DistilBERT, MPNet, RoBERTa (including XLM-RoBERTa — multilingual-e5-base/-large, bge-m3, the same encoder with the SentencePiece tokenizer), and ModernBERT encoders, plus two decoder-derived embedders: Qwen3-Embedding (a causal decoder) and EmbeddingGemma (a bidirectional Gemma 3 backbone — the current MMTEB state of the art for its size). sentence-transformers format: mean, CLS, or last-token pooling, with an optional Dense projection head (EmbeddingGemma); WordPiece, byte-level BPE, SentencePiece Unigram (the XLM-R tokenizer — multilingual models work, 100+ languages), or the Gemma byte-fallback BPE; absolute positions (plus MPNet's bucketed relative-position bias) OR rotary positions (RoPE, single- or dual-theta — ModernBERT, Qwen3, EmbeddingGemma); alternating global/local sliding-window attention (ModernBERT, EmbeddingGemma), full causal attention (Qwen3), or bidirectional attention with grouped-query attention and QK-norm (Qwen3, EmbeddingGemma); exact GELU, tanh-GELU, GeGLU, and SwiGLU; LayerNorm and RMSNorm (unit-offset for Gemma); F32/F16/BF16 safetensors. Validated end-to-end against each model's own ONNX Runtime reference (ModernBERT and Qwen3 against the canonical PyTorch ModernBertModel / Qwen3Model, since their ONNX exports bundle or omit the pooling rembed reproduces; XLM-RoBERTa against PyTorch XLMRobertaModel, and EmbeddingGemma against PyTorch Gemma3TextModel with the sentence-transformers pool+Dense+normalize head, since neither reliably ships ONNX):

model	pooling	dtype	fp32 vs ONNX	int8
sentence-transformers/all-MiniLM-L6-v2	mean	F32	1.5e-7	cosine ≥ 0.9991
sentence-transformers/all-MiniLM-L12-v2	mean	F32	1.9e-7	in bounds
sentence-transformers/paraphrase-MiniLM-L3-v2	mean	F32	< 1e-4	—
BAAI/bge-small-en-v1.5	cls	F32	< 1e-4	in bounds
sentence-transformers/all-mpnet-base-v2	mean	F32	3.3e-7	cosine ≥ 0.9978
sentence-transformers/all-distilroberta-v1	mean	F32	3.2e-7	cosine ≥ 0.9985
sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2	mean	F32	7e-7	cosine ≥ 0.9995
intfloat/multilingual-e5-small	mean	F32	1.7e-7	cosine ≥ 0.9995
intfloat/multilingual-e5-base	mean	F32	2.9e-7 (vs PyTorch)	cosine ≥ 0.999
BAAI/bge-base-en-v1.5	cls	F32	7.4e-7	cosine ≥ 0.995
thenlper/gte-base	mean	F32	3.8e-6	cosine ≥ 0.988
sentence-transformers/paraphrase-mpnet-base-v2	mean	F32	1.2e-6	cosine ≥ 0.9945
sentence-transformers/multi-qa-MiniLM-L6-cos-v1	mean	F32	2.1e-7	cosine ≥ 0.998
Snowflake/snowflake-arctic-embed-s	cls	F32	2.5e-7	cosine ≥ 0.995
sentence-transformers/multi-qa-distilbert-cos-v1	mean	F32	2.5e-7	cosine ≥ 0.999
nomic-ai/modernbert-embed-base	mean	F32	< 1e-4 (vs PyTorch)	cosine ≥ 0.998
Qwen/Qwen3-Embedding-0.6B	lasttoken	BF16	< 1e-4 (vs PyTorch)	cosine ≥ 0.997
google/embeddinggemma-300m	mean + Dense	F32	< 1e-4 (vs PyTorch)	cosine ≥ 0.998
thenlper/gte-small	mean	F16	2e-3 maxAbs + cosine ≥ 0.9999 + meanAbs ≤ 2e-4 (the repo's ONNX export is fp32 while its safetensors are f16, so maxAbs is dominated by the checkpoint's own rounding; the cosine/mean bounds are what actually constrain rembed)	—

On CPUs with AVX-VNNI — Intel Alder Lake (2021) onward and Sapphire Rapids+ servers, AMD Zen 5+; note that AVX-512-VNNI-only parts like Ice Lake-SP and Zen 4 do NOT have it — WithInt8Activations selects full int8 inference (u8 activations × s8 weights via VPDPBUSD) for a further ~1.3× over weight-only int8. The accuracy trade is real, PER-MODEL, and test-enforced (worst golden cosine, full int8 vs weight-only):

model	full int8	weight-only int8
MiniLM-L6 / L12 / L3	0.9917 / 0.9932 / 0.9979	≥ 0.9990
mpnet-base / paraphrase-mpnet	0.9912 / 0.9867	≥ 0.9945
multilingual MiniLM / multilingual-e5	0.9982 / 0.9988	≥ 0.9995
multilingual-e5-base (xlm-roberta)	0.9849	0.9992
gte-small / gte-base	0.9991 / 0.9741	≥ 0.9880
multi-qa MiniLM / distilbert	0.9949 / 0.9854	≥ 0.9940
arctic-embed-s	0.9932	0.9953
distilroberta	0.9747	0.9987
modernbert-embed	0.9660	0.9984
qwen3-embedding-0.6B	0.9747	0.9978
embeddinggemma-300m	0.9938	0.9981
bge-base	0.9593	0.9957

Activation outliers are a PER-CHECKPOINT property, not an architecture one: bge-base (a plain BERT) measures worst of all at 0.9593, below distilroberta's 0.9747 and modernbert-embed's 0.9660 (whose GeGLU gate activations have a range the per-row u8 scale can't hold), while bge-base's sibling bge-small is unremarkable. Qwen3-Embedding compounds this: last-token pooling reads a single position, so there is no averaging across tokens to soften activation-quantization error — prefer WithInt8 (weight-only) there. Check the table before enabling full int8 for a model — anything below ~0.99 is a real retrieval-quality risk — and prefer WithInt8 (weight-only, ≥ 0.988 everywhere) when in doubt. Every figure above is enforced in the golden matrix.

Cross-engine, measured on a Zen 4 cloud box with a both-orders/median protocol (bench/RESULTS.md has the full data and every noise flag): rembed fp32 sits at parity with ONNX Runtime fp32, and rembed full int8 beat ORT fp32 in every round — the flag-free rounds measured 0.70× and 0.75× (5.9 ms vs 7.9 ms on mpnet) — while trading blows at parity with ORT's own AVX-512-VNNI int8 graphs.

Disk-backed weights (run larger than RAM). WithDiskWeights() memory-maps the weights from a pack file instead of loading them into RAM: the OS pages weights in on access and evicts under pressure, so resident memory tracks the working set and a model larger than RAM runs (disk-bandwidth-bound when it does not fit, full speed with a warm page cache when it does — the same trade ORT's mmap mode makes). On first use the safetensors (single-file or sharded) are streamed to a pack file one tensor at a time, so even the pack step fits a small box. This is what lets Qwen3-Embedding-4B run cgo-free on a laptop that cannot hold it in fp32 RAM. Close the Embedder to unmap. Numerics are unchanged — only where the bytes live. (Currently wired for qwen3.)

Expected compatible (same architecture, no committed golden yet): the remaining e5 sizes, the largest BGE/GTE variants, the msmarco families, other BERT/DistilBERT-based sentence-transformers checkpoints, and the larger Qwen3-Embedding sizes (4B/8B — same architecture, far larger). Caveat for retrieval models: e5 requires "query: "/"passage: " prefixes, Qwen3-Embedding expects an instruction on queries only ("Instruct: {task}\nQuery:{text}", with documents left bare), and some models (e.g. arctic) declare prompt handling in their pooling config — rembed embeds exactly the text you pass and does not add prefixes; add them yourself or retrieval quality silently degrades. Validated models in that category: the e5 family ("query: "/"passage: "), bge ("Represent this sentence for searching relevant passages: " on queries), and arctic-embed (its own query prefix) — rembed embeds exactly the text you pass.

One deliberate tokenizer divergence: on NFD (decomposed) Hangul/kana — routine output from macOS — HF's fast tokenizer skips ≥6-byte grapheme clusters during normalization and shreds Korean into jamo; rembed matches the sentencepiece C++ reference instead, which composes NFD back so decomposed and composed text embed identically. 65k-input fuzzing against the reference: zero mismatches.

Dev: golden reference generation

The validation harness's golden files come from ONNX Runtime in Python (ModernBERT from the canonical PyTorch ModernBertModel instead; this is a dev-time tool; users never need it):

cd models
python3 -m venv .venv && .venv/bin/pip install -r requirements.txt
.venv/bin/python convert.py sentence-transformers/all-MiniLM-L6-v2

Python

The same in-process engine is callable from Python through a C-shared build (ctypes, ~µs call overhead; the Go library itself stays cgo-free — the shared object is a separate artifact for foreign callers):

python/build.sh   # needs a C toolchain; produces python/rembed/librembed.so

import sys; sys.path.insert(0, "python")
from rembed import Embedder

emb = Embedder("models/all-MiniLM-L6-v2")           # fp32
emb = Embedder("models/all-MiniLM-L6-v2", int8=True)  # weight-only int8
vecs = emb.embed(["hello world"])                    # (n, dim) float32 numpy

Validated against the same golden reference as the Go tests (python/test_rembed.py); vectors cross the ABI bit-identically.

CLI

go run ./cmd/rembed embed    -model models/all-MiniLM-L6-v2 "some text"
go run ./cmd/rembed validate -model models/all-MiniLM-L6-v2
go run ./cmd/rembed bench    -model models/all-MiniLM-L6-v2

License

Apache-2.0

Documentation

Overview

Package rembed is a pure-Go embedding inference engine for BERT-family encoder models: text in, L2-normalized embedding vectors out. No cgo, no ONNX Runtime.

emb, err := rembed.Load("models/all-MiniLM-L6-v2")
vecs, err := emb.Embed(ctx, []string{"hello world"})

Index

• type Embedder
  • func Load(ref string, opts ...Option) (*Embedder, error)
  • func (e *Embedder) Close() error
  • func (e *Embedder) Dim() int
  • func (e *Embedder) Embed(ctx context.Context, texts []string) ([][]float32, error)
  • func (e *Embedder) EmbedTokens(ctx context.Context, texts []string) ([]TokenEmbeddings, error)
  • func (e *Embedder) Model() string
  • func (e *Embedder) Quantized() bool
  • func (e *Embedder) QuantizedActivations() bool
  • func (e *Embedder) Tokenize(text string) []int64
• type Option
  • func WithDim(d int) Option
  • func WithDiskWeights() Option
  • func WithInt8() Option
  • func WithInt8Activations() Option
  • func WithWorkers(n int) Option
• type TokenEmbeddings

Types

type Embedder

type Embedder struct {
	// contains filtered or unexported fields
}

Embedder turns texts into fixed-size embedding vectors. It is safe for concurrent use. Latency is the default optimization target: each Embed call fans out across up to GOMAXPROCS cores and keeps a spinning worker pool for its duration, which trades idle-core burn for wall time — see WithWorkers to cap that for throughput-saturated servers.

func Load

func Load(ref string, opts ...Option) (*Embedder, error)

Load opens a model. ref may be:

• a local directory: either a converted dir (manifest.json) or a plain Hugging Face repo checkout (config.json + 1_Pooling/... — the manifest is derived on the fly);
• a Hugging Face model id like "sentence-transformers/all-MiniLM-L6-v2" (or explicitly "hf:org/name", which always means the Hub): the files are downloaded straight from the Hub in pure Go into $REMBED_CACHE (default: the user cache dir) and reused on later loads. Set HF_TOKEN for gated repos.

A ref that could be BOTH — an org/name that does not exist locally but whose first segment IS a local directory (e.g. a missing "models/all-MiniLM-L6-v2") — is treated as a missing LOCAL path: a typo must fail as one, not turn into silent network egress carrying HF_TOKEN. Use the "hf:" prefix to force the Hub in that situation.

func (*Embedder) Close

func (e *Embedder) Close() error

Close releases resources held by the Embedder — the memory-mapped weights file when loaded WithDiskWeights. It is a safe no-op for RAM-loaded models. After Close, the Embedder must not be used.

func (*Embedder) Dim

func (e *Embedder) Dim() int

Dim returns the embedding dimensionality — the WithDim truncation when set, else the model's full hidden size.

func (*Embedder) Embed

func (e *Embedder) Embed(ctx context.Context, texts []string) ([][]float32, error)

Embed returns one embedding per input text, each of length Dim(), L2-normalized when the model manifest says so (true for the sentence-transformers models). Texts are embedded independently, and a batch of several texts is fanned out ACROSS texts (each forward pass serial or lightly parallel) — near-linear throughput scaling with zero padding waste, and results bit-identical to embedding one at a time. ctx is checked before each text's forward pass; a forward already in flight runs to completion. NOTE: a batch call can hold up to min(GOMAXPROCS, WithWorkers) scratch buffers (~25 MB each at max sequence length) simultaneously — size servers with WithWorkers.

func (*Embedder) EmbedTokens

func (e *Embedder) EmbedTokens(ctx context.Context, texts []string) ([]TokenEmbeddings, error)

EmbedTokens returns per-token hidden states for each text — the raw material for rerankers, late-interaction (ColBERT-style) retrieval, and custom pooling. Embed remains the API for sentence vectors; nothing here is pooled or normalized. Batches fan out across texts exactly like Embed. All results are held live for the whole call — a 256-text batch of long inputs is ~200 MB of hidden states.

func (*Embedder) Model

func (e *Embedder) Model() string

Model returns the model name from the manifest.

func (*Embedder) Quantized

func (e *Embedder) Quantized() bool

Quantized reports whether an int8 path is actually active (WithInt8 or WithInt8Activations requested AND every dense weight packed as int8 — the engine falls back to fp32 per-matrix when the CPU or a shape cannot take it).

func (*Embedder) QuantizedActivations

func (e *Embedder) QuantizedActivations() bool

QuantizedActivations reports whether the full u8-activation VNNI path is active for every dense weight.

func (*Embedder) Tokenize

func (e *Embedder) Tokenize(text string) []int64

Tokenize exposes the tokenizer's input_ids for one text. It exists for the validation harness (attributing golden mismatches to tokenization vs numerics) and debugging; it is not a stable part of the embedding API.

type Option

type Option func(*loadOptions)

Option configures Load.

func WithDim

func WithDim(d int) Option

WithDim truncates each embedding to its first d dimensions and re-L2-normalizes — Matryoshka Representation Learning (MRL). Models trained for it (EmbeddingGemma: 768→512/256/128; nomic-embed) keep most of their quality at a fraction of the storage and search cost. d must be in [1, full dim]; d equal to the full dim is a no-op. The truncation is a deterministic slice-then-renormalize of the final vector, so it is meaningful only for MRL-trained models — on others it still runs but the shorter vector is not a faithful embedding. Dim() reflects the truncated size.

func WithDiskWeights

func WithDiskWeights() Option

WithDiskWeights memory-maps the model's weights from a pack file on disk instead of loading them into RAM: the OS pages weights in on access and evicts them under memory pressure, so a model larger than RAM runs (disk-bandwidth-bound, but it runs) and resident memory tracks the working set. On first use the safetensors are packed to <dir>/weights.rembedpack (streamed one tensor at a time, so the pack step also fits a small box); later loads mmap it directly. Currently supported for qwen3 (the decoder embedder whose 4B/8B sizes motivate it). The Embedder MUST be Closed to unmap. Numerics are unchanged — only where the bytes live changes.

func WithInt8

func WithInt8() Option

WithInt8 selects weight-only int8 inference: transformer dense weights are quantized at load (per-output-channel symmetric scales; activations stay float32), cutting per-embed weight traffic ~4× (~42 MB → ~10.5 MB for MiniLM). Token embeddings stay fp32, so RESIDENT model memory drops ~1.5×, not 4×. Embeddings differ slightly from fp32 — see the int8 golden test for the measured bound. On CPUs without AVX2+FMA the engine silently falls back to fp32; check Quantized() when the mode matters.

func WithInt8Activations

func WithInt8Activations() Option

WithInt8Activations selects FULL int8 inference on CPUs with AVX-VNNI (Alder Lake 2021+, and AVX-512-VNNI servers): weights per-channel int8 AND activations quantized per row to u8 at matmul time, so VPDPBUSD does four multiply-accumulates per lane per instruction. Accuracy trades a little further than weight-only int8 — see the measured bound in the golden test — and the engine falls back to weight-only int8 (then fp32) where VNNI or a shape is unavailable; check QuantizedActivations() when the mode matters. Implies WithInt8.

func WithWorkers

func WithWorkers(n int) Option

WithWorkers caps the number of CPU workers one Embed call uses. The default (0) uses GOMAXPROCS, minimizing single-request latency by keeping every core busy — including a spinning fork-join pool that burns idle-core cycles for the duration of each call (~10× the useful CPU at low concurrency). A server saturating many concurrent Embed calls should set a small cap; WithWorkers(1) is fully serial with zero spinning.

The cap governs every fan-out: the packed SIMD path, the attention and GELU phases, and the unpacked fallback matmul (non-amd64, or weight shapes the packer rejects).

type TokenEmbeddings

type TokenEmbeddings struct {
	IDs     []int64     // input token ids, including [CLS]/[SEP] framing
	Vectors [][]float32 // len(IDs) rows of Dim()
}

TokenEmbeddings is one text's token-level output: the final encoder layer's hidden state for every token (ONNX Runtime's last_hidden_state), unpooled and unnormalized.

Vectors' rows are views into a single backing array — one allocation of len(IDs)×Dim() float32s per text (~786 KB for a 512-token input). Retaining one row keeps the whole text's allocation alive, and rows are capacity-clamped so append on a row cannot bleed into its neighbor.