searchi/crates/searchi-semantic/src/minilm.rs

//! 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());
    }
}