feat: add desktop document search and verified native CPU embeddings

crates/searchi-semantic/src/lib.rs

@@ -0,0 +1,290 @@
+//! 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;
+/// 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());
+    }
+}