src/vector/embed_support.rs
//! Shared embedder support — the **selected** model's identity, [`ModelProfile`]
//! construction, and the tokenize-independent post-processing (mean-pool +
//! L2-normalize) used by **both** the tract (`embed`) and ort (`embed_ort`)
//! backends.
//!
//! Keeping this in one place guarantees the two runtimes produce the *same*
//! `model_profile` and the same vectors for the same text, so embeddings made
//! by one backend are reusable by the other (the warehouse dedup key is
//! backend-agnostic).
//!
//! ## Model is now selectable (plan item G3)
//!
//! The model is **no longer a hard-wired constant**. It is chosen from
//! [`super::embed_registry`] via `$NORNIR_EMBED_MODEL` (default
//! `jina-v2-base-code`, 768-dim). [`active_model`] returns the selected
//! [`EmbedModel`]; its `dim` / `max_tokens` / `model_name` flow everywhere a
//! constant used to. The same env var is read by `build.rs`, so the fetched
//! weights match the loaded model.
//!
//! ## Re-embed awareness
//!
//! The selected model's name + `dim` enter [`ModelProfile`] and therefore the
//! warehouse `model_profile` key. **Different models ⇒ different keys ⇒ vectors
//! never mix**; switching the model requires a fresh `nornir vector index`
//! before semantic search returns hits. See [`super::embed_registry`].
//!
//! Compiled when either embed feature is on.
use super::embed_registry::{self, EmbedModel};
use super::store::ModelProfile;
/// The model the embedder will load/run, resolved from `$NORNIR_EMBED_MODEL`
/// (or the registry default). Panics with a clear message on an unknown id —
/// the same id `build.rs` validated, so a mismatch is a misconfiguration the
/// operator must fix, not something to paper over with the wrong model.
pub fn active_model() -> &'static EmbedModel {
embed_registry::selected().unwrap_or_else(|e| panic!("{e}"))
}
/// Selected model's fully-qualified name (e.g. `jinaai/jina-embeddings-v2-base-code`).
pub fn model_name() -> &'static str {
active_model().model_name
}
/// Selected model's output dimensionality. **Dynamic** — formerly the `DIM`
/// const; now sourced from the registry so a different model carries a
/// different dim through the index and the warehouse `dim` column.
pub fn dim() -> usize {
active_model().dim
}
/// Selected model's max input tokens; longer inputs are truncated.
pub fn max_tokens() -> usize {
active_model().max_tokens
}
/// Where `build.rs` placed the weights of the **selected** model, and their
/// content hashes. build.rs fetches the env-selected model into its own cache
/// subdir and exports that dir + SHAs here, so these always describe
/// [`active_model`].
pub const MODEL_DIR: &str = env!("NORNIR_MODEL_DIR");
pub const WEIGHTS_SHA: &str = env!("NORNIR_MODEL_WEIGHTS_SHA");
pub const TOKENIZER_SHA: &str = env!("NORNIR_MODEL_TOKENIZER_SHA");
/// Resolve the model directory **at runtime** (the baked [`MODEL_DIR`] points at
/// the *building* user's `~/.cache`, which a service user — e.g. `nornir` —
/// can't read; that EACCES is exactly what broke server-side Vector.Search).
/// Precedence:
/// 1. `$NORNIR_MODEL_DIR` (runtime override — the service install can set this),
/// 2. `/opt/nornir/models` if it holds the model (shared, world-readable,
/// mirrors the `/opt/nornir/cuda` convention),
/// 3. the RUNNING user's `~/.cache/nornir/models/<subdir>` if it holds the
/// model — a binary built in the sealed toolchain container baked the
/// container's cache path, which exists nowhere else (2026-09-26),
/// 4. the build-time cache (interactive use on the building user's box).
pub fn model_dir() -> std::path::PathBuf {
if let Ok(d) = std::env::var("NORNIR_MODEL_DIR") {
if !d.is_empty() {
return std::path::PathBuf::from(d);
}
}
let opt = std::path::Path::new("/opt/nornir/models");
if opt.join("tokenizer.json").exists() {
return opt.to_path_buf();
}
if let Some(home) = std::env::var_os("HOME") {
let mine = std::path::PathBuf::from(home)
.join(".cache/nornir/models")
.join(active_model().cache_subdir);
if mine.join("tokenizer.json").exists() {
return mine;
}
}
std::path::PathBuf::from(MODEL_DIR)
}
/// The selected model's identity shared by both backends. `model_name` + `dim`
/// come from the registry; `weights_sha` is the SHA of whichever model artifact
/// build.rs fetched (the ONNX export) — the same file feeds tract and ort, so
/// the profile (and thus the dedup key) matches across backends but **differs
/// across models** (so their vectors never collide in the warehouse).
pub fn profile() -> ModelProfile {
ModelProfile {
model_name: model_name().to_string(),
weights_sha: WEIGHTS_SHA.to_string(),
tokenizer_sha: TOKENIZER_SHA.to_string(),
pooling: "mean".to_string(),
normalize: true,
dim: dim(),
dtype: "f32".to_string(),
}
}
/// Mean-pool a `[1, n_tokens, dim]` row-major hidden state (flattened, so
/// element `(t, d)` is `hidden[t * dim + d]`) over tokens, then L2-normalize.
/// Returns the `dim`-length embedding. `dim` is the selected model's
/// dimensionality ([`dim`]); callers pass it so this stays model-agnostic.
pub fn pool_and_normalize(hidden: &[f32], n_tokens: usize, dim: usize) -> Vec<f32> {
debug_assert_eq!(hidden.len(), n_tokens * dim);
let mut v = vec![0f32; dim];
for t in 0..n_tokens {
let row = &hidden[t * dim..(t + 1) * dim];
for (acc, &x) in v.iter_mut().zip(row) {
*acc += x;
}
}
let inv = 1.0 / n_tokens.max(1) as f32;
for x in &mut v {
*x *= inv;
}
l2_normalize(&mut v);
v
}
/// L2-normalize in place. A zero vector is left unchanged.
pub fn l2_normalize(v: &mut [f32]) {
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in v {
*x /= norm;
}
}
}
/// Truncate raw token ids to `max` (and to at least one token), as `i64` for
/// the ONNX `input_ids` tensor, with an all-ones attention mask. `max` is the
/// selected model's [`max_tokens`]. Variable length — for DYNAMIC EPs (CPU /
/// tract / CUDA) which return one hidden row per input token.
pub fn prepare_tokens(raw: &[u32], max: usize) -> (Vec<i64>, Vec<i64>) {
let n = raw.len().clamp(1, max.max(1));
let ids: Vec<i64> = raw[..n].iter().map(|&x| x as i64).collect();
let mask = vec![1i64; n];
(ids, mask)
}
/// Ids only (no mask) — the batched path masks via row length, so building the
/// throwaway `vec![1i64; n]` [`prepare_tokens`] returns is pure waste there. Same
/// truncation/clamp as [`prepare_tokens`], so the ids are identical.
pub fn prepare_ids(raw: &[u32], max: usize) -> Vec<i64> {
let n = raw.len().clamp(1, max.max(1));
raw[..n].iter().map(|&x| x as i64).collect()
}
/// Fixed-length bucket ladder for STATIC-shape execution providers (MIGraphX,
/// ONNX Runtime 1.24's AMD GPU path). MIGraphX graph-compiles for a constant
/// input shape and reuses that program, so a transformer's variable sequence
/// length must be padded up to one of a small fixed set — AMD's own recommended
/// "bucketing / size-splitting" approach for variable-length inputs (full dynamic
/// shapes are still in design in MIGraphX). One compiled program (one `.mxr`) per
/// bucket actually used; the count is bounded by this ladder.
pub const SHAPE_BUCKETS: &[usize] = &[64, 128, 256, 512, 1024, 2048, 4096, 8192];
/// Smallest bucket `≥ n` (capped at `max`), for padding a static-shape input.
/// `n` over the largest bucket clamps to `max` (the input is then truncated).
pub fn bucket_len(n: usize, max: usize) -> usize {
let cap = max.max(1);
SHAPE_BUCKETS
.iter()
.copied()
.find(|&b| b >= n && b <= cap)
.unwrap_or(cap)
}
/// Default max rows per batched forward — the batch (`B`) dimension cap. Enough
/// to amortize per-call / per-launch overhead and share the GEMM across rows,
/// small enough that one long text in a bucket doesn't blow up the padded work.
/// Override with `$NORNIR_EMBED_BATCH_ROWS`. Shared by BOTH backends (tract CPU
/// and ort GPU/CPU) so the `[B, L]` submit is planned identically.
pub const MAX_BATCH_ROWS: usize = 16;
/// Default padded-token budget per batch (`rows * pad_len`). Caps the total work
/// of a batch so a bucket of long texts forms smaller batches than a bucket of
/// short ones. Override with `$NORNIR_EMBED_BATCH_TOKENS`.
pub const MAX_BATCH_TOKENS: usize = 8192;
/// Rows-per-batch cap that keeps a SMALL input from starving the CPU fork-join
/// fan-out. `plan_batches` packs up to `max_rows` rows per batch, and the tract
/// backend then runs `min(workers, batches)` threads — so an input of 7 texts
/// packed into ONE 7-row batch runs on ONE core while 15 sit idle.
///
/// That is exactly the shape of the budget gate's probe (8 stride-sampled
/// chunks), and it is how the live gate on t14s measured **0.14 chunk/s** and
/// refused every repo (journal 2026-08-26 13:11): the probe timed a
/// single-threaded forward and extrapolated it over a corpus the real run would
/// embed with every core busy. Capping rows at `ceil(n / workers)` splits a
/// small input into one batch per worker (never more than `max_rows`, never
/// zero), so the probe — and any small re-populate — runs at the same
/// parallelism the big run gets. Large inputs are unaffected:
/// `ceil(2258/16) = 142 > 16` keeps the cap at `max_rows`.
pub fn fanout_rows_cap(n: usize, max_rows: usize, workers: usize) -> usize {
max_rows.min(n.div_ceil(workers.max(1))).max(1)
}
/// Read a positive `usize` from `var`, falling back to `default` when unset,
/// empty, or unparseable/zero. The env escape hatch for the batch-size knobs.
pub fn env_usize(var: &str, default: usize) -> usize {
std::env::var(var)
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|&x| x > 0)
.unwrap_or(default)
}
/// Group text indices (given their token lengths) into batches for ONE `[B, L]`
/// forward each: sort ascending by token count (length bucketing → similar
/// lengths batch together, minimal padding), then greedily pack up to `max_rows`
/// rows without the padded token count (`rows * running_max_len`) exceeding
/// `max_tokens`. A single text longer than the budget forms its own batch;
/// `order`-stable within a batch. Shared by both embedder backends so the batched
/// submit is planned identically — the znippy "N workers do all the work" model
/// applied to the accelerator: one real batch, not N locked singletons.
pub fn plan_batches(lens: &[usize], max_rows: usize, max_tokens: usize) -> Vec<Vec<usize>> {
let mut order: Vec<usize> = (0..lens.len()).collect();
order.sort_by_key(|&i| lens[i]);
let max_rows = max_rows.max(1);
let max_tokens = max_tokens.max(1);
let mut batches: Vec<Vec<usize>> = Vec::new();
let mut cur: Vec<usize> = Vec::new();
let mut cur_max = 0usize;
for &i in &order {
let l = lens[i].max(1);
let new_max = cur_max.max(l);
let would_pad = new_max * (cur.len() + 1);
// Start a new batch when adding this row would exceed either cap — but
// never emit an empty batch (a lone oversized row still goes in).
if !cur.is_empty() && (cur.len() >= max_rows || would_pad > max_tokens) {
batches.push(std::mem::take(&mut cur));
cur_max = 0;
}
cur_max = cur_max.max(l);
cur.push(i);
}
if !cur.is_empty() {
batches.push(cur);
}
batches
}
/// Like [`prepare_tokens`] but pads to a FIXED length `l` (pad id `0`, attention
/// mask `0` on the padding, `1` on the `n` real tokens). For static-shape EPs
/// that require a constant input shape. Returns `(ids[l], mask[l], n_real)`.
/// `n_real` is the real token count (≤ `l`), so the caller can mask-pool exactly
/// the real rows and ignore padding. Truncates to `l` if the input is longer.
pub fn prepare_tokens_fixed(raw: &[u32], max: usize, l: usize) -> (Vec<i64>, Vec<i64>, usize) {
let l = l.max(1);
let n = raw.len().clamp(1, l.min(max.max(1)));
let mut ids = vec![0i64; l];
let mut mask = vec![0i64; l];
for i in 0..n {
ids[i] = raw[i] as i64;
mask[i] = 1;
}
(ids, mask, n)
}
/// Mask-aware mean-pool a `[1, n_rows, dim]` flattened hidden state over the
/// tokens whose `mask` entry is non-zero (the real, non-padding tokens), then
/// L2-normalize. For the static-shape / padded path: padding rows contribute
/// nothing, so the padded embedding equals the unpadded one. `mask.len()` is the
/// row count; rows beyond `hidden.len()/dim` are ignored defensively.
pub fn pool_and_normalize_masked(hidden: &[f32], mask: &[i64], dim: usize) -> Vec<f32> {
let rows = hidden.len() / dim;
let mut v = vec![0f32; dim];
let mut count = 0usize;
for t in 0..rows {
if mask.get(t).copied().unwrap_or(0) == 0 {
continue;
}
count += 1;
let row = &hidden[t * dim..(t + 1) * dim];
for (acc, &x) in v.iter_mut().zip(row) {
*acc += x;
}
}
let inv = 1.0 / count.max(1) as f32;
for x in &mut v {
*x *= inv;
}
l2_normalize(&mut v);
v
}
#[cfg(test)]
mod tests {
use super::*;
/// Inject-and-assert: with no env override the default model is selected,
/// and `dim()`/`model_name()`/`profile()` all agree with the registry
/// default (768-dim jina). NOT a "didn't panic" test — it asserts the
/// concrete dim and name the profile reports.
#[test]
fn default_model_profile_matches_registry() {
// No NORNIR_EMBED_MODEL in the test env ⇒ default.
let m = super::embed_registry::default_model();
assert_eq!(dim(), m.dim);
assert_eq!(dim(), 768, "default jina stays 768-dim");
assert_eq!(model_name(), m.model_name);
let p = profile();
assert_eq!(p.dim, m.dim);
assert_eq!(p.model_name, m.model_name);
}
/// Pooling honours the passed `dim` (not a global const), so a 384-dim
/// model pools correctly.
#[test]
fn pool_respects_dim() {
let dim = 4;
// two tokens, each [1,1,1,1] → mean [1,1,1,1] → normalized 0.5 each.
let hidden = vec![1.0f32; 2 * dim];
let v = pool_and_normalize(&hidden, 2, dim);
assert_eq!(v.len(), dim);
for x in &v {
assert!((x - 0.5).abs() < 1e-6, "{x}");
}
}
/// THE 2026-08-26 DEFECT, at plan level: the budget gate's 7-chunk timed
/// sample must plan into one batch PER WORKER (7 single-row batches on a
/// 16-core box), not one 7-row batch on one core. Composing `plan_batches`
/// with `fanout_rows_cap` is exactly what `JinaEmbedder::embed` does, so
/// this pins the shape the probe's measurement runs at.
#[test]
fn small_input_fans_out_one_batch_per_worker() {
// The gate's probe: 7 timed chunks, 16 workers ⇒ cap 1 ⇒ 7 batches.
assert_eq!(fanout_rows_cap(7, MAX_BATCH_ROWS, 16), 1);
let lens = vec![100usize; 7];
let batches = plan_batches(&lens, fanout_rows_cap(7, MAX_BATCH_ROWS, 16), MAX_BATCH_TOKENS);
assert_eq!(batches.len(), 7, "7 probe chunks must form 7 batches: {batches:?}");
// A 40-chunk incremental re-populate on 16 cores: ceil(40/16)=3 rows/batch
// ⇒ at least 14 batches ⇒ every core engaged.
assert_eq!(fanout_rows_cap(40, MAX_BATCH_ROWS, 16), 3);
// Large corpus: the cap must NOT shrink batches (16-row batches stand).
assert_eq!(fanout_rows_cap(2258, MAX_BATCH_ROWS, 16), MAX_BATCH_ROWS);
// Degenerate boxes/inputs never yield a zero cap.
assert_eq!(fanout_rows_cap(1, MAX_BATCH_ROWS, 16), 1);
assert_eq!(fanout_rows_cap(7, MAX_BATCH_ROWS, 1), 7);
assert_eq!(fanout_rows_cap(0, MAX_BATCH_ROWS, 0), 1);
}
/// The static-shape bucket ladder picks the smallest bucket ≥ n, capped at
/// `max`. This is what keeps MIGraphX to a handful of compiled programs.
#[test]
fn bucket_len_picks_smallest_covering_bucket() {
assert_eq!(bucket_len(1, 8192), 64);
assert_eq!(bucket_len(64, 8192), 64);
assert_eq!(bucket_len(65, 8192), 128);
assert_eq!(bucket_len(200, 8192), 256);
assert_eq!(bucket_len(513, 8192), 1024);
assert_eq!(bucket_len(8192, 8192), 8192);
// Over the top bucket / over max → clamp to max (input then truncated).
assert_eq!(bucket_len(99999, 8192), 8192);
// A smaller model max caps the bucket (e.g. a 256-token model).
assert_eq!(bucket_len(300, 256), 256);
}
/// Fixed-shape prep pads to exactly `l`: real tokens up front (mask 1), pad
/// id 0 / mask 0 behind, and reports the real-token count. Truncates if long.
#[test]
fn prepare_tokens_fixed_pads_and_masks() {
let (ids, mask, n_real) = prepare_tokens_fixed(&[7, 8, 9], 8192, 8);
assert_eq!(ids, vec![7, 8, 9, 0, 0, 0, 0, 0]);
assert_eq!(mask, vec![1, 1, 1, 0, 0, 0, 0, 0]);
assert_eq!(n_real, 3);
// Longer than the bucket → truncated to l, all real.
let (ids, mask, n_real) = prepare_tokens_fixed(&[1, 2, 3, 4, 5], 8192, 4);
assert_eq!(ids, vec![1, 2, 3, 4]);
assert_eq!(mask, vec![1, 1, 1, 1]);
assert_eq!(n_real, 4);
}
/// Mask-aware pooling over a PADDED hidden state equals the plain mean over
/// just the real rows — padding rows (mask 0) contribute nothing, so a padded
/// MIGraphX embed matches the unpadded/dynamic one.
#[test]
fn masked_pool_ignores_padding_rows() {
let dim = 4;
// 2 real rows of [1,1,1,1], then 2 padding rows of garbage (mask 0).
let mut hidden = vec![1.0f32; 2 * dim];
hidden.extend(std::iter::repeat(9.0f32).take(2 * dim)); // padding garbage
let mask = vec![1i64, 1, 0, 0];
let masked = pool_and_normalize_masked(&hidden, &mask, dim);
// Same as a plain mean over the 2 real rows (the dynamic path).
let plain = pool_and_normalize(&hidden[..2 * dim], 2, dim);
assert_eq!(masked.len(), dim);
for (a, b) in masked.iter().zip(&plain) {
assert!((a - b).abs() < 1e-6, "masked {a} != plain {b}");
}
// And the garbage padding did NOT leak in (each component is 0.5, not >0.5).
for x in &masked {
assert!((x - 0.5).abs() < 1e-6, "padding leaked: {x}");
}
}
}