feat: integrate Google semantic search and branded desktop launcher

crates/searchi-semantic/src/gemma.rs

@@ -0,0 +1,398 @@
+use crate::TextWindow;
+use anyhow::{ensure, Context, Result};
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+use sha2::{Digest, Sha256};
+use std::{
+    fs::File,
+    io::{BufRead, BufReader, Read, Write},
+    path::{Path, PathBuf},
+    process::{Child, ChildStdin, Command, Stdio},
+    sync::mpsc::{self, Receiver},
+    thread,
+    time::Duration,
+};
+
+pub const GEMMA_DIMENSIONS: usize = 768;
+pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
+const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf";
+const MODEL_BYTES: u64 = 309855456;
+const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135";
+const MAX_REQUEST: usize = 2 * 1024 * 1024;
+const MAX_RESPONSE: usize = 4 * 1024 * 1024;
+
+#[derive(Debug, Clone, Copy)]
+pub enum Device {
+    Auto,
+    Cpu,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct WorkerInfo {
+    pub ready: bool,
+    pub model: String,
+    pub dimensions: usize,
+    pub device: String,
+    pub threads: usize,
+    pub runtime: String,
+}
+
+/// Persistent isolated native worker. No HTTP or remote model fallback.
+pub struct GemmaEmbedder {
+    child: Child,
+    stdin: ChildStdin,
+    responses: Receiver<std::result::Result<Value, String>>,
+    pub info: WorkerInfo,
+    healthy: bool,
+    force_cpu: bool,
+    model_dir: PathBuf,
+    worker_path: PathBuf,
+}
+
+impl Drop for GemmaEmbedder {
+    fn drop(&mut self) {
+        let _ = self.child.kill();
+        let _ = self.child.wait();
+    }
+}
+
+impl GemmaEmbedder {
+    /// Load the verified Google model and explicitly supplied local worker.
+    /// Production selects available acceleration with CPU fallback. Tests set
+    /// `SEARCHI_FORCE_CPU=1` or call `load_with_device(..., Device::Cpu)`.
+    pub fn load(
+        model_dir: impl AsRef<Path>,
+        worker: impl AsRef<Path>,
+        threads: usize,
+    ) -> Result<Self> {
+        Self::load_with_device(model_dir, worker, threads, Device::Auto)
+    }
+
+    pub fn load_with_device(
+        model_dir: impl AsRef<Path>,
+        worker: impl AsRef<Path>,
+        threads: usize,
+        device: Device,
+    ) -> Result<Self> {
+        ensure!(
+            (1..=8).contains(&threads),
+            "CPU threads must be between 1 and 8"
+        );
+        let model_dir = model_dir.as_ref().to_path_buf();
+        let model = model_dir.join(MODEL_FILE);
+        verify_model(&model)?;
+        let worker = worker
+            .as_ref()
+            .canonicalize()
+            .context("Cannot resolve local Gemma worker executable")?;
+        ensure!(
+            worker.is_file(),
+            "Gemma worker must be a local executable file"
+        );
+        let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1");
+        let cpu = force_cpu || matches!(device, Device::Cpu);
+        let mut command = Command::new(&worker);
+        command
+            .args(["--model"])
+            .arg(model)
+            .args([
+                "--threads",
+                &threads.to_string(),
+                "--device",
+                if cpu { "cpu" } else { "auto" },
+            ])
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::inherit());
+        // Background helpers must not create a visible console on Windows.
+        #[cfg(windows)]
+        {
+            use std::os::windows::process::CommandExt;
+            command.creation_flags(0x08000000); // CREATE_NO_WINDOW
+        }
+        let mut child = command.spawn().context("Cannot start local Gemma worker")?;
+        let stdin = child.stdin.take().context("Worker stdin missing")?;
+        let stdout = child.stdout.take().context("Worker stdout missing")?;
+        let (sender, responses) = mpsc::sync_channel(1);
+        thread::spawn(move || {
+            let mut reader = BufReader::new(stdout);
+            loop {
+                let response = read_response(&mut reader);
+                let failed = response.is_err();
+                if sender.send(response).is_err() || failed {
+                    break;
+                }
+            }
+        });
+        let ready = match responses.recv_timeout(Duration::from_secs(120)) {
+            Ok(Ok(value)) => value,
+            other => {
+                let _ = child.kill();
+                let _ = child.wait();
+                anyhow::bail!("Gemma worker did not initialize: {other:?}");
+            }
+        };
+        let info: WorkerInfo = match serde_json::from_value(ready) {
+            Ok(info) => info,
+            Err(error) => {
+                let _ = child.kill();
+                let _ = child.wait();
+                return Err(error).context("Invalid Gemma worker handshake");
+            }
+        };
+        if !info.ready
+            || info.model != "google/embeddinggemma-2"
+            || info.dimensions != GEMMA_DIMENSIONS
+            || info.runtime != "b11461-no-logits-v1"
+            || info.threads != threads
+            || !["cpu", "gpu"].contains(&info.device.as_str())
+            || (cpu && info.device != "cpu")
+        {
+            let _ = child.kill();
+            let _ = child.wait();
+            anyhow::bail!("Gemma worker handshake violates configured model/device contract");
+        }
+        Ok(Self {
+            child,
+            stdin,
+            responses,
+            info,
+            healthy: true,
+            force_cpu: cpu,
+            model_dir,
+            worker_path: worker,
+        })
+    }
+
+    fn request(&mut self, request: Value) -> Result<Value> {
+        ensure!(
+            self.healthy,
+            "Gemma worker session failed; reopen it explicitly"
+        );
+        let encoded = serde_json::to_vec(&request)?;
+        ensure!(
+            encoded.len() < MAX_REQUEST,
+            "Gemma request exceeds 2 MiB protocol budget"
+        );
+        let response = match self.exchange(&encoded) {
+            Ok(response) => response,
+            Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
+                // GPU runtimes can abort rather than return an encoder error.
+                // Restart once as CPU and retry this same bounded request.
+                let _ = self.child.wait();
+                let replacement = Self::load_with_device(
+                    &self.model_dir,
+                    &self.worker_path,
+                    self.info.threads,
+                    Device::Cpu,
+                )
+                .with_context(|| {
+                    format!("CPU worker restart failed after GPU transport error: {error}")
+                })?;
+                *self = replacement;
+                self.exchange(&encoded)?
+            }
+            Err(error) => return Err(error),
+        };
+        if let Some(device) = response["device"].as_str() {
+            ensure!(
+                ["cpu", "gpu"].contains(&device) && (!self.force_cpu || device == "cpu"),
+                "Worker response violates configured device policy"
+            );
+            self.info.device = device.into();
+        }
+        if response["fatal"] == true {
+            self.healthy = false;
+            self.info.ready = false;
+            let _ = self.child.kill();
+        }
+        ensure!(
+            response["ok"] == true,
+            "Gemma worker error: {}",
+            response["error"].as_str().unwrap_or("unknown error")
+        );
+        Ok(response)
+    }
+
+    fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
+        let sent = (|| {
+            self.stdin.write_all(encoded)?;
+            self.stdin.write_all(b"\n")?;
+            self.stdin.flush()
+        })()
+        .context("Cannot write local Gemma worker request");
+        if let Err(error) = sent {
+            self.healthy = false;
+            self.info.ready = false;
+            let _ = self.child.kill();
+            return Err(error);
+        }
+        let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
+            Ok(Ok(value)) => value,
+            other => {
+                self.healthy = false;
+                self.info.ready = false;
+                let _ = self.child.kill();
+                anyhow::bail!("Gemma worker request failed: {other:?}");
+            }
+        };
+        Ok(response)
+    }
+
+    pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
+        ensure!(
+            !text.trim().is_empty() && text.len() <= 1024 * 1024,
+            "Passage must be nonempty and at most 1 MiB"
+        );
+        let response = self.request(json!({"op":"windows","text":text,"title":title}))?;
+        let windows: Vec<TextWindow> = serde_json::from_value(response["windows"].clone())?;
+        ensure!(
+            !windows.is_empty() && windows.len() <= 256,
+            "Invalid Gemma window count"
+        );
+        for window in &windows {
+            ensure!(
+                text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
+                "Gemma window does not match original source range"
+            );
+        }
+        ensure!(
+            windows
+                .last()
+                .is_some_and(|window| window.end_byte == text.len()),
+            "Gemma windows did not cover the passage tail"
+        );
+        Ok(windows)
+    }
+
+    pub fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
+        ensure!(
+            !texts.is_empty() && texts.len() <= 8,
+            "Document batch size must be between 1 and 8"
+        );
+        ensure!(
+            texts
+                .iter()
+                .all(|text| !text.trim().is_empty() && text.len() <= 65536),
+            "Document window must be nonempty and at most 64 KiB"
+        );
+        let response = self.request(json!({"op":"documents","texts":texts,"title":title}))?;
+        extract_vectors(response, texts.len())
+    }
+
+    pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        ensure!(
+            !query.trim().is_empty() && query.len() <= 4096,
+            "Query must be nonempty and at most 4096 UTF-8 bytes"
+        );
+        let response = self.request(json!({"op":"query","text":query}))?;
+        Ok(extract_vectors(response, 1)?.remove(0))
+    }
+}
+
+fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
+    actual_device == "gpu" && !force_cpu
+}
+
+fn read_response(reader: &mut impl BufRead) -> std::result::Result<Value, String> {
+    let mut line = Vec::new();
+    match reader
+        .take((MAX_RESPONSE + 1) as u64)
+        .read_until(b'\n', &mut line)
+    {
+        Ok(0) => Err("Gemma worker closed its output".into()),
+        Ok(_) if line.len() > MAX_RESPONSE || line.last() != Some(&b'\n') => {
+            Err("Gemma worker response exceeds protocol byte budget".into())
+        }
+        Ok(_) => serde_json::from_slice(&line)
+            .map_err(|error| format!("Invalid Gemma worker JSON: {error}")),
+        Err(error) => Err(format!("Gemma worker read failed: {error}")),
+    }
+}
+
+fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
+    let vectors: Vec<Vec<f32>> = serde_json::from_value(response["vectors"].clone())?;
+    ensure!(
+        vectors.len() == count,
+        "Gemma returned incorrect embedding batch length"
+    );
+    for vector in &vectors {
+        ensure!(
+            vector.len() == GEMMA_DIMENSIONS && vector.iter().all(|value| value.is_finite()),
+            "Invalid Gemma embedding shape or values"
+        );
+        let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
+        ensure!(
+            (norm - 1.0).abs() < 1e-4,
+            "Gemma embedding is not L2 normalized"
+        );
+    }
+    Ok(vectors)
+}
+
+fn verify_model(path: &Path) -> Result<()> {
+    let file = File::open(path).with_context(|| {
+        format!(
+            "Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
+            path.display()
+        )
+    })?;
+    ensure!(
+        file.metadata()?.len() == MODEL_BYTES,
+        "EmbeddingGemma 2 model size mismatch"
+    );
+    let mut reader = BufReader::new(file);
+    let mut hash = Sha256::new();
+    let mut buffer = [0; 65536];
+    loop {
+        let count = reader.read(&mut buffer)?;
+        if count == 0 {
+            break;
+        }
+        hash.update(&buffer[..count]);
+    }
+    ensure!(
+        format!("{:x}", hash.finalize()) == MODEL_SHA256,
+        "EmbeddingGemma 2 model SHA256 mismatch"
+    );
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn missing_google_model_fails_without_starting_worker() {
+        assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
+    }
+
+    #[test]
+    fn retry_policy_never_enables_gpu_or_retries_cpu_as_gpu() {
+        assert!(can_retry_cpu("gpu", false));
+        assert!(!can_retry_cpu("gpu", true));
+        assert!(!can_retry_cpu("cpu", false));
+        assert!(!can_retry_cpu("cpu", true));
+    }
+
+    #[test]
+    fn protocol_rejects_oversized_truncated_and_malformed_frames() {
+        assert!(read_response(&mut std::io::Cursor::new(vec![b'x'; MAX_RESPONSE + 1])).is_err());
+        assert!(read_response(&mut std::io::Cursor::new(b"{\"ok\":true}")).is_err());
+        assert!(read_response(&mut std::io::Cursor::new(b"not-json\n")).is_err());
+        assert_eq!(
+            read_response(&mut std::io::Cursor::new(b"{\"ok\":true}\n")).unwrap()["ok"],
+            true
+        );
+    }
+
+    #[test]
+    fn protocol_rejects_wrong_embedding_shape_count_and_normalization() {
+        assert!(extract_vectors(json!({"vectors":[]}), 1).is_err());
+        assert!(extract_vectors(json!({"vectors":[[1.0,0.0]]}), 1).is_err());
+        assert!(extract_vectors(json!({"vectors":[vec![0.0;GEMMA_DIMENSIONS]]}), 1).is_err());
+        let mut invalid = vec![0.0; GEMMA_DIMENSIONS];
+        invalid[0] = 1.0;
+        invalid[1] = 1.0;
+        assert!(extract_vectors(json!({"vectors":[invalid]}), 1).is_err());
+    }
+}

crates/searchi-semantic/src/lib.rs

@@ -1,290 +1,19 @@
-//! Offline text embeddings using a pinned quantized multilingual model.
-//!
-//! Model provisioning is explicit and separate. This crate never downloads a
-//! model or runtime and exposes no GPU execution-provider configuration.
+//! Local Google EmbeddingGemma 2 inference through an isolated native worker.
+//! The rejected MiniLM experiment is available only under `minilm-proof`.
+mod gemma;
+pub use gemma::{Device, GemmaEmbedder, WorkerInfo, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
 
-use std::{
-    fs::File,
-    io::{BufReader, Read},
-    path::{Path, PathBuf},
-    sync::Mutex,
+/// A verified source byte range covered by a model token window.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct TextWindow {
+    pub text: String,
+    pub start_byte: usize,
+    pub end_byte: usize,
+}
+
+#[cfg(feature = "minilm-proof")]
+mod minilm;
+#[cfg(feature = "minilm-proof")]
+pub use minilm::{
+    CpuEmbedder, DIMENSIONS, MAX_BATCH, MAX_TEXT_BYTES, MAX_TOKENS, MODEL_FINGERPRINT,
 };
-
-use anyhow::{bail, ensure, Context, Result};
-use ort::{ep, session::Session, value::Tensor};
-use sha2::{Digest, Sha256};
-use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
-
-/// Embedding dimensionality; stored vectors must carry the model fingerprint.
-pub const DIMENSIONS: usize = 384;
-pub const MAX_TOKENS: usize = 128;
-pub const MAX_BATCH: usize = 8;
-pub const MAX_TEXT_BYTES: usize = 65536;
-/// Includes model revision, quantization, pooling, normalization and token cap.
-pub const MODEL_FINGERPRINT: &str =
-    "minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
-
-const ARTIFACTS: [(&str, u64, &str); 2] = [
-    (
-        "onnx/model_quantized.onnx",
-        118308126,
-        "66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
-    ),
-    (
-        "tokenizer.json",
-        17082913,
-        "b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
-    ),
-];
-// ONNX Runtime's library and environment are process-global. Prevent concurrent
-// initializers from selecting different libraries or silently overriding state.
-static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
-
-/// Local, mutable inference session. Call from a bounded background worker.
-pub struct CpuEmbedder {
-    session: Session,
-    tokenizer: Tokenizer,
-}
-
-impl CpuEmbedder {
-    /// Verify local artifacts, then load the explicitly supplied runtime.
-    ///
-    /// The runtime must be initialized through this adapter before other `ort`
-    /// code in this process. Threads are limited to 1–8; default caller policy
-    /// should use one thread in a VM. Missing or damaged assets are errors.
-    pub fn load(
-        model_dir: impl AsRef<Path>,
-        runtime_library: impl AsRef<Path>,
-        threads: usize,
-    ) -> Result<Self> {
-        ensure!(
-            (1..=8).contains(&threads),
-            "CPU threads must be between 1 and 8"
-        );
-        let model_dir = model_dir.as_ref();
-        verify_artifacts(model_dir)?;
-        let runtime_path = runtime_library
-            .as_ref()
-            .canonicalize()
-            .context("Cannot resolve explicitly supplied ONNX Runtime library")?;
-        ensure!(
-            runtime_path.is_file(),
-            "ONNX Runtime library must be a file"
-        );
-        {
-            let mut initialized = RUNTIME_PATH
-                .lock()
-                .map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
-            if let Some(existing) = initialized.as_ref() {
-                ensure!(
-                    *existing == runtime_path,
-                    "ONNX Runtime already initialized with a different library"
-                );
-            } else {
-                let committed = ort::init_from(&runtime_path)?
-                    .with_telemetry(false)
-                    .with_execution_providers([ep::CPU::default().build()])
-                    .commit();
-                ensure!(
-                    committed,
-                    "ONNX Runtime environment was already initialized outside the CPU adapter"
-                );
-                *initialized = Some(runtime_path);
-            }
-        }
-        let session = Session::builder()?
-            .with_execution_providers([ep::CPU::default().build()])
-            .map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
-            .with_intra_threads(threads)
-            .map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
-            .with_inter_threads(1)
-            .map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
-            .with_parallel_execution(false)
-            .map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
-            .commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
-        let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
-            .map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
-        tokenizer
-            .with_truncation(Some(TruncationParams {
-                max_length: MAX_TOKENS,
-                ..Default::default()
-            }))
-            .map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
-        tokenizer.with_padding(Some(PaddingParams {
-            pad_id: 1,
-            pad_token: "<pad>".into(),
-            ..Default::default()
-        }));
-        Ok(Self { session, tokenizer })
-    }
-
-    /// Produce finite, L2-normalized vectors from genuine local model inference.
-    ///
-    /// Inputs above 128 tokens are truncated. Production indexing should chunk
-    /// documents first and keep chunk locations; never silently embed entire PDFs.
-    pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
-        validate_texts(texts)?;
-        // Encode sequentially so tokenizers cannot create an unbounded global
-        // Rayon worker pool independently of the caller's inference budget.
-        let mut encoded = texts
-            .iter()
-            .map(|text| self.tokenizer.encode(*text, true))
-            .collect::<std::result::Result<Vec<_>, _>>()
-            .map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
-        let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
-        for item in &mut encoded {
-            item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
-        }
-        ensure!(
-            sequence_len > 0 && sequence_len <= MAX_TOKENS,
-            "Invalid token sequence length"
-        );
-        let ids: Vec<i64> = encoded
-            .iter()
-            .flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
-            .collect();
-        let masks: Vec<i64> = encoded
-            .iter()
-            .flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
-            .collect();
-        let types: Vec<i64> = encoded
-            .iter()
-            .flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
-            .collect();
-        let shape = [texts.len(), sequence_len];
-        let outputs = self.session.run(ort::inputs![
-            "input_ids" => Tensor::from_array((shape, ids))?,
-            "attention_mask" => Tensor::from_array((shape, masks.clone()))?,
-            "token_type_ids" => Tensor::from_array((shape, types))?
-        ])?;
-        let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
-        ensure!(
-            **output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
-            "Unexpected embedding output shape"
-        );
-        pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
-    }
-}
-
-fn validate_texts(texts: &[&str]) -> Result<()> {
-    ensure!(
-        !texts.is_empty() && texts.len() <= MAX_BATCH,
-        "Batch size must be between 1 and {MAX_BATCH}"
-    );
-    ensure!(
-        texts
-            .iter()
-            .all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
-        "Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
-    );
-    Ok(())
-}
-
-fn verify_artifacts(model_dir: &Path) -> Result<()> {
-    for (relative, bytes, expected) in ARTIFACTS {
-        let path = model_dir.join(relative);
-        let file = File::open(&path).with_context(|| {
-            format!(
-                "Missing local model artifact {}; run the explicit model downloader first",
-                path.display()
-            )
-        })?;
-        ensure!(
-            file.metadata()?.len() == bytes,
-            "Model artifact size mismatch: {}",
-            path.display()
-        );
-        let mut reader = BufReader::new(file);
-        let mut hash = Sha256::new();
-        let mut buffer = [0_u8; 65536];
-        loop {
-            let read = reader.read(&mut buffer)?;
-            if read == 0 {
-                break;
-            }
-            hash.update(&buffer[..read]);
-        }
-        let actual = format!("{:x}", hash.finalize());
-        ensure!(
-            actual == expected,
-            "Model artifact SHA256 mismatch: {}",
-            path.display()
-        );
-    }
-    Ok(())
-}
-
-fn pool_and_normalize(
-    hidden: &[f32],
-    masks: &[i64],
-    batch: usize,
-    tokens: usize,
-) -> Result<Vec<Vec<f32>>> {
-    ensure!(
-        hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
-        "Invalid pooling input dimensions"
-    );
-    let mut vectors = Vec::with_capacity(batch);
-    for item in 0..batch {
-        let mut vector = vec![0_f32; DIMENSIONS];
-        let mut count = 0;
-        for token in 0..tokens {
-            if masks[item * tokens + token] == 0 {
-                continue;
-            }
-            count += 1;
-            let offset = (item * tokens + token) * DIMENSIONS;
-            for dimension in 0..DIMENSIONS {
-                vector[dimension] += hidden[offset + dimension];
-            }
-        }
-        ensure!(count > 0, "Cannot pool a fully masked sequence");
-        for value in &mut vector {
-            *value /= count as f32;
-        }
-        let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
-        if !norm.is_finite() || norm <= f32::EPSILON {
-            bail!("Model produced nonfinite or zero embeddings");
-        }
-        for value in &mut vector {
-            *value /= norm;
-        }
-        ensure!(
-            vector.iter().all(|value| value.is_finite()),
-            "Model produced nonfinite embeddings"
-        );
-        vectors.push(vector);
-    }
-    Ok(vectors)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn padding_is_excluded_and_output_is_normalized() {
-        let mut hidden = vec![0.0; 2 * DIMENSIONS];
-        hidden[0] = 3.0;
-        hidden[1] = 4.0;
-        hidden[DIMENSIONS] = 999.0;
-        let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
-        assert!((pooled[0][0] - 0.6).abs() < 1e-6);
-        assert!((pooled[0][1] - 0.8).abs() < 1e-6);
-    }
-
-    #[test]
-    fn rejects_unbounded_inputs_and_invalid_outputs() {
-        assert!(validate_texts(&[]).is_err());
-        assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
-        assert!(validate_texts(&[" "]).is_err());
-        assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
-        assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
-        assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
-    }
-
-    #[test]
-    fn missing_model_does_not_load_or_download_runtime() {
-        assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
-    }
-}

crates/searchi-semantic/src/minilm.rs

@@ -0,0 +1,347 @@
+//! Offline text embeddings using a pinned quantized multilingual model.
+//!
+//! Model provisioning is explicit and separate. This crate never downloads a
+//! model or runtime and exposes no GPU execution-provider configuration.
+
+use std::{
+    fs::File,
+    io::{BufReader, Read},
+    path::{Path, PathBuf},
+    sync::Mutex,
+};
+
+use anyhow::{bail, ensure, Context, Result};
+use ort::{ep, session::Session, value::Tensor};
+use sha2::{Digest, Sha256};
+use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
+
+/// Embedding dimensionality; stored vectors must carry the model fingerprint.
+pub const DIMENSIONS: usize = 384;
+pub const MAX_TOKENS: usize = 128;
+pub const MAX_BATCH: usize = 8;
+pub const MAX_TEXT_BYTES: usize = 65536;
+pub const MAX_PASSAGE_BYTES: usize = 1024 * 1024;
+pub const MAX_WINDOWS: usize = 256;
+pub const WINDOW_OVERLAP_TOKENS: usize = 24;
+/// A source byte range covered by one overlapping model token window.
+use crate::TextWindow;
+/// Includes model revision, quantization, pooling, normalization and token cap.
+pub const MODEL_FINGERPRINT: &str =
+    "minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
+
+const ARTIFACTS: [(&str, u64, &str); 2] = [
+    (
+        "onnx/model_quantized.onnx",
+        118308126,
+        "66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
+    ),
+    (
+        "tokenizer.json",
+        17082913,
+        "b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
+    ),
+];
+// ONNX Runtime's library and environment are process-global. Prevent concurrent
+// initializers from selecting different libraries or silently overriding state.
+static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
+
+/// Local, mutable inference session. Call from a bounded background worker.
+pub struct CpuEmbedder {
+    session: Session,
+    tokenizer: Tokenizer,
+}
+
+impl CpuEmbedder {
+    /// Split a bounded passage into overlapping token windows, retaining its tail.
+    /// Special tokens have no source range; source offsets remain UTF-8 byte
+    /// offsets into the original passage for accurate preview highlighting.
+    pub fn split_windows(&self, text: &str) -> Result<Vec<TextWindow>> {
+        ensure!(
+            !text.trim().is_empty() && text.len() <= MAX_PASSAGE_BYTES,
+            "Passage must be nonempty and at most {MAX_PASSAGE_BYTES} UTF-8 bytes"
+        );
+        let mut tokenizer = self.tokenizer.clone();
+        tokenizer
+            .with_truncation(Some(TruncationParams {
+                max_length: MAX_TOKENS,
+                stride: WINDOW_OVERLAP_TOKENS,
+                ..Default::default()
+            }))
+            .map_err(|error| anyhow::anyhow!("Window tokenizer configuration: {error}"))?;
+        tokenizer.with_padding(None);
+        let encoded = tokenizer
+            .encode(text, true)
+            .map_err(|error| anyhow::anyhow!("Window tokenization failed: {error}"))?;
+        ensure!(
+            encoded.get_overflowing().len() < MAX_WINDOWS,
+            "Passage exceeds {MAX_WINDOWS} model windows; split it before embedding"
+        );
+        let mut windows = Vec::new();
+        for item in std::iter::once(&encoded).chain(encoded.get_overflowing()) {
+            let offsets: Vec<_> = item
+                .get_offsets()
+                .iter()
+                .zip(item.get_special_tokens_mask())
+                .filter(|((start, end), special)| **special == 0 && end > start)
+                .map(|(offset, _)| *offset)
+                .collect();
+            let start_byte = offsets.iter().map(|(start, _)| *start).min().unwrap_or(0);
+            let end_byte = offsets.iter().map(|(_, end)| *end).max().unwrap_or(0);
+            ensure!(start_byte < end_byte, "Token window has no source text");
+            let chunk = text
+                .get(start_byte..end_byte)
+                .context("Invalid UTF-8 tokenizer offsets")?;
+            ensure!(
+                chunk.len() <= MAX_TEXT_BYTES,
+                "Token window exceeds embedding text byte budget"
+            );
+            windows.push(TextWindow {
+                text: chunk.to_owned(),
+                start_byte,
+                end_byte,
+            });
+        }
+        Ok(windows)
+    }
+
+    /// Verify local artifacts, then load the explicitly supplied runtime.
+    ///
+    /// The runtime must be initialized through this adapter before other `ort`
+    /// code in this process. Threads are limited to 1–8; default caller policy
+    /// should use one thread in a VM. Missing or damaged assets are errors.
+    pub fn load(
+        model_dir: impl AsRef<Path>,
+        runtime_library: impl AsRef<Path>,
+        threads: usize,
+    ) -> Result<Self> {
+        ensure!(
+            (1..=8).contains(&threads),
+            "CPU threads must be between 1 and 8"
+        );
+        let model_dir = model_dir.as_ref();
+        verify_artifacts(model_dir)?;
+        let runtime_path = runtime_library
+            .as_ref()
+            .canonicalize()
+            .context("Cannot resolve explicitly supplied ONNX Runtime library")?;
+        ensure!(
+            runtime_path.is_file(),
+            "ONNX Runtime library must be a file"
+        );
+        {
+            let mut initialized = RUNTIME_PATH
+                .lock()
+                .map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
+            if let Some(existing) = initialized.as_ref() {
+                ensure!(
+                    *existing == runtime_path,
+                    "ONNX Runtime already initialized with a different library"
+                );
+            } else {
+                let committed = ort::init_from(&runtime_path)?
+                    .with_telemetry(false)
+                    .with_execution_providers([ep::CPU::default().build()])
+                    .commit();
+                ensure!(
+                    committed,
+                    "ONNX Runtime environment was already initialized outside the CPU adapter"
+                );
+                *initialized = Some(runtime_path);
+            }
+        }
+        let session = Session::builder()?
+            .with_execution_providers([ep::CPU::default().build()])
+            .map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
+            .with_intra_threads(threads)
+            .map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
+            .with_inter_threads(1)
+            .map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
+            .with_parallel_execution(false)
+            .map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
+            .commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
+        let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
+            .map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
+        tokenizer
+            .with_truncation(Some(TruncationParams {
+                max_length: MAX_TOKENS,
+                ..Default::default()
+            }))
+            .map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
+        tokenizer.with_padding(Some(PaddingParams {
+            pad_id: 1,
+            pad_token: "<pad>".into(),
+            ..Default::default()
+        }));
+        Ok(Self { session, tokenizer })
+    }
+
+    /// Produce finite, L2-normalized vectors from genuine local model inference.
+    ///
+    /// Inputs above 128 tokens are truncated. Production indexing should chunk
+    /// documents first and keep chunk locations; never silently embed entire PDFs.
+    pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
+        validate_texts(texts)?;
+        // Encode sequentially so tokenizers cannot create an unbounded global
+        // Rayon worker pool independently of the caller's inference budget.
+        let mut encoded = texts
+            .iter()
+            .map(|text| self.tokenizer.encode(*text, true))
+            .collect::<std::result::Result<Vec<_>, _>>()
+            .map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
+        let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
+        for item in &mut encoded {
+            item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
+        }
+        ensure!(
+            sequence_len > 0 && sequence_len <= MAX_TOKENS,
+            "Invalid token sequence length"
+        );
+        let ids: Vec<i64> = encoded
+            .iter()
+            .flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
+            .collect();
+        let masks: Vec<i64> = encoded
+            .iter()
+            .flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
+            .collect();
+        let types: Vec<i64> = encoded
+            .iter()
+            .flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
+            .collect();
+        let shape = [texts.len(), sequence_len];
+        let outputs = self.session.run(ort::inputs![
+            "input_ids" => Tensor::from_array((shape, ids))?,
+            "attention_mask" => Tensor::from_array((shape, masks.clone()))?,
+            "token_type_ids" => Tensor::from_array((shape, types))?
+        ])?;
+        let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
+        ensure!(
+            **output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
+            "Unexpected embedding output shape"
+        );
+        pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
+    }
+}
+
+fn validate_texts(texts: &[&str]) -> Result<()> {
+    ensure!(
+        !texts.is_empty() && texts.len() <= MAX_BATCH,
+        "Batch size must be between 1 and {MAX_BATCH}"
+    );
+    ensure!(
+        texts
+            .iter()
+            .all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
+        "Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
+    );
+    Ok(())
+}
+
+fn verify_artifacts(model_dir: &Path) -> Result<()> {
+    for (relative, bytes, expected) in ARTIFACTS {
+        let path = model_dir.join(relative);
+        let file = File::open(&path).with_context(|| {
+            format!(
+                "Missing local model artifact {}; run the explicit model downloader first",
+                path.display()
+            )
+        })?;
+        ensure!(
+            file.metadata()?.len() == bytes,
+            "Model artifact size mismatch: {}",
+            path.display()
+        );
+        let mut reader = BufReader::new(file);
+        let mut hash = Sha256::new();
+        let mut buffer = [0_u8; 65536];
+        loop {
+            let read = reader.read(&mut buffer)?;
+            if read == 0 {
+                break;
+            }
+            hash.update(&buffer[..read]);
+        }
+        let actual = format!("{:x}", hash.finalize());
+        ensure!(
+            actual == expected,
+            "Model artifact SHA256 mismatch: {}",
+            path.display()
+        );
+    }
+    Ok(())
+}
+
+fn pool_and_normalize(
+    hidden: &[f32],
+    masks: &[i64],
+    batch: usize,
+    tokens: usize,
+) -> Result<Vec<Vec<f32>>> {
+    ensure!(
+        hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
+        "Invalid pooling input dimensions"
+    );
+    let mut vectors = Vec::with_capacity(batch);
+    for item in 0..batch {
+        let mut vector = vec![0_f32; DIMENSIONS];
+        let mut count = 0;
+        for token in 0..tokens {
+            if masks[item * tokens + token] == 0 {
+                continue;
+            }
+            count += 1;
+            let offset = (item * tokens + token) * DIMENSIONS;
+            for dimension in 0..DIMENSIONS {
+                vector[dimension] += hidden[offset + dimension];
+            }
+        }
+        ensure!(count > 0, "Cannot pool a fully masked sequence");
+        for value in &mut vector {
+            *value /= count as f32;
+        }
+        let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
+        if !norm.is_finite() || norm <= f32::EPSILON {
+            bail!("Model produced nonfinite or zero embeddings");
+        }
+        for value in &mut vector {
+            *value /= norm;
+        }
+        ensure!(
+            vector.iter().all(|value| value.is_finite()),
+            "Model produced nonfinite embeddings"
+        );
+        vectors.push(vector);
+    }
+    Ok(vectors)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn padding_is_excluded_and_output_is_normalized() {
+        let mut hidden = vec![0.0; 2 * DIMENSIONS];
+        hidden[0] = 3.0;
+        hidden[1] = 4.0;
+        hidden[DIMENSIONS] = 999.0;
+        let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
+        assert!((pooled[0][0] - 0.6).abs() < 1e-6);
+        assert!((pooled[0][1] - 0.8).abs() < 1e-6);
+    }
+
+    #[test]
+    fn rejects_unbounded_inputs_and_invalid_outputs() {
+        assert!(validate_texts(&[]).is_err());
+        assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
+        assert!(validate_texts(&[" "]).is_err());
+        assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
+        assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
+        assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
+    }
+
+    #[test]
+    fn missing_model_does_not_load_or_download_runtime() {
+        assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
+    }
+}