src/vector/embed_static.rs
//! **Static token→vector table embedder (model2vec / potion) — no graph, no runtime.**
//!
//! The third [`Embedder`] impl behind the same two-method trait as the tract and ort
//! backends, and the only one that executes **no model**. Where
//! [`super::embed::JinaEmbedder`] runs a 12-layer transformer per text, this looks each
//! token up in a matrix and averages the rows. That is the whole forward pass.
//!
//! ## Why it is ~500× and what that number really is
//!
//! It is not a faster implementation of the same arithmetic — the arithmetic is gone.
//! A transformer forward over `L` tokens costs `O(L²·d)` attention plus `O(L·d²)`
//! feed-forward across every layer: hundreds of GFLOPs per document for jina, and
//! 612 MB of f32 weights streamed through DRAM to do it. This costs `L·d` additions:
//! **kilo**flops, over a 28.8 MB table. The measured end-to-end speedup lands far below
//! the ~10⁶ arithmetic ratio because tokenization then dominates — which is the honest
//! way to read the 500× claim: the model stops being the bottleneck, the tokenizer
//! becomes it.
//!
//! ## What it costs, stated plainly
//!
//! A token's vector is **identical regardless of its neighbours**. There is no
//! attention, so no context, and no positional encoding, so no order:
//! `a.copy_to(b)` and `b.copy_to(a)` produce the same embedding, bit for bit. For
//! lexical/topical similarity over code that is often fine; for anything that turns on
//! argument order or control flow it is not. This is structural and cannot be tuned.
//!
//! ⚠ **It is therefore NOT a drop-in for the duplicates lens.** `mimir_analyze`'s
//! candidate gate is an ABSOLUTE `cosine >= 0.92`, and static-table cosines have a
//! different distribution from contextual ones, so that constant needs re-measuring
//! per model before this is trusted there. A weaker embedder does not error — it
//! returns an empty, green duplicate report.
//!
//! ## Format (hand-parsed on purpose)
//!
//! safetensors is 8 bytes of little-endian header length, then that many bytes of JSON,
//! then packed tensor data. `potion-base-8M` is exactly one tensor — verified against
//! the published file 2026-08-14:
//!
//! ```text
//! {"embeddings":{"dtype":"F32","shape":[29528,256],"data_offsets":[0,30236672]}}
//! ```
//!
//! 29528 × 256 × 4 = 30 236 672, so the rows are contiguous f32-LE and the reader is a
//! chunked cast. Parsing it here rather than adding the `safetensors` crate keeps the
//! dependency count flat (workspace law: pure Rust, and no dep for 40 lines) — the
//! header goes through `serde_json`, which is already a dependency. The tradeoff is
//! named: we support exactly the F32 single-tensor layout and **refuse** anything else
//! by name rather than mis-reading it. potion also publishes F16/I8 variants; those are
//! rejected here until someone needs them.
use anyhow::{Context, Result};
use tokenizers::Tokenizer;
use super::embed_support as es;
use super::store::{Embedder, ModelProfile};
/// The tensor holding the token→vector table in a model2vec safetensors file.
const TABLE_TENSOR: &str = "embeddings";
/// A static-table embedder: one `[vocab, dim]` f32 matrix plus the tokenizer that
/// indexes it.
pub struct StaticEmbedder {
/// Row-major `[vocab, dim]`, so token `t`'s vector is `table[t*dim..(t+1)*dim]`.
table: Vec<f32>,
vocab: usize,
dim: usize,
tokenizer: Tokenizer,
}
impl StaticEmbedder {
/// Load the table + tokenizer from the runtime model dir ([`es::model_dir`], so a
/// service user reads a copy it can actually read rather than the builder's
/// `~/.cache`). The weights filename comes from the registry
/// (`weights_local_name`), which is the same value `build.rs` wrote — one authority
/// for the fetch and the load.
pub fn load() -> Result<Self> {
let dir = es::model_dir();
let model = es::active_model();
let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json"))
.map_err(|e| anyhow::anyhow!("load tokenizer.json: {e}"))?;
let path = dir.join(model.weights_local_name);
let bytes = std::fs::read(&path)
.with_context(|| format!("read static table {}", path.display()))?;
let (table, vocab, dim) = parse_safetensors_f32_matrix(&bytes, TABLE_TENSOR)
.with_context(|| format!("parse {}", path.display()))?;
// The registry's `dim` is what the warehouse column, the HNSW index and the
// ModelProfile all key off. If the file disagrees, every vector written would be
// rejected downstream by `ensure_dim` — so refuse HERE, where the cause is still
// visible, instead of at the first append.
anyhow::ensure!(
dim == model.dim,
"static table {} is {dim}-dim but registry entry `{}` declares dim {} — \
the warehouse, the index and the model_profile all key off the registry, \
so a mismatch would fail at write time with the cause already lost",
path.display(),
model.id,
model.dim
);
Ok(Self {
table,
vocab,
dim,
tokenizer,
})
}
/// Rows in the table (the tokenizer's vocabulary size, as distilled).
pub fn vocab(&self) -> usize {
self.vocab
}
/// Embed one text: tokenize → gather rows → mean-pool → L2-normalize.
///
/// Out-of-range ids are SKIPPED rather than clamped: model2vec distils its own
/// vocabulary, and it can be smaller than the tokenizer's (potion-base-8M is 29 528
/// rows against a 30 522-token BGE vocab). Clamping would silently fold every
/// unknown token onto whatever row happened to be last.
///
/// A text whose every token is out of range yields a ZERO vector, not an error —
/// but it is counted, because a corpus quietly embedding to zeros is exactly the
/// kind of hollow success this codebase keeps finding.
fn embed_one(&self, text: &str) -> Result<(Vec<f32>, bool)> {
let enc = self
.tokenizer
.encode(text, true)
.map_err(|e| anyhow::anyhow!("tokenize: {e}"))?;
let ids = enc.get_ids();
let mut acc = vec![0f32; self.dim];
let mut hits = 0usize;
for &id in ids {
let row = id as usize;
if row >= self.vocab {
continue;
}
let base = row * self.dim;
for (a, w) in acc.iter_mut().zip(&self.table[base..base + self.dim]) {
*a += *w;
}
hits += 1;
}
if hits == 0 {
return Ok((acc, true));
}
let inv = 1.0 / hits as f32;
for a in acc.iter_mut() {
*a *= inv;
}
// model2vec's config says `normalize: true` and its weights are NOT pre-
// normalized, so this is load-bearing: the whole scoring path treats cosine as a
// dot product, and `ModelProfile.normalize` is written as `true` regardless.
es::l2_normalize(&mut acc);
Ok((acc, false))
}
}
impl Embedder for StaticEmbedder {
/// Identity comes from the SAME shared helper the graph backends use, so the
/// `model_profile` is derived one way for every impl. The table's own sha (via
/// `NORNIR_MODEL_WEIGHTS_SHA`) plus `model_name` plus `dim` are what keep these
/// 256-dim vectors from ever being compared against jina's 768-dim ones.
fn profile(&self) -> ModelProfile {
es::profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
// No batching, no padding, no length bucketing: there is no graph whose shape
// could be amortized, so a batch is just a loop. Parallelism is over texts via
// gatling (workspace law: gatling, never rayon); the work per text is tiny, so
// the chunk size is what keeps this from being pure scheduling overhead.
let threads = std::thread::available_parallelism().map_or(1, |n| n.get());
let idx: Vec<usize> = (0..texts.len()).collect();
let out: Vec<Result<(usize, Vec<f32>, bool)>> = if threads <= 1 || texts.len() < 32 {
idx.iter()
.map(|&i| self.embed_one(texts[i]).map(|(v, z)| (i, v, z)))
.collect()
} else {
znippy_zoomies::gatling_forkjoin::gatling_map_balanced(
&idx,
threads,
32,
|&i: &usize| texts[i].len() as u64,
|_, &i| self.embed_one(texts[i]).map(|(v, z)| (i, v, z)),
)
};
let mut vecs: Vec<Vec<f32>> = vec![Vec::new(); texts.len()];
let mut zeros = 0usize;
for r in out {
let (i, v, z) = r?;
if z {
zeros += 1;
}
vecs[i] = v;
}
if zeros > 0 {
eprintln!(
"⚠ static-table embed: {zeros}/{} text(s) had NO in-vocabulary token and \
embedded to a ZERO vector — they will match nothing and rank nowhere. \
Wrong tokenizer for this table, or non-text input.",
texts.len()
);
}
Ok(vecs)
}
}
/// Parse a single named F32 matrix out of a safetensors buffer.
///
/// Returns `(data, rows, cols)`. Every refusal is BY NAME: a truncated file, a missing
/// tensor, a non-F32 dtype, a non-2D shape, or a byte length that disagrees with the
/// shape. Silently accepting any of those would produce vectors that are wrong rather
/// than absent, which is the harder failure to notice.
fn parse_safetensors_f32_matrix(bytes: &[u8], name: &str) -> Result<(Vec<f32>, usize, usize)> {
anyhow::ensure!(bytes.len() >= 8, "not safetensors: {} bytes", bytes.len());
let hlen = u64::from_le_bytes(bytes[..8].try_into().unwrap()) as usize;
let start = 8usize
.checked_add(hlen)
.context("safetensors header length overflows")?;
anyhow::ensure!(
start <= bytes.len(),
"safetensors header claims {hlen} bytes but the file holds {}",
bytes.len().saturating_sub(8)
);
let header: serde_json::Value =
serde_json::from_slice(&bytes[8..start]).context("safetensors header is not JSON")?;
let t = header.get(name).with_context(|| {
let names: Vec<&str> = header
.as_object()
.map(|o| {
o.keys()
.map(String::as_str)
.filter(|k| *k != "__metadata__")
.collect()
})
.unwrap_or_default();
format!("no `{name}` tensor in the table; it holds {names:?}")
})?;
let dtype = t.get("dtype").and_then(|d| d.as_str()).unwrap_or("");
anyhow::ensure!(
dtype == "F32",
"`{name}` is {dtype}, and only F32 is supported here — potion also publishes \
F16/I8 variants, which would need a converting reader rather than a cast"
);
let shape: Vec<usize> = t
.get("shape")
.and_then(|s| s.as_array())
.map(|a| {
a.iter()
.filter_map(|v| v.as_u64())
.map(|v| v as usize)
.collect()
})
.unwrap_or_default();
anyhow::ensure!(
shape.len() == 2,
"`{name}` has shape {shape:?}; a token table must be 2-D [vocab, dim]"
);
let (rows, cols) = (shape[0], shape[1]);
anyhow::ensure!(
rows > 0 && cols > 0,
"`{name}` has an empty axis: {shape:?}"
);
let off: Vec<usize> = t
.get("data_offsets")
.and_then(|s| s.as_array())
.map(|a| {
a.iter()
.filter_map(|v| v.as_u64())
.map(|v| v as usize)
.collect()
})
.unwrap_or_default();
anyhow::ensure!(
off.len() == 2 && off[1] >= off[0],
"`{name}` has bad data_offsets"
);
let (lo, hi) = (start + off[0], start + off[1]);
anyhow::ensure!(
hi <= bytes.len(),
"`{name}` data ends at {hi} but the file is {} bytes — truncated download?",
bytes.len()
);
let want = rows
.checked_mul(cols)
.and_then(|n| n.checked_mul(4))
.context("shape overflows")?;
anyhow::ensure!(
hi - lo == want,
"`{name}` spans {} bytes but shape {shape:?} needs {want}",
hi - lo
);
let data = bytes[lo..hi]
.chunks_exact(4)
.map(|c| f32::from_le_bytes(c.try_into().unwrap()))
.collect();
Ok((data, rows, cols))
}
#[cfg(test)]
mod tests {
use super::*;
/// Build a minimal, VALID safetensors buffer holding one F32 matrix — the fixture
/// every refusal test below mutates one field of.
fn st(name: &str, dtype: &str, shape: &[usize], floats: &[f32]) -> Vec<u8> {
let header = format!(
r#"{{"{name}":{{"dtype":"{dtype}","shape":[{}],"data_offsets":[0,{}]}}}}"#,
shape
.iter()
.map(|d| d.to_string())
.collect::<Vec<_>>()
.join(","),
floats.len() * 4
);
let mut out = (header.len() as u64).to_le_bytes().to_vec();
out.extend(header.as_bytes());
for f in floats {
out.extend(f.to_le_bytes());
}
out
}
/// PURE: the happy path reads the matrix back exactly, in row-major order.
#[test]
fn parses_a_single_f32_matrix() {
let buf = st("embeddings", "F32", &[2, 3], &[1., 2., 3., 4., 5., 6.]);
let (v, rows, cols) = parse_safetensors_f32_matrix(&buf, "embeddings").unwrap();
assert_eq!((rows, cols), (2, 3));
assert_eq!(v, vec![1., 2., 3., 4., 5., 6.]);
}
/// Every malformation is REFUSED BY NAME rather than mis-read. A cast that accepts a
/// truncated or wrongly-typed table yields plausible garbage vectors, which is
/// strictly worse than an error: nothing downstream can tell them from real ones.
#[test]
fn every_malformation_is_refused_by_name() {
let good = st("embeddings", "F32", &[2, 3], &[1., 2., 3., 4., 5., 6.]);
// Wrong tensor name — and the message must LIST what is actually there.
let e = parse_safetensors_f32_matrix(&good, "nope")
.unwrap_err()
.to_string();
assert!(e.contains("no `nope` tensor"), "{e}");
assert!(e.contains("embeddings"), "must name what IS present: {e}");
// F16 is a real potion variant, so this refusal is load-bearing, not theoretical.
let f16 = st("embeddings", "F16", &[2, 3], &[1., 2., 3., 4., 5., 6.]);
let e = parse_safetensors_f32_matrix(&f16, "embeddings")
.unwrap_err()
.to_string();
assert!(e.contains("F16") && e.contains("F32"), "{e}");
// 1-D is not a token table.
let d1 = st("embeddings", "F32", &[6], &[1., 2., 3., 4., 5., 6.]);
let e = parse_safetensors_f32_matrix(&d1, "embeddings")
.unwrap_err()
.to_string();
assert!(e.contains("must be 2-D"), "{e}");
// Shape says 2x3 but only 5 floats are present — the truncated-download case.
let short = st("embeddings", "F32", &[2, 3], &[1., 2., 3., 4., 5.]);
let e = parse_safetensors_f32_matrix(&short, "embeddings")
.unwrap_err()
.to_string();
assert!(e.contains("needs 24") || e.contains("truncated"), "{e}");
// Header length larger than the file.
let mut bad = good.clone();
bad[..8].copy_from_slice(&(9999u64).to_le_bytes());
let e = parse_safetensors_f32_matrix(&bad, "embeddings")
.unwrap_err()
.to_string();
assert!(e.contains("header claims"), "{e}");
// Not safetensors at all.
let e = parse_safetensors_f32_matrix(b"abc", "embeddings")
.unwrap_err()
.to_string();
assert!(e.contains("not safetensors"), "{e}");
}
/// The pooling is a plain mean of the IN-VOCAB rows, then L2-normalized — and an
/// out-of-range id is SKIPPED, not clamped onto the last row.
///
/// Built without a tokenizer so it stays a pure unit test: it drives the same
/// gather/mean/normalize arithmetic `embed_one` performs.
#[test]
fn pooling_skips_out_of_vocab_and_normalizes() {
let dim = 2usize;
let vocab = 2usize;
// row 0 = (3,4) → |5|; row 1 = (0,1)
let table = vec![3., 4., 0., 1.];
let pool = |ids: &[usize]| -> (Vec<f32>, bool) {
let mut acc = vec![0f32; dim];
let mut hits = 0;
for &r in ids {
if r >= vocab {
continue;
}
for (a, w) in acc.iter_mut().zip(&table[r * dim..(r + 1) * dim]) {
*a += *w;
}
hits += 1;
}
if hits == 0 {
return (acc, true);
}
let inv = 1.0 / hits as f32;
for a in acc.iter_mut() {
*a *= inv;
}
es::l2_normalize(&mut acc);
(acc, false)
};
// Row 0 alone normalizes to (0.6, 0.8).
let (v, zero) = pool(&[0]);
assert!(!zero);
assert!(
(v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6,
"{v:?}"
);
// An out-of-vocab id (7) must not change the answer — if it were CLAMPED to the
// last row it would fold row 1 in and this assert would fail.
let (v_oov, _) = pool(&[0, 7]);
assert!(
(v_oov[0] - v[0]).abs() < 1e-6 && (v_oov[1] - v[1]).abs() < 1e-6,
"{v_oov:?}"
);
// All-out-of-vocab is a ZERO vector and reports itself as such.
let (v0, zero) = pool(&[7, 9]);
assert!(zero, "must report the zero-vector case");
assert_eq!(v0, vec![0., 0.]);
// Order does not matter — this IS the model's documented limitation, asserted so
// nobody later mistakes it for a bug in this file.
let (ab, _) = pool(&[0, 1]);
let (ba, _) = pool(&[1, 0]);
assert_eq!(
ab, ba,
"a static table has no positional information, by design"
);
}
}