src/vector/embed_registry.rs
// Embedding-**model** registry — the interchangeable set of models behind the
// [`super::store::Embedder`] trait (plan item G3).
//
// Historically the embedder was hard-wired to `jina-embeddings-v2-base-code`
// (768-dim). This registry makes the *model* a selectable axis: a small table
// of known models, each carrying everything the rest of the system needs —
// the HuggingFace repo + file paths (for `build.rs` to fetch), the output
// dimensionality (so the index + warehouse are dimension-aware), the max
// context, and a stable `profile_id` that namespaces the model in the
// `model_profile` warehouse key.
//
// ## Why this file is `include!`-able with **no deps**
//
// `build.rs` needs the *same* registry the crate uses (so the model it fetches
// is exactly the model the embedder loads). build scripts can't depend on the
// crate's modules, so the registry is written as a self-contained, `std`-only
// table that `build.rs` pulls in via `include!`. Keep it that way: no `use` of
// sibling modules, no external crates — only `&'static str` / `usize` data and
// pure lookups.
//
// ## Selection
//
// [`selected`] resolves the active model from `$NORNIR_EMBED_MODEL` (a model
// `id`, e.g. `jina-v2-base-code`), falling back to [`DEFAULT_MODEL_ID`]. The
// **same** env var is read by `build.rs`, so the fetched weights and the
// loaded model never disagree.
//
// ## Re-embed awareness (IMPORTANT)
//
// Each model has a distinct `profile_id`, which flows into
// [`super::store::ModelProfile`] and thus into the warehouse `model_profile`
// key (a hash over `model_name` + weight/tokenizer SHAs + `dim` + …). **Vectors
// from two different models therefore land under different `model_profile`
// keys and never mix** — a search reconstructs only the snapshot whose
// `model_profile` matches the query embedder. The practical consequence:
// **switching `$NORNIR_EMBED_MODEL` requires a fresh `nornir vector index`** at
// the new model before semantic search returns results; old vectors aren't
// reused (different dim, different semantics) and aren't silently blended in.
/// What KIND of thing the weights are — which decides whether any inference runtime
/// is involved at all.
///
/// This is the axis that makes a static embedder expressible. It is deliberately part
/// of the **model** registry rather than the backend selector: whether a forward pass
/// exists is a property of the weights, not of the silicon. `nornir::vector::select`
/// reads it (`select::model_backend`) and lets it outrank the hardware probe.
///
/// `#[repr(u8)]`-free, derive-only, `&'static`-data compatible — see the
/// `include!`-with-no-deps constraint at the top of this file.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelKind {
/// A transformer ONNX graph, executed by tract or onnxruntime. Attention +
/// feed-forward per token; hundreds of GFLOPs per document.
Onnx,
/// A **static token→vector table** (model2vec / potion): the source transformer was
/// run ONCE offline and its output distilled into per-token vectors. Inference is a
/// gather plus a mean-pool — kiloflops, no graph, no runtime.
///
/// The tradeoff is structural and not tunable: a token's vector is identical
/// regardless of its neighbours, so word order and context are gone. `a.copy_to(b)`
/// and `b.copy_to(a)` embed identically.
StaticTable,
}
/// One interchangeable embedding model. All fields are build-time-stable data;
/// the runtime weight/tokenizer SHAs (which also enter the `model_profile`) are
/// resolved separately by `build.rs` and passed via `rustc-env`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EmbedModel {
/// Short, stable selector id (the value of `$NORNIR_EMBED_MODEL`).
pub id: &'static str,
/// Fully-qualified model name (HuggingFace `org/name`), recorded in the
/// `ModelProfile` and used in diagnostics.
pub model_name: &'static str,
/// HuggingFace repo to fetch the weights + tokenizer from.
pub hf_repo: &'static str,
/// Whether the weights are a graph to execute or a table to look up.
pub kind: ModelKind,
/// Repo-relative path of the weights file in `hf_repo`. An ONNX graph for
/// [`ModelKind::Onnx`], a safetensors table for [`ModelKind::StaticTable`].
///
/// Renamed from `onnx_path` 2026-08-14: the field held the local filename
/// `"model.onnx"` implicitly (hardcoded in `build.rs`), which a non-ONNX model
/// cannot honour. See [`Self::weights_local_name`].
pub weights_path: &'static str,
/// The filename the weights are cached under locally. Was hardcoded to
/// `"model.onnx"` in `build.rs`; it is data now because a static table is a
/// `.safetensors` and the loader must find it by the same name the fetch wrote.
pub weights_local_name: &'static str,
/// Repo-relative path of the tokenizer.
pub tokenizer_path: &'static str,
/// Output embedding dimensionality. The index, the warehouse `dim` column,
/// and the pooling all key off this — it is **not** a global constant.
pub dim: usize,
/// Max input tokens (longer inputs are truncated).
pub max_tokens: usize,
/// Local cache sub-directory name (under the model cache root). Keeping one
/// dir per model lets several models coexist in the cache.
pub cache_subdir: &'static str,
/// **The content pin**: sha256 of the weights file and of `tokenizer.json`, or
/// `None` for a model nobody has pinned yet. Two jobs (2026-09-26): `build.rs`
/// REFUSES fetched or cached bytes that differ from a pin (the fetch reads
/// `resolve/main`, which moves), and a SEALED build (the pinned toolchain
/// container, `--network none`) that cannot fetch compiles the pin in as the
/// model's identity instead of dying, so the files are found at run time
/// (`embed_support::model_dir`). Unpinned models still need the fetch.
pub weights_sha256: Option<&'static str>,
/// See [`Self::weights_sha256`].
pub tokenizer_sha256: Option<&'static str>,
}
/// Default model id. Unchanged behaviour: jina-embeddings-v2-base-code, 768-dim.
pub const DEFAULT_MODEL_ID: &str = "jina-v2-base-code";
/// The known embedding models. The first entry is the default and the only one
/// `build.rs` fetches unless `$NORNIR_EMBED_MODEL` selects another. Additional
/// entries make the registry + selection plumbing real and dimension-aware even
/// before their weights are fetched.
pub const MODELS: &[EmbedModel] = &[
// ---- default: code-specialised, 768-dim (QK-LayerNorm "code" arch) ----
EmbedModel {
id: "jina-v2-base-code",
model_name: "jinaai/jina-embeddings-v2-base-code",
hf_repo: "jinaai/jina-embeddings-v2-base-code",
kind: ModelKind::Onnx,
weights_path: "onnx/model.onnx",
weights_local_name: "model.onnx",
tokenizer_path: "tokenizer.json",
dim: 768,
max_tokens: 8192,
cache_subdir: "jina-embeddings-v2-base-code",
// MEASURED 2026-09-26 on loki: sha256 of ~/.cache/nornir/models/jina-embeddings-v2-base-code
// (fetched 2026-06-26), the bytes every nornir built on this box has carried.
weights_sha256: Some("63363fc178428b74620c6f3780cbc7191883fa5c7f84c0945c45eb5c4256733b"),
tokenizer_sha256: Some("b01c78a902aa4facb2f47f95449f48e2f7bbfea5d2472ee2f6ce92323c6f86e5"),
},
// ---- the SAME model, int8-quantized: 4x less bus for an identical contract ----
//
// Added 2026-08-14 because the constraint on `dwarf` is MEMORY BANDWIDTH, not compute,
// and this is the only lever that pulls it without changing anything else.
//
// MEASURED on the box: ort-CPU embeds jina fp32 at 3 chunk/s (production log,
// `nordisk/Skidbladnir [1/19] 2247 chunk(s) … 673.5s`), which puts one pass over the
// 19-member nordisk workspace at 3.5 HOURS as a floor — Skidbladnir is a leaf
// library, not an average member. That is 3.5 h of DRAM at 75-85 C on a box whose
// DIMMs hit their 85 C crit under exactly this load, with no ECC, while writing the
// warehouse. It is not a throughput problem that a faster runtime fixes: ort-CPU is
// already 4.18x tract (`examples/embed_cpu_bakeoff.rs`). It is 611 MB of f32 weights
// crossing a 77 GB/s bus, once per batch, 19 times over.
//
// `onnx/model_quantized.onnx` is 154 MB against fp32's 611 — a 4.0x cut in the one
// number that matters. Everything else is IDENTICAL: same repo, same weights lineage,
// same code specialisation, same 8192-token context, same 768 dims, same tokenizer.
// So chunk boundaries, the HNSW index, `mimir_analyze`'s absolute `cosine >= 0.92`
// duplicates gate and every calibration built on them stay valid — which is exactly
// what a smaller MODEL (potion at 256 dims, or minilm at a 256-token context that
// would silently truncate these ~120-line chunks) does NOT give you.
//
// ⚠ `weights_local_name` is deliberately `model.onnx`, not `model_quantized.onnx`:
// `embed.rs` and `embed_ort.rs` still hardcode `dir.join("model.onnx")` when they
// build a session. The distinct `cache_subdir` is what separates the two on disk, so
// this entry needs NO loader change. If those two call sites ever learn to read
// `weights_local_name` (they should), this can take its true name.
//
// ⚠ NO VNNI ON THIS CPU (Ryzen 5 7535HS is AVX2+FMA only), so int8 will NOT give the
// ~4x COMPUTE win it gives on VNNI silicon — the gain here is the bus, and the bus is
// the bottleneck. Anyone reading this on a VNNI box should expect more, not less.
//
// MEASURED 2026-08-15 on `dwarf`, ort-CPU, 48 chunks of `src/vector`, both arms back to
// back on the shipped `select::load()` path (`examples/quantized_ab.rs`):
//
// fp32 611 MB 2.38 chunk/s 421 ms/chunk load 1652 ms
// q8 154 MB 4.93 chunk/s 203 ms/chunk load 773 ms
// → 2.07x throughput, 2.14x load. Agreement: mean cosine 0.9850, worst 0.9738.
// → RANKING: 48/48 texts still rank themselves first across the two models.
//
// So the 4.0x cut in bytes buys 2.07x wall-clock, not 4x — exactly the VNNI caveat
// above: the bus win is real but compute is no longer free once the weights fit. The
// ~3.5 h nordisk floor becomes ~1.7 h. Semantics survive: nothing re-ranked.
//
// ⚠ STILL UNCALIBRATED: 0.9850 is same-text-ACROSS-models agreement, which is NOT the
// quantity `mimir_analyze`'s absolute `cosine >= 0.92` duplicates gate thresholds
// (that one compares DIFFERENT texts within one model). int8 compresses the similarity
// distribution, so that constant must be re-measured against q8 rows before the gate
// is trusted here — it has not been.
//
// Vectors land in their own space automatically: `ModelProfile::id()` hashes the
// weights SHA, so q8 rows can never be compared against fp32 rows even though the
// dim and model_name match.
EmbedModel {
id: "jina-v2-base-code-q8",
model_name: "jinaai/jina-embeddings-v2-base-code-q8",
hf_repo: "jinaai/jina-embeddings-v2-base-code",
kind: ModelKind::Onnx,
weights_path: "onnx/model_quantized.onnx",
weights_local_name: "model.onnx",
tokenizer_path: "tokenizer.json",
dim: 768,
max_tokens: 8192,
cache_subdir: "jina-embeddings-v2-base-code-q8",
weights_sha256: None,
tokenizer_sha256: None,
},
// ---- alt 1: smaller general-text model, 384-dim ----
// Registry data only by default (not fetched unless selected). Demonstrates
// a *different dimensionality* flowing through the index + warehouse key.
EmbedModel {
id: "minilm-l6-v2",
model_name: "sentence-transformers/all-MiniLM-L6-v2",
hf_repo: "sentence-transformers/all-MiniLM-L6-v2",
kind: ModelKind::Onnx,
weights_path: "onnx/model.onnx",
weights_local_name: "model.onnx",
tokenizer_path: "tokenizer.json",
dim: 384,
max_tokens: 256,
cache_subdir: "all-MiniLM-L6-v2",
weights_sha256: None,
tokenizer_sha256: None,
},
// ---- alt 2: larger general-text model, 768-dim ----
EmbedModel {
id: "bge-base-en-v1.5",
model_name: "BAAI/bge-base-en-v1.5",
hf_repo: "BAAI/bge-base-en-v1.5",
kind: ModelKind::Onnx,
weights_path: "onnx/model.onnx",
weights_local_name: "model.onnx",
tokenizer_path: "tokenizer.json",
dim: 768,
max_tokens: 512,
cache_subdir: "bge-base-en-v1.5",
weights_sha256: None,
tokenizer_sha256: None,
},
// ---- alt 3: STATIC TABLE — model2vec / potion, 256-dim, NO transformer ----
//
// The first non-graph entry, and the reason [`ModelKind`] exists. Every value below
// is read off the model's own published config
// (huggingface.co/minishlab/potion-base-8M/config.json, checked 2026-08-14):
// {"model_type":"model2vec","apply_pca":256,"apply_zipf":true,
// "hidden_dim":256,"seq_length":1000000,"normalize":true,
// "tokenizer_name":"baai/bge-base-en-v1.5"}
//
// `max_tokens` is that `seq_length`, i.e. effectively unbounded — and honestly so: a
// static table has NO positional encoding and therefore no context limit to exceed.
// Nothing is truncated, which is a real behavioural difference from every entry
// above (jina truncates at 8192).
//
// ⚠ `normalize: true` in the config is the model's own instruction, and it matches
// the `ModelProfile.normalize` this crate already writes — but model2vec's weights
// are NOT pre-normalized, so the loader must L2-normalize the pooled vector itself
// or cosine-as-dot-product silently stops holding.
EmbedModel {
id: "potion-base-8m",
model_name: "minishlab/potion-base-8M",
hf_repo: "minishlab/potion-base-8M",
kind: ModelKind::StaticTable,
weights_path: "model.safetensors",
weights_local_name: "model.safetensors",
tokenizer_path: "tokenizer.json",
dim: 256,
max_tokens: 1_000_000,
cache_subdir: "potion-base-8M",
weights_sha256: None,
tokenizer_sha256: None,
},
];
/// Look a model up by its `id`. `None` if no such id is registered.
pub fn by_id(id: &str) -> Option<&'static EmbedModel> {
MODELS.iter().find(|m| m.id == id)
}
/// The default model (first registry entry). Infallible — the table is
/// non-empty by construction.
pub fn default_model() -> &'static EmbedModel {
&MODELS[0]
}
/// The active model's stable registry `id` — the READABLE name that lands in the
/// warehouse `embeddings.model_id` column, next to (never instead of) the opaque
/// `model_profile` sha256.
///
/// Why it exists: `model_profile` is the weights-SHA-sensitive DEDUP key, which
/// is correct for identity but unreadable, so "which model produced these
/// vectors?" was a question the warehouse could not answer. A cross-language
/// cosine threshold is a property of the MODEL, so a calibration has to be keyed
/// on something readable and stable across a re-export. Falls back to
/// [`DEFAULT_MODEL_ID`] when the env names an unknown id, so this never fails a
/// write (the loud error belongs to [`selected`], at load time).
pub fn selected_model_id() -> String {
selected()
.map(|m| m.id)
.unwrap_or(DEFAULT_MODEL_ID)
.to_string()
}
/// True when the selected model is a [`ModelKind::StaticTable`] — i.e. there is no
/// ONNX graph on disk for it and no inference runtime can load it.
///
/// Falls back to the DEFAULT model's kind on an unknown id, matching
/// [`selected_model_id`]: the loud error for a bad id belongs to [`selected`] at load
/// time, and answering "is it static?" with a panic here would move that error to
/// every caller.
pub fn selected_kind_is_static() -> bool {
matches!(
selected().unwrap_or_else(|_| default_model()).kind,
ModelKind::StaticTable
)
}
/// Resolve the env-selected model id (trimmed, lower-cased). Returns the raw
/// value so callers can produce a precise error for an unknown id. `None` ⇒ use
/// the default.
pub fn selected_id_from_env() -> Option<String> {
match std::env::var("NORNIR_EMBED_MODEL") {
Ok(v) if !v.trim().is_empty() => Some(v.trim().to_ascii_lowercase()),
_ => None,
}
}
/// The active model: `$NORNIR_EMBED_MODEL` (by `id`) or the default. Errors with
/// the list of known ids if the env names an unknown model.
pub fn selected() -> Result<&'static EmbedModel, String> {
match selected_id_from_env() {
None => Ok(default_model()),
Some(id) => by_id(&id).ok_or_else(|| {
let known: Vec<&str> = MODELS.iter().map(|m| m.id).collect();
format!(
"NORNIR_EMBED_MODEL=`{id}` is not a known embedding model; \
known ids: {}",
known.join(", ")
)
}),
}
}