chore: import upstream snapshot with attribution

crates/jcode-embedding/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "jcode-embedding"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+name = "jcode_embedding"
+path = "src/lib.rs"
+
+[dependencies]
+anyhow = "1"
+reqwest = { version = "0.12", default-features = false, features = ["blocking", "charset", "http2", "system-proxy", "rustls-tls", "rustls-tls-native-roots"] }
+tokenizers = { version = "0.21", default-features = false, features = ["onig"] }
+tract-hir = "0.21"
+tract-onnx = "0.21"

crates/jcode-embedding/src/lib.rs

@@ -0,0 +1,582 @@
+use anyhow::{Context, Result};
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
+use std::io::Write;
+use std::path::Path;
+use tokenizers::Tokenizer;
+use tract_hir::infer::Factoid as _;
+use tract_hir::prelude::*;
+
+pub const MODEL_NAME: &str = "all-MiniLM-L6-v2";
+type RunnableEmbeddingModel =
+    SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
+
+#[derive(Debug)]
+struct TopKItem<T> {
+    score: f32,
+    ordinal: usize,
+    value: T,
+}
+
+impl<T> PartialEq for TopKItem<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.score.to_bits() == other.score.to_bits() && self.ordinal == other.ordinal
+    }
+}
+
+impl<T> Eq for TopKItem<T> {}
+
+impl<T> PartialOrd for TopKItem<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<T> Ord for TopKItem<T> {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.score
+            .total_cmp(&other.score)
+            .then_with(|| self.ordinal.cmp(&other.ordinal))
+    }
+}
+
+fn top_k_scored<T, I>(items: I, limit: usize) -> Vec<(T, f32)>
+where
+    I: IntoIterator<Item = (T, f32)>,
+{
+    if limit == 0 {
+        return Vec::new();
+    }
+
+    let mut heap: BinaryHeap<Reverse<TopKItem<T>>> = BinaryHeap::new();
+    for (ordinal, (value, score)) in items.into_iter().enumerate() {
+        let candidate = Reverse(TopKItem {
+            score,
+            ordinal,
+            value,
+        });
+
+        if heap.len() < limit {
+            heap.push(candidate);
+            continue;
+        }
+
+        let replace = heap
+            .peek()
+            .map(|smallest| score > smallest.0.score)
+            .unwrap_or(false);
+        if replace {
+            heap.pop();
+            heap.push(candidate);
+        }
+    }
+
+    let mut results: Vec<_> = heap
+        .into_iter()
+        .map(|Reverse(item)| (item.value, item.score, item.ordinal))
+        .collect();
+    results.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.2.cmp(&b.2)));
+    results
+        .into_iter()
+        .map(|(value, score, _)| (value, score))
+        .collect()
+}
+const EMBEDDING_DIM: usize = 384;
+const MAX_SEQ_LENGTH: usize = 256;
+
+const MODEL_URL: &str =
+    "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/onnx/model.onnx";
+const TOKENIZER_URL: &str =
+    "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json";
+
+pub type EmbeddingVec = Vec<f32>;
+
+pub struct Embedder {
+    model: RunnableEmbeddingModel,
+    tokenizer: Tokenizer,
+    /// Per-input binding plan: (role, dtype) in the model's DECLARED input order.
+    /// Exporters differ in both input ORDER (MiniLM puts input_ids first; e5/bge
+    /// put attention_mask first) and DTYPE (f32 vs i64), so we bind by name and
+    /// feed each input its model-declared dtype instead of assuming a position.
+    input_plan: Vec<(InputRole, DatumType)>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+enum InputRole {
+    InputIds,
+    AttentionMask,
+    TokenTypeIds,
+}
+
+fn classify_input(name: &str) -> InputRole {
+    let n = name.to_ascii_lowercase();
+    if n.contains("attention") || n.contains("mask") {
+        InputRole::AttentionMask
+    } else if n.contains("token_type") || n.contains("segment") {
+        InputRole::TokenTypeIds
+    } else {
+        InputRole::InputIds
+    }
+}
+
+impl Embedder {
+    pub fn load_from_dir(model_dir: &Path) -> Result<Self> {
+        let model_path = model_dir.join("model.onnx");
+        let tokenizer_path = model_dir.join("tokenizer.json");
+
+        if !model_path.exists() || !tokenizer_path.exists() {
+            download_model(model_dir)?;
+        }
+
+        let tokenizer = Tokenizer::from_file(&tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
+
+        let raw = tract_onnx::onnx()
+            .model_for_path(&model_path)
+            .context("Failed to load ONNX model")?;
+
+        // Determine each input's role (by name) and dtype (declared, else i64).
+        let input_outlets = raw
+            .input_outlets()
+            .context("Failed to read model input outlets")?
+            .to_vec();
+        let mut input_plan: Vec<(InputRole, DatumType)> = Vec::with_capacity(input_outlets.len());
+        for (ix, outlet) in input_outlets.iter().enumerate() {
+            let role = classify_input(&raw.node(outlet.node).name);
+            let dt = raw
+                .input_fact(ix)
+                .ok()
+                .and_then(|f| f.datum_type.concretize())
+                .unwrap_or(DatumType::I64);
+            input_plan.push((role, dt));
+        }
+
+        let mut model = raw;
+        for (ix, (_, dt)) in input_plan.iter().enumerate() {
+            model = model.with_input_fact(ix, InferenceFact::dt_shape(*dt, [1, MAX_SEQ_LENGTH]))?;
+        }
+        let model = model
+            .into_optimized()
+            .context("Failed to optimize model")?
+            .into_runnable()
+            .context("Failed to make model runnable")?;
+
+        Ok(Self {
+            model,
+            tokenizer,
+            input_plan,
+        })
+    }
+
+    pub fn embed(&self, text: &str) -> Result<EmbeddingVec> {
+        let encoding = self
+            .tokenizer
+            .encode(text, true)
+            .map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;
+
+        let mut input_ids = vec![0i64; MAX_SEQ_LENGTH];
+        let mut attention_mask = vec![0i64; MAX_SEQ_LENGTH];
+        let token_type_ids = vec![0i64; MAX_SEQ_LENGTH];
+
+        let ids = encoding.get_ids();
+        let len = ids.len().min(MAX_SEQ_LENGTH);
+
+        for i in 0..len {
+            input_ids[i] = ids[i] as i64;
+            attention_mask[i] = 1;
+        }
+
+        // Build each input tensor by role, cast to the model's declared dtype.
+        let make = |data: &[i64], dt: DatumType| -> Result<Tensor> {
+            let t: Tensor =
+                tract_ndarray::Array2::from_shape_vec((1, MAX_SEQ_LENGTH), data.to_vec())?.into();
+            Ok(t.cast_to_dt(dt)?.into_owned())
+        };
+
+        let mut inputs: TVec<TValue> = tvec![];
+        for (role, dt) in &self.input_plan {
+            let data: &[i64] = match role {
+                InputRole::InputIds => &input_ids,
+                InputRole::AttentionMask => &attention_mask,
+                InputRole::TokenTypeIds => &token_type_ids,
+            };
+            inputs.push(make(data, *dt)?.into());
+        }
+
+        let outputs = self.model.run(inputs)?;
+
+        let output = outputs[0].to_array_view::<f32>()?.to_owned();
+
+        let shape = output.shape();
+        if shape.len() == 3 {
+            let seq_len = shape[1];
+            let hidden_dim = shape[2];
+            let mut embedding = vec![0f32; hidden_dim];
+
+            let valid_tokens = len.min(seq_len);
+
+            for i in 0..valid_tokens {
+                for j in 0..hidden_dim {
+                    embedding[j] += output[[0, i, j]];
+                }
+            }
+
+            for val in &mut embedding {
+                *val /= valid_tokens.max(1) as f32;
+            }
+
+            let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
+            if norm > 0.0 {
+                for val in &mut embedding {
+                    *val /= norm;
+                }
+            }
+
+            Ok(embedding)
+        } else {
+            anyhow::bail!("Unexpected output shape: {:?}", shape);
+        }
+    }
+
+    pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingVec>> {
+        texts.iter().map(|t| self.embed(t)).collect()
+    }
+}
+
+/// Build the run() inputs for a BERT-style model from token ids, honoring the
+/// model's declared input order/role/dtype. Shared by Embedder and CrossEncoder.
+fn build_bert_inputs(
+    input_plan: &[(InputRole, DatumType)],
+    ids: &[u32],
+    type_ids: &[u32],
+) -> Result<TVec<TValue>> {
+    let mut input_ids = vec![0i64; MAX_SEQ_LENGTH];
+    let mut attention_mask = vec![0i64; MAX_SEQ_LENGTH];
+    let mut token_type_ids = vec![0i64; MAX_SEQ_LENGTH];
+    let len = ids.len().min(MAX_SEQ_LENGTH);
+    for i in 0..len {
+        input_ids[i] = ids[i] as i64;
+        attention_mask[i] = 1;
+        if i < type_ids.len() {
+            token_type_ids[i] = type_ids[i] as i64;
+        }
+    }
+    let make = |data: &[i64], dt: DatumType| -> Result<Tensor> {
+        let t: Tensor =
+            tract_ndarray::Array2::from_shape_vec((1, MAX_SEQ_LENGTH), data.to_vec())?.into();
+        Ok(t.cast_to_dt(dt)?.into_owned())
+    };
+    let mut inputs: TVec<TValue> = tvec![];
+    for (role, dt) in input_plan {
+        let data: &[i64] = match role {
+            InputRole::InputIds => &input_ids,
+            InputRole::AttentionMask => &attention_mask,
+            InputRole::TokenTypeIds => &token_type_ids,
+        };
+        inputs.push(make(data, *dt)?.into());
+    }
+    Ok(inputs)
+}
+
+/// A cross-encoder reranker: scores a (query, passage) pair jointly and returns
+/// a single relevance logit. Used to reorder a candidate set after first-stage
+/// retrieval (recall-5). Higher score = more relevant.
+pub struct CrossEncoder {
+    model: RunnableEmbeddingModel,
+    tokenizer: Tokenizer,
+    input_plan: Vec<(InputRole, DatumType)>,
+}
+
+impl CrossEncoder {
+    pub fn load_from_dir(model_dir: &Path) -> Result<Self> {
+        let model_path = model_dir.join("model.onnx");
+        let tokenizer_path = model_dir.join("tokenizer.json");
+        let tokenizer = Tokenizer::from_file(&tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
+
+        let raw = tract_onnx::onnx()
+            .model_for_path(&model_path)
+            .context("Failed to load cross-encoder ONNX model")?;
+        let input_outlets = raw.input_outlets().context("read input outlets")?.to_vec();
+        let mut input_plan: Vec<(InputRole, DatumType)> = Vec::with_capacity(input_outlets.len());
+        for (ix, outlet) in input_outlets.iter().enumerate() {
+            let role = classify_input(&raw.node(outlet.node).name);
+            let dt = raw
+                .input_fact(ix)
+                .ok()
+                .and_then(|f| f.datum_type.concretize())
+                .unwrap_or(DatumType::I64);
+            input_plan.push((role, dt));
+        }
+        let mut model = raw;
+        for (ix, (_, dt)) in input_plan.iter().enumerate() {
+            model = model.with_input_fact(ix, InferenceFact::dt_shape(*dt, [1, MAX_SEQ_LENGTH]))?;
+        }
+        let model = model
+            .into_optimized()
+            .context("optimize cross-encoder")?
+            .into_runnable()
+            .context("make cross-encoder runnable")?;
+        Ok(Self {
+            model,
+            tokenizer,
+            input_plan,
+        })
+    }
+
+    /// Relevance score for a (query, passage) pair. Higher = more relevant.
+    pub fn score(&self, query: &str, passage: &str) -> Result<f32> {
+        let encoding = self
+            .tokenizer
+            .encode((query, passage), true)
+            .map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;
+        let inputs = build_bert_inputs(
+            &self.input_plan,
+            encoding.get_ids(),
+            encoding.get_type_ids(),
+        )?;
+        let outputs = self.model.run(inputs)?;
+        let view = outputs[0].to_array_view::<f32>()?;
+        // logits shape is [1, 1] (relevance) or [1, N]; take the first/primary.
+        view.iter()
+            .next()
+            .copied()
+            .ok_or_else(|| anyhow::anyhow!("empty cross-encoder output"))
+    }
+
+    /// Rerank `(id, text)` candidates by cross-encoder score against `query`.
+    /// Returns ids sorted by descending relevance.
+    pub fn rerank(
+        &self,
+        query: &str,
+        candidates: &[(String, String)],
+    ) -> Result<Vec<(String, f32)>> {
+        let mut scored: Vec<(String, f32)> = Vec::with_capacity(candidates.len());
+        for (id, text) in candidates {
+            let s = self.score(query, text)?;
+            scored.push((id.clone(), s));
+        }
+        scored.sort_by(|a, b| b.1.total_cmp(&a.1));
+        Ok(scored)
+    }
+}
+
+pub const fn embedding_dim() -> usize {
+    EMBEDDING_DIM
+}
+
+pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
+    if a.len() != b.len() || a.is_empty() {
+        return 0.0;
+    }
+
+    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
+    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
+    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+    if norm_a == 0.0 || norm_b == 0.0 {
+        return 0.0;
+    }
+
+    dot / (norm_a * norm_b)
+}
+
+pub fn batch_cosine_similarity(query: &[f32], candidates: &[&[f32]]) -> Vec<f32> {
+    let dim = query.len();
+    if dim == 0 || candidates.is_empty() {
+        return vec![0.0; candidates.len()];
+    }
+
+    candidates
+        .iter()
+        .map(|c| {
+            if c.len() != dim {
+                0.0
+            } else {
+                c.iter().zip(query.iter()).map(|(a, b)| a * b).sum()
+            }
+        })
+        .collect()
+}
+
+pub fn find_similar(
+    query: &[f32],
+    candidates: &[EmbeddingVec],
+    threshold: f32,
+    top_k: usize,
+) -> Vec<(usize, f32)> {
+    let refs: Vec<&[f32]> = candidates.iter().map(|v| v.as_slice()).collect();
+    let scores = batch_cosine_similarity(query, &refs);
+
+    top_k_scored(
+        scores
+            .into_iter()
+            .enumerate()
+            .filter(|(_, score)| *score >= threshold),
+        top_k,
+    )
+}
+
+pub fn is_model_available(model_dir: &Path) -> bool {
+    model_dir.join("model.onnx").exists() && model_dir.join("tokenizer.json").exists()
+}
+
+fn download_model(model_dir: &Path) -> Result<()> {
+    let model_dir = model_dir.to_path_buf();
+    match std::thread::spawn(move || download_model_blocking(&model_dir)).join() {
+        Ok(result) => result,
+        Err(panic) => {
+            let panic_msg = if let Some(msg) = panic.downcast_ref::<&str>() {
+                (*msg).to_string()
+            } else if let Some(msg) = panic.downcast_ref::<String>() {
+                msg.clone()
+            } else {
+                "unknown panic payload".to_string()
+            };
+            anyhow::bail!("Embedding model download thread panicked: {}", panic_msg);
+        }
+    }
+}
+
+fn download_model_blocking(model_dir: &Path) -> Result<()> {
+    let client = reqwest::blocking::Client::builder()
+        .user_agent(concat!("jcode-embedding/", env!("CARGO_PKG_VERSION")))
+        .timeout(std::time::Duration::from_secs(300))
+        .build()?;
+
+    std::fs::create_dir_all(model_dir)?;
+
+    let model_path = model_dir.join("model.onnx");
+    if !model_path.exists() {
+        let response = client.get(MODEL_URL).send()?;
+        if !response.status().is_success() {
+            anyhow::bail!("Failed to download model: {}", response.status());
+        }
+        let bytes = response.bytes()?;
+        let mut file = std::fs::File::create(&model_path)?;
+        file.write_all(&bytes)?;
+    }
+
+    let tokenizer_path = model_dir.join("tokenizer.json");
+    if !tokenizer_path.exists() {
+        let response = client.get(TOKENIZER_URL).send()?;
+        if !response.status().is_success() {
+            anyhow::bail!("Failed to download tokenizer: {}", response.status());
+        }
+        let bytes = response.bytes()?;
+        let mut file = std::fs::File::create(&tokenizer_path)?;
+        file.write_all(&bytes)?;
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn cosine_similarity_handles_basic_cases() {
+        let a = vec![1.0, 0.0, 0.0];
+        let b = vec![1.0, 0.0, 0.0];
+        let c = vec![0.0, 1.0, 0.0];
+        let d = vec![-1.0, 0.0, 0.0];
+
+        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
+        assert!((cosine_similarity(&a, &c) - 0.0).abs() < 0.001);
+        assert!((cosine_similarity(&a, &d) - (-1.0)).abs() < 0.001);
+    }
+
+    #[test]
+    fn find_similar_returns_only_top_k_sorted_hits() {
+        let query = vec![1.0, 0.0, 0.0];
+        let candidates = vec![
+            vec![0.2, 0.0, 0.0],
+            vec![0.9, 0.0, 0.0],
+            vec![0.7, 0.0, 0.0],
+            vec![0.8, 0.0, 0.0],
+        ];
+
+        let hits = find_similar(&query, &candidates, 0.1, 2);
+
+        assert_eq!(hits, vec![(1, 0.9), (3, 0.8)]);
+    }
+
+    fn related_beats_unrelated(model_dir: &Path) {
+        let e = Embedder::load_from_dir(model_dir).expect("load model");
+        let q = e.embed("how do I set the cargo build profile").unwrap();
+        let related = e.embed("The build uses the selfdev cargo profile").unwrap();
+        let unrelated = e.embed("Bees pollinate flowers in spring").unwrap();
+        let sim_rel = cosine_similarity(&q, &related);
+        let sim_unrel = cosine_similarity(&q, &unrelated);
+        assert!(
+            sim_rel > sim_unrel + 0.05,
+            "expected related ({sim_rel:.3}) >> unrelated ({sim_unrel:.3}) for {}",
+            model_dir.display()
+        );
+    }
+
+    /// Regression test for the input-binding fix: the real MiniLM model must
+    /// produce meaningfully higher similarity for related vs unrelated text.
+    /// Skipped automatically if the model isn't present locally.
+    #[test]
+    fn minilm_related_beats_unrelated_if_present() {
+        let dir = std::env::var_os("HOME")
+            .map(|h| std::path::PathBuf::from(h).join(".jcode/models/all-MiniLM-L6-v2"))
+            .filter(|d| is_model_available(d));
+        match dir {
+            Some(d) => related_beats_unrelated(&d),
+            None => eprintln!("skip: MiniLM model not present locally"),
+        }
+    }
+
+    /// Exploratory check for an alternate model with a DIFFERENT input
+    /// order/dtype (e5-small-v2: attention_mask declared first). Reports the
+    /// related vs unrelated gap but does NOT hard-fail: some models need
+    /// model-specific pooling/normalization (e.g. CLS pooling) that the shared
+    /// mean-pool path does not yet provide. Skipped if not present locally.
+    #[test]
+    fn alt_model_related_beats_unrelated_if_present() {
+        let dir = std::env::var_os("HOME")
+            .map(|h| std::path::PathBuf::from(h).join("jcode-memory-bench/models/e5-small-v2"))
+            .filter(|d| is_model_available(d));
+        let Some(d) = dir else {
+            eprintln!("skip: e5-small-v2 model not present locally");
+            return;
+        };
+        let e = Embedder::load_from_dir(&d).expect("load model");
+        let q = e.embed("how do I set the cargo build profile").unwrap();
+        let related = e.embed("The build uses the selfdev cargo profile").unwrap();
+        let unrelated = e.embed("Bees pollinate flowers in spring").unwrap();
+        let sim_rel = cosine_similarity(&q, &related);
+        let sim_unrel = cosine_similarity(&q, &unrelated);
+        eprintln!(
+            "e5-small-v2: related={sim_rel:.4} unrelated={sim_unrel:.4} gap={:.4} (informational; mean-pool may need CLS for this family)",
+            sim_rel - sim_unrel
+        );
+    }
+
+    /// Cross-encoder reranker must score a relevant (query, passage) pair higher
+    /// than an irrelevant one. Skipped if the model isn't present locally.
+    #[test]
+    fn cross_encoder_scores_relevant_higher_if_present() {
+        let dir = std::env::var_os("HOME")
+            .map(|h| std::path::PathBuf::from(h).join("jcode-memory-bench/models/ce-minilm-l6"))
+            .filter(|d| d.join("model.onnx").exists() && d.join("tokenizer.json").exists());
+        let Some(d) = dir else {
+            eprintln!("skip: cross-encoder model not present locally");
+            return;
+        };
+        let ce = CrossEncoder::load_from_dir(&d).expect("load cross-encoder");
+        let q = "how do I set the cargo build profile";
+        let rel = ce
+            .score(q, "The build uses the selfdev cargo profile")
+            .unwrap();
+        let unrel = ce.score(q, "Bees pollinate flowers in spring").unwrap();
+        eprintln!("cross-encoder: relevant={rel:.3} irrelevant={unrel:.3}");
+        assert!(
+            rel > unrel,
+            "cross-encoder must score relevant ({rel:.3}) > irrelevant ({unrel:.3})"
+        );
+    }
+}

crates/jcode-embedding/tests/embed_latency_probe.rs

@@ -0,0 +1,67 @@
+//! Ignored embed-latency probe for the tract inference stack.
+//!
+//! Motivation: at opt-level 0 (plain dev/selfdev profiles before the
+//! workspace pinned tract-* and tokenizers to opt-level 3) a single MiniLM
+//! embed measured ~666 ms of interpreter overhead on the shared server. That
+//! latency kept the embedding model "busy" through recurring memory
+//! maintenance, so the 15-minute idle unloader never fired and the model's
+//! ~100 MB stayed resident indefinitely.
+//!
+//! Run with:
+//!   cargo test -p jcode-embedding --test embed_latency_probe -- --ignored --nocapture
+//!
+//! Requires the MiniLM model to be installed (~/.jcode/models/all-MiniLM-L6-v2).
+
+use std::path::PathBuf;
+
+fn model_dir() -> Option<PathBuf> {
+    let home = std::env::var_os("HOME").map(PathBuf::from)?;
+    let dir = home
+        .join(".jcode")
+        .join("models")
+        .join(jcode_embedding::MODEL_NAME);
+    dir.join("model.onnx").exists().then_some(dir)
+}
+
+#[test]
+#[ignore = "perf probe; requires installed model; run with --ignored --nocapture"]
+fn embed_latency_probe() {
+    let Some(dir) = model_dir() else {
+        eprintln!("model not installed; skipping");
+        return;
+    };
+    let load_start = std::time::Instant::now();
+    let embedder = jcode_embedding::Embedder::load_from_dir(&dir).expect("load model");
+    let load = load_start.elapsed();
+
+    // Warm once (first run pays one-time plan setup).
+    let _ = embedder
+        .embed("warmup sentence for the embedding probe")
+        .expect("warm embed");
+
+    const ITERS: usize = 10;
+    let texts: Vec<String> = (0..ITERS)
+        .map(|i| format!("memory recall probe sentence number {i} with a few extra tokens"))
+        .collect();
+    let start = std::time::Instant::now();
+    for text in &texts {
+        let v = embedder.embed(text).expect("embed");
+        assert_eq!(v.len(), 384);
+    }
+    let per_embed_ms = start.elapsed().as_secs_f64() * 1000.0 / ITERS as f64;
+
+    println!("embed latency probe:");
+    println!("  model load: {load:?}");
+    println!("  per-embed:  {per_embed_ms:.1} ms (over {ITERS} iters)");
+
+    // Regression guard: with the tract stack pinned to opt-level 3 this
+    // measures ~300 ms on this hardware (the model always runs a full
+    // 256-token forward pass); the opt-level 0 regression measured ~666 ms.
+    // The bound sits between the two so profile regressions fail loudly
+    // without flaking on normal variance.
+    assert!(
+        per_embed_ms < 450.0,
+        "embed took {per_embed_ms:.1} ms; tract opt-level regression? \
+         (check [profile.*.package.tract-*] pins in the workspace Cargo.toml)"
+    );
+}