src/vector/embed.rs
//! tract-onnx embedder (CPU, pure Rust) — `jina-embeddings-v2-base-code`.
//!
//! Implements [`super::store::Embedder`] by running the model's ONNX export on
//! CPU via `tract-onnx`. This is the **pure-Rust** backend (the runtime has no
//! C deps; only the shared `tokenizers` crate does — the accepted exception).
//! The ort backend ([`super::embed_ort`]) is the GPU/CPU alternative; both run
//! the same model and share [`super::embed_support`], so their vectors match.
//!
//! candle's built-in jina model is the *English* architecture and cannot load
//! the code model's QK-LayerNorm weights, so we run the official ONNX export.
//!
//! The graph takes `input_ids` + `attention_mask` (i64 `[batch, sequence]` —
//! **both dims are symbolic/dynamic** in the ONNX export) and outputs the last
//! hidden state `[batch, sequence, dim]` (dim = the **selected** model's dim,
//! 768 for the jina default); we mean-pool + L2-normalize (shared
//! [`super::embed_support::pool_and_normalize`]). The model is selectable via
//! the registry ([`super::embed_registry`]); this backend runs whichever ONNX
//! export `build.rs` fetched.
//!
//! ## Batched forward (CPU throughput)
//!
//! Because the batch dim is dynamic, tract runs one forward at shape
//! `(B, max_len)` just as it runs `(1, n)`. [`JinaEmbedder::embed`] exploits
//! this: it length-buckets the input texts (sort by token count so each batch
//! pads to a similar length — minimal padding waste), forms batches up to
//! [`MAX_BATCH_ROWS`] rows AND a [`MAX_BATCH_TOKENS`] padded-token budget, pads
//! each batch to its own max length, and runs ONE forward per batch. Padding
//! tokens carry attention-mask `0`, and pooling is **mask-aware**
//! ([`es::pool_and_normalize_masked`]) so padding never leaks into a pooled
//! vector. Results are reassembled in the original input order. This is
//! bit-identical to the one-at-a-time path (see the `batched_equals_one_at_a_time`
//! test): the attention mask makes real-token hidden states independent of the
//! padding, and the matmul contraction dim (the model width) is unchanged by
//! batching, so no float ops are reordered.
//!
//! Cargo feature: `embed-tract`.
use std::sync::Arc;
use anyhow::{Context, Result};
use tokenizers::Tokenizer;
use tract_onnx::prelude::*;
use super::embed_support as es;
// Batch planning + the batch-size knobs are SHARED with the ort backend (both
// build one `[B, L]` submit), so they live in `embed_support`.
use super::embed_support::{MAX_BATCH_ROWS, MAX_BATCH_TOKENS, plan_batches};
use super::store::{Embedder, ModelProfile};
type OnnxModel = TypedRunnableModel;
/// A loaded jina-v2-base-code ONNX model + tokenizer (tract / CPU).
pub struct JinaEmbedder {
model: Arc<OnnxModel>,
tokenizer: Tokenizer,
}
impl JinaEmbedder {
/// Load + optimize the ONNX model and tokenizer. The model dir is resolved
/// at runtime ([`es::model_dir`]) so a service user reads a readable copy
/// (`$NORNIR_MODEL_DIR` / `/opt/nornir/models`) rather than the builder's
/// `~/.cache`.
pub fn load() -> Result<Self> {
let dir = es::model_dir();
let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json"))
.map_err(|e| anyhow::anyhow!("load tokenizer.json: {e}"))?;
let onnx = dir.join("model.onnx");
let model = tract_onnx::onnx()
.model_for_path(&onnx)
.with_context(|| format!("load onnx {}", onnx.display()))?
.into_optimized()
.context("optimize onnx graph")?
.into_runnable()
.context("make onnx graph runnable")?;
Ok(Self { model, tokenizer })
}
/// Embed a single string: tokenize → forward → mean-pool → L2-normalize.
fn embed_one(&self, text: &str) -> Result<Vec<f32>> {
let enc = self
.tokenizer
.encode(text, true)
.map_err(|e| anyhow::anyhow!("tokenize: {e}"))?;
let (ids, mask) = es::prepare_tokens(enc.get_ids(), es::max_tokens());
let n = ids.len();
let input_ids = tract_ndarray::Array2::from_shape_vec((1, n), ids)?.into_tensor();
let attn = tract_ndarray::Array2::from_shape_vec((1, n), mask)?.into_tensor();
let outputs = self
.model
.run(tvec!(input_ids.into(), attn.into()))
.context("onnx forward")?;
let hidden = outputs[0]
.to_plain_array_view::<f32>()
.context("read hidden state")?; // [1, n, 768]
let dim = es::dim();
let shape = hidden.shape();
anyhow::ensure!(
shape.len() == 3 && shape[2] == dim,
"unexpected output shape {shape:?} (expected last dim {dim})"
);
let flat = hidden.as_slice().context("hidden state not contiguous")?;
Ok(es::pool_and_normalize(flat, n, dim))
}
/// Tokenize one text into truncated `i64` ids (length ∈ `[1, max_tokens]`),
/// the same ids `embed_one` would feed — so a batched forward built from
/// these is equivalent to running `embed_one` per text.
fn tokenize(&self, text: &str) -> Result<Vec<i64>> {
let enc = self
.tokenizer
.encode(text, true)
.map_err(|e| anyhow::anyhow!("tokenize: {e}"))?;
Ok(es::prepare_ids(enc.get_ids(), es::max_tokens()))
}
/// Run ONE forward over a batch of pre-tokenized texts (`toks[i]` for each
/// `i` in `batch`), padding every row to the batch's longest row. Returns
/// `(original_index, embedding)` per row. Padding tokens (id `0`, mask `0`)
/// are ignored by mask-aware pooling, so each row's vector equals its
/// one-at-a-time `embed_one` result.
fn forward_batch(&self, batch: &[usize], toks: &[Vec<i64>]) -> Result<Vec<(usize, Vec<f32>)>> {
let b = batch.len();
debug_assert!(b > 0, "empty batch");
let pad_len = batch
.iter()
.map(|&i| toks[i].len())
.max()
.unwrap_or(1)
.max(1);
let dim = es::dim();
// Build the (B, pad_len) `input_ids` + `attention_mask` tensors by writing the
// real tokens DIRECTLY into tract's aligned buffers — no intermediate `Vec`
// and no alloc-then-copy into the tensor (`into_tensor` copies). `Tensor::zero`
// pre-zeros both, so padding ids stay 0 and padding mask stays 0 for free; we
// only touch the real front cells of each row.
let mut ids_t = Tensor::zero::<i64>(&[b, pad_len])?;
let mut mask_t = Tensor::zero::<i64>(&[b, pad_len])?;
{
// SAFETY: both tensors were just created contiguous, dtype i64, length
// exactly b*pad_len; we hold the sole &mut, so the unchecked slices are
// sound (tract has no safe checked `as_slice_mut`, only the pointer form).
let ids_s = unsafe { ids_t.as_slice_mut_unchecked::<i64>() };
let mask_s = unsafe { mask_t.as_slice_mut_unchecked::<i64>() };
for (row, &i) in batch.iter().enumerate() {
let base = row * pad_len;
for (j, &id) in toks[i].iter().enumerate() {
ids_s[base + j] = id;
mask_s[base + j] = 1;
}
}
}
let outputs = self
.model
.run(tvec!(ids_t.into(), mask_t.into()))
.context("onnx batched forward")?;
let hidden = outputs[0]
.to_plain_array_view::<f32>()
.context("read hidden state")?; // [b, pad_len, dim]
let shape = hidden.shape();
anyhow::ensure!(
shape.len() == 3 && shape[0] == b && shape[1] == pad_len && shape[2] == dim,
"unexpected batched output shape {shape:?} (expected [{b}, {pad_len}, {dim}])"
);
let flat = hidden.as_slice().context("hidden state not contiguous")?;
// Pool over the FRONT `n_real` rows of each row's hidden block. Real tokens are
// front-contiguous and the attention mask zeroed padding in the forward, so
// this is bit-identical to the mask-aware pool — and needs no mask buffer at
// pooling time (so `mask_t` was free to move into the forward above).
let mut out = Vec::with_capacity(b);
for (row, &i) in batch.iter().enumerate() {
let n_real = toks[i].len().max(1);
let start = row * pad_len * dim;
let row_hidden = &flat[start..start + n_real * dim];
out.push((i, es::pool_and_normalize(row_hidden, n_real, dim)));
}
Ok(out)
}
}
impl Embedder for JinaEmbedder {
fn profile(&self) -> ModelProfile {
es::profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
// tract runs one forward per call single-threaded. Two levers give CPU
// throughput here: (1) each forward runs a whole BATCH of texts at shape
// (B, max_len) — the ONNX batch dim is dynamic, so this is one GEMM-heavy
// pass instead of B tiny ones; (2) batches are spread across cores. The
// model + tokenizer are shared read-only across scoped threads (no clone
// of the graph). Output is reassembled in the original `texts` order, so
// it is identical to the one-at-a-time map.
let n = texts.len();
if n == 0 {
return Ok(Vec::new());
}
if n == 1 {
return Ok(vec![self.embed_one(texts[0])?]);
}
// Tokenize once (cheap vs the forward); lengths drive the bucketing.
let toks: Vec<Vec<i64>> = texts
.iter()
.map(|t| self.tokenize(t))
.collect::<Result<_>>()?;
let lens: Vec<usize> = toks.iter().map(|t| t.len()).collect();
// THE MEMORY CEILING (gap 3): the plan — fan-out width × per-batch token
// cap — comes from `vector::memory`, which fits it under `--max-memory` /
// `NORNIR_EMBED_MAX_MEMORY_MB` when one is declared and otherwise hands
// back the legacy knobs (every core, `NORNIR_EMBED_BATCH_ROWS/TOKENS`).
// Every in-flight forward owns its own `[B, L]` activations, so `workers`
// is what peak RSS is linear in; `taskset` used to be the only lever.
let workers_available = std::thread::available_parallelism()
.map(|x| x.get())
.unwrap_or(1);
let limits = super::memory::EmbedLimits::from_env(
workers_available,
MAX_BATCH_ROWS,
MAX_BATCH_TOKENS,
);
limits.announce_once();
let (max_rows, max_tokens, workers) = (limits.max_rows, limits.max_tokens, limits.workers);
// Small inputs split into one batch per worker (`fanout_rows_cap`), so a
// handful of texts — the budget gate's probe, an incremental re-populate —
// engages every core instead of packing into ONE single-threaded batch.
// This is the 2026-08-26 defect: the gate timed its 7-chunk sample as one
// 1-thread forward (0.14 chunk/s live on t14s) and refused corpora the
// 16-wide real run embeds several times faster. Large inputs are
// unchanged (the cap stays at `max_rows`).
let batches = plan_batches(&lens, es::fanout_rows_cap(n, max_rows, workers), max_tokens);
let threads = workers.min(batches.len());
// Run each batch as one forward; distribute batches across cores,
// heaviest-first by padded-token weight so a few big batches don't
// strand a worker.
let batch_results: Vec<Result<Vec<(usize, Vec<f32>)>>> = if threads <= 1 {
batches
.iter()
.map(|b| self.forward_batch(b, &toks))
.collect()
} else {
znippy_zoomies::gatling_forkjoin::gatling_map_balanced(
&batches,
threads,
1,
|b: &Vec<usize>| {
let pad = b.iter().map(|&i| lens[i]).max().unwrap_or(0);
(b.len() * pad.max(1)) as u64
},
|_, b| self.forward_batch(b, &toks),
)
};
// Scatter each batch's rows back to their original positions.
let mut out: Vec<Vec<f32>> = vec![Vec::new(); n];
for br in batch_results {
for (i, v) in br? {
out[i] = v;
}
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Bucketing is pure over token lengths — no model needed. It sorts ascending
/// (so similar lengths batch together), caps rows at `max_rows`, caps padded
/// tokens (`rows * running_max_len`) at `max_tokens`, covers every index
/// exactly once, and lets a lone oversized text form its own batch.
#[test]
fn plan_batches_buckets_by_length_and_respects_caps() {
// 7 texts, mixed lengths. rows cap 3, token budget 12.
let lens = vec![10usize, 2, 9, 1, 3, 8, 2];
let batches = plan_batches(&lens, 3, 12);
// Every index appears exactly once.
let mut seen: Vec<usize> = batches.iter().flatten().copied().collect();
seen.sort_unstable();
assert_eq!(seen, (0..lens.len()).collect::<Vec<_>>());
for b in &batches {
assert!(!b.is_empty(), "no empty batch");
assert!(b.len() <= 3, "row cap honored: {b:?}");
let pad = b.iter().map(|&i| lens[i]).max().unwrap();
// Either within budget, or a single row that alone exceeds it.
assert!(
pad * b.len() <= 12 || b.len() == 1,
"token budget honored (or lone oversized row): pad {pad} * {} rows",
b.len()
);
// Ascending-sorted input ⇒ within a batch lengths are non-decreasing.
let ls: Vec<usize> = b.iter().map(|&i| lens[i]).collect();
let mut sorted = ls.clone();
sorted.sort_unstable();
assert_eq!(ls, sorted, "batch lengths are contiguous in sorted order");
}
// The lone long text (len 10) can't share a 12-token budget with anyone
// (2 rows * 10 = 20 > 12), so it sits in a 1-row batch.
assert!(
batches.iter().any(|b| b.len() == 1 && lens[b[0]] == 10),
"the len-10 text forms its own batch: {batches:?}"
);
}
/// EXACTNESS GUARD (real model). Embedding a set of varied-length texts
/// (including duplicates and lengths that force multi-row buckets + padding)
/// via the batched [`JinaEmbedder::embed`] MUST equal running `embed_one`
/// one-at-a-time — bit-identical (`to_bits()`). The attention mask makes each
/// real token's hidden state independent of padding, and batching leaves the
/// matmul contraction dim unchanged, so no float ops are reordered. A tiny
/// batch budget forces several multi-row, padded batches. Ignored by default
/// (needs the staged model); run with:
/// `cargo test --features embed-tract -- --ignored batched_equals_one_at_a_time`
#[test]
#[ignore = "loads the real ONNX model (needs build-time weight cache); run with --features embed-tract -- --ignored"]
fn batched_equals_one_at_a_time() {
let e = JinaEmbedder::load().expect("load model");
// Varied lengths incl. exact duplicates and near-duplicates; short and
// long so buckets both pad and split.
let texts: Vec<String> = vec![
"fn a() {}".into(),
"fn add(a: i32, b: i32) -> i32 { a + b }".into(),
"x".into(),
"the quick brown fox jumps over the lazy dog".into(),
"fn add(a: i32, b: i32) -> i32 { a + b }".into(), // dup of [1]
"pub struct Foo { bar: Vec<u8>, baz: Option<String>, qux: [f64; 16] }".into(),
"y".into(),
"// a fairly long comment line that tokenizes to a good handful of tokens indeed"
.into(),
"let z = compute_the_thing(alpha, beta, gamma, delta, epsilon, zeta);".into(),
"fn a() {}".into(), // dup of [0]
];
// Reference: strictly one-at-a-time forwards at shape (1, n).
let reference: Vec<Vec<f32>> = texts.iter().map(|t| e.embed_one(t).unwrap()).collect();
// Force multi-row, padded batches: tiny caps.
// TODO: Audit that the environment access only happens in single-threaded code.
unsafe {
std::env::set_var("NORNIR_EMBED_BATCH_ROWS", "4");
std::env::set_var("NORNIR_EMBED_BATCH_TOKENS", "64");
}
let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
let batched = e.embed(&refs).unwrap();
assert_eq!(batched.len(), reference.len());
let mut max_abs = 0.0f32;
let mut exact = true;
for (idx, (r, b)) in reference.iter().zip(&batched).enumerate() {
assert_eq!(r.len(), b.len(), "dim mismatch at {idx}");
for (k, (&rv, &bv)) in r.iter().zip(b).enumerate() {
if rv.to_bits() != bv.to_bits() {
exact = false;
}
max_abs = max_abs.max((rv - bv).abs());
// Hard cap even if not bit-identical: never accept larger drift.
assert!(
(rv - bv).abs() <= 1e-5,
"batched vs one-at-a-time drift too large at text {idx} dim {k}: \
{rv} vs {bv}"
);
}
}
eprintln!("batched_equals_one_at_a_time: bit_identical={exact} max_abs_diff={max_abs:e}");
// Duplicate texts must produce identical vectors to each other.
assert_eq!(batched[0], batched[9], "dup text [0]==[9]");
assert_eq!(batched[1], batched[4], "dup text [1]==[4]");
}
/// Loads the real ONNX model (needs the build-time weight cache). Ignored
/// by default so the normal test run stays fast and offline; run with
/// `cargo test --features embed-tract -- --ignored embed`.
#[test]
#[ignore = "loads the real ONNX model (needs build-time weight cache); run with --features embed-tract -- --ignored"]
fn loads_and_embeds() {
let e = JinaEmbedder::load().expect("load model");
let p = e.profile();
assert_eq!(p.dim, es::dim());
assert_eq!(p.model_name, es::model_name());
let v = e.embed(&["fn main() {}"]).unwrap();
assert_eq!(v.len(), 1);
assert_eq!(v[0].len(), es::dim());
let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
assert!((norm - 1.0).abs() < 1e-3, "norm {norm}");
// Code semantics: two equivalent Rust fns are closer to each other
// than either is to unrelated prose.
let a = e
.embed(&["fn add(a: i32, b: i32) -> i32 { a + b }"])
.unwrap();
let b = e
.embed(&["pub fn sum(x: i32, y: i32) -> i32 { x + y }"])
.unwrap();
let c = e
.embed(&["the quick brown fox jumps over the lazy dog"])
.unwrap();
let dot = |x: &[f32], y: &[f32]| x.iter().zip(y).map(|(p, q)| p * q).sum::<f32>();
assert!(
dot(&a[0], &b[0]) > dot(&a[0], &c[0]),
"two fns ({}) closer than fn-vs-prose ({})",
dot(&a[0], &b[0]),
dot(&a[0], &c[0])
);
}
/// Full pipeline with the real model: index 3 code files into the
/// warehouse, then a natural-language query retrieves the right one.
#[test]
#[ignore = "real model + full warehouse index/search pipeline; run with --features embed-tract -- --ignored"]
fn end_to_end_semantic_search() {
use crate::vector::chunk::ChunkOptions;
use crate::vector::store::{RepoRef, index_repo, search};
use crate::warehouse::iceberg::IcebergWarehouse;
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
let embedder = JinaEmbedder::load().unwrap();
let files = vec![
(
"math.rs".to_string(),
"pub fn add(a: i32, b: i32) -> i32 { a + b }".to_string(),
),
(
"io.rs".to_string(),
"fn read_file(path: &str) -> std::io::Result<String> { std::fs::read_to_string(path) }".to_string(),
),
(
"net.rs".to_string(),
"async fn fetch(url: &str) -> Result<String> { http_get(url).await }".to_string(),
),
];
let snap = index_repo(
&wh,
&RepoRef {
workspace: "ws",
repo: "demo",
git_sha: "sha1",
branch: "main",
complete: true,
},
&files,
&ChunkOptions::default(),
&embedder,
)
.unwrap();
assert_eq!(snap.new_vectors, 3);
let mp = embedder.profile().id();
let q = embedder
.embed(&["a function that adds two integers together"])
.unwrap();
let hits = search(&wh, "demo", Some("sha1"), &mp, &q[0], 3).unwrap();
assert_eq!(
hits[0].1.file, "math.rs",
"NL query about adding integers should retrieve the add fn"
);
}
}