rightkit_embed/
embedder.rs

1use std::path::PathBuf;
2use std::sync::Mutex;
3
4use ort::session::Session;
5use rightkit_ort::{
6    configure_runtime, init_environment, CandidateSource, EnvironmentOptions, ExecutionProvider,
7    GlobalPool, RuntimeCandidate, SessionOptions,
8};
9use thiserror::Error;
10use tokenizers::{EncodeInput, Tokenizer};
11
12pub const EMBEDDINGGEMMA_MODEL_ID: &str = "embeddinggemma-300m-q4";
13/// EmbeddingGemma 2 text-only Q4 export (b64, accuracy level 4). Selected when
14/// the model directory's top-level `config.json` `model_type` is
15/// `embedding_gemma2`.
16pub const EMBEDDINGGEMMA2_MODEL_ID: &str = "embeddinggemma2-q4";
17/// Literal special tokens that encode to EmbeddingGemma 2's image, video and
18/// audio sentinel ids; the text-only graph is exact only without them.
19const GEMMA2_MODALITY_SENTINELS: [&str; 3] = ["<|image|>", "<|video|>", "<|audio|>"];
20pub const EMBEDDINGGEMMA_NATIVE_DIM: usize = 768;
21pub const DEFAULT_TRUNCATE_DIM: usize = 256;
22pub const MAX_SEQUENCE_TOKENS: usize = 2048;
23
24#[derive(Debug, Error)]
25pub enum EmbedError {
26    #[error("model directory is missing: {0}")]
27    MissingModelDirectory(PathBuf),
28    #[error("model artifact is missing: {0}")]
29    MissingArtifact(PathBuf),
30    #[error("invalid model configuration: {0}")]
31    InvalidConfig(String),
32    #[error("tokenizer: {0}")]
33    Tokenizer(String),
34    #[error("ONNX Runtime: {0}")]
35    Runtime(String),
36    #[error("model output: {0}")]
37    Output(String),
38}
39
40#[derive(Debug, Clone)]
41pub struct EmbedderConfig {
42    pub model_dir: PathBuf,
43    pub runtime_path: PathBuf,
44    pub provider: ExecutionProvider,
45    pub truncate_dim: usize,
46    pub max_tokens: usize,
47}
48
49impl EmbedderConfig {
50    pub fn new(model_dir: impl Into<PathBuf>, runtime_path: impl Into<PathBuf>) -> Self {
51        Self {
52            model_dir: model_dir.into(),
53            runtime_path: runtime_path.into(),
54            provider: ExecutionProvider::Cpu,
55            truncate_dim: DEFAULT_TRUNCATE_DIM,
56            max_tokens: MAX_SEQUENCE_TOKENS,
57        }
58    }
59
60    pub fn with_provider(mut self, provider: ExecutionProvider) -> Self {
61        self.provider = provider;
62        self
63    }
64
65    pub fn with_truncate_dim(mut self, dim: usize) -> Self {
66        self.truncate_dim = dim;
67        self
68    }
69
70    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
71        self.max_tokens = max_tokens;
72        self
73    }
74}
75
76pub struct Embedder {
77    tokenizer: Tokenizer,
78    session: Mutex<Session>,
79    native_dim: usize,
80    truncate_dim: usize,
81    max_tokens: usize,
82    model_id: String,
83    gemma2: bool,
84}
85
86impl Embedder {
87    pub fn load(config: EmbedderConfig) -> Result<Self, EmbedError> {
88        validate_config(&config)?;
89        let gemma2 = bundle_is_gemma2(&config.model_dir)?;
90        let tokenizer = Tokenizer::from_file(config.model_dir.join("tokenizer.json"))
91            .map_err(|e| EmbedError::Tokenizer(e.to_string()))?;
92        let selection = configure_runtime([RuntimeCandidate::new(
93            config.runtime_path.clone(),
94            CandidateSource::Explicit,
95        )])
96        .map_err(|e| EmbedError::Runtime(e.to_string()))?;
97        init_environment(
98            &selection,
99            &EnvironmentOptions {
100                global_pool: Some(GlobalPool::from_budget()),
101            },
102        )
103        .map_err(EmbedError::Runtime)?;
104        let session = SessionOptions::new(config.provider)
105            .with_threads(Some(1))
106            .build_from_file(&config.model_dir.join("model_q4.onnx"))
107            .map_err(|e| EmbedError::Runtime(e.to_string()))?;
108        Ok(Self {
109            tokenizer,
110            session: Mutex::new(session),
111            native_dim: EMBEDDINGGEMMA_NATIVE_DIM,
112            truncate_dim: config.truncate_dim,
113            max_tokens: config.max_tokens,
114            model_id: if gemma2 {
115                EMBEDDINGGEMMA2_MODEL_ID
116            } else {
117                EMBEDDINGGEMMA_MODEL_ID
118            }
119            .to_owned(),
120            gemma2,
121        })
122    }
123
124    pub fn model_id(&self) -> &str {
125        &self.model_id
126    }
127
128    pub fn native_dim(&self) -> usize {
129        self.native_dim
130    }
131
132    pub fn dimension(&self) -> usize {
133        self.truncate_dim
134    }
135
136    pub fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
137        self.embed_prompt(&format!("title: none | text: {}", self.model_text(text)))
138    }
139
140    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
141        self.embed_prompt(&format!(
142            "task: search result | query: {}",
143            self.model_text(text)
144        ))
145    }
146
147    pub fn embed_batch<'a, I>(&self, texts: I) -> Result<Vec<Vec<f32>>, EmbedError>
148    where
149        I: IntoIterator<Item = &'a str>,
150    {
151        texts.into_iter().map(|text| self.embed(text)).collect()
152    }
153
154    pub fn embed_query_batch<'a, I>(&self, texts: I) -> Result<Vec<Vec<f32>>, EmbedError>
155    where
156        I: IntoIterator<Item = &'a str>,
157    {
158        texts
159            .into_iter()
160            .map(|text| self.embed_query(text))
161            .collect()
162    }
163
164    /// EmbeddingGemma 2's text-only graph must never see modality sentinels;
165    /// their literal spellings are defanged before tokenization.
166    fn model_text<'a>(&self, text: &'a str) -> std::borrow::Cow<'a, str> {
167        if !self.gemma2 || !GEMMA2_MODALITY_SENTINELS.iter().any(|s| text.contains(s)) {
168            return std::borrow::Cow::Borrowed(text);
169        }
170        let mut owned = text.to_owned();
171        for sentinel in GEMMA2_MODALITY_SENTINELS {
172            owned = owned.replace(sentinel, &sentinel.replace('|', "\u{2502}"));
173        }
174        std::borrow::Cow::Owned(owned)
175    }
176
177    fn embed_prompt(&self, prompt: &str) -> Result<Vec<f32>, EmbedError> {
178        let encoding = self
179            .tokenizer
180            .encode(EncodeInput::Single(prompt.to_owned().into()), true)
181            .map_err(|e| EmbedError::Tokenizer(e.to_string()))?;
182        let ids: Vec<i64> = encoding
183            .get_ids()
184            .iter()
185            .take(self.max_tokens)
186            .map(|id| i64::from(*id))
187            .collect();
188        if ids.is_empty() {
189            return Err(EmbedError::Tokenizer("encoding returned no tokens".into()));
190        }
191        let mask: Vec<i64> = encoding
192            .get_attention_mask()
193            .iter()
194            .take(ids.len())
195            .map(|value| i64::from(*value))
196            .collect();
197        let mut session = self
198            .session
199            .lock()
200            .map_err(|_| EmbedError::Runtime("model mutex poisoned".into()))?;
201        let outputs = session
202            .run(ort::inputs![
203                "input_ids" => ort::value::Tensor::from_array(([1usize, ids.len()], ids))
204                    .map_err(|e| EmbedError::Runtime(e.to_string()))?,
205                "attention_mask" => ort::value::Tensor::from_array(([1usize, mask.len()], mask.clone()))
206                    .map_err(|e| EmbedError::Runtime(e.to_string()))?
207            ])
208            .map_err(|e| EmbedError::Runtime(e.to_string()))?;
209        let output = outputs
210            .get("sentence_embedding")
211            .ok_or_else(|| EmbedError::Output("missing sentence_embedding output".into()))?;
212        let (shape, values) = output
213            .try_extract_tensor::<f32>()
214            .map_err(|e| EmbedError::Output(e.to_string()))?;
215        let native = sentence_embedding(shape.as_ref(), values, &mask, self.native_dim)?;
216        Ok(truncate_and_normalize(&native, self.truncate_dim))
217    }
218}
219
220/// Reads only the top-level `model_type`; nested sub-configs (audio, text,
221/// vision) carry their own types and are ignored. A missing config.json is the
222/// original EmbeddingGemma bundle.
223fn bundle_is_gemma2(model_dir: &std::path::Path) -> Result<bool, EmbedError> {
224    let path = model_dir.join("config.json");
225    let Ok(bytes) = std::fs::read(&path) else {
226        return Ok(false);
227    };
228    let config: serde_json::Value = serde_json::from_slice(&bytes)
229        .map_err(|e| EmbedError::InvalidConfig(format!("{}: {e}", path.display())))?;
230    Ok(config["model_type"] == "embedding_gemma2")
231}
232
233fn validate_config(config: &EmbedderConfig) -> Result<(), EmbedError> {
234    if !config.model_dir.is_dir() {
235        return Err(EmbedError::MissingModelDirectory(config.model_dir.clone()));
236    }
237    if config.truncate_dim == 0 || config.truncate_dim > EMBEDDINGGEMMA_NATIVE_DIM {
238        return Err(EmbedError::InvalidConfig(format!(
239            "truncate dimension must be 1..={EMBEDDINGGEMMA_NATIVE_DIM}"
240        )));
241    }
242    if config.max_tokens == 0 || config.max_tokens > MAX_SEQUENCE_TOKENS {
243        return Err(EmbedError::InvalidConfig(format!(
244            "max tokens must be 1..={MAX_SEQUENCE_TOKENS}"
245        )));
246    }
247    for name in ["model_q4.onnx", "model_q4.onnx_data", "tokenizer.json"] {
248        let path = config.model_dir.join(name);
249        if !path.is_file() {
250            return Err(EmbedError::MissingArtifact(path));
251        }
252    }
253    if !config.runtime_path.is_file() {
254        return Err(EmbedError::MissingArtifact(config.runtime_path.clone()));
255    }
256    Ok(())
257}
258
259fn sentence_embedding(
260    shape: &[i64],
261    values: &[f32],
262    mask: &[i64],
263    native_dim: usize,
264) -> Result<Vec<f32>, EmbedError> {
265    if shape == [1, native_dim as i64] && values.len() == native_dim {
266        return Ok(normalize(values.to_vec()));
267    }
268    if shape.len() == 3
269        && shape[0] == 1
270        && shape[2] == native_dim as i64
271        && shape[1] > 0
272        && values.len() == shape[1] as usize * native_dim
273    {
274        let tokens = shape[1] as usize;
275        let mut pooled = vec![0.0; native_dim];
276        let mut count = 0.0f32;
277        for token in 0..tokens {
278            if mask.get(token).copied().unwrap_or(1) == 0 {
279                continue;
280            }
281            count += 1.0;
282            for (dim, value) in pooled.iter_mut().enumerate() {
283                *value += values[token * native_dim + dim];
284            }
285        }
286        if count == 0.0 {
287            return Err(EmbedError::Output(
288                "attention mask contains no tokens".into(),
289            ));
290        }
291        for value in &mut pooled {
292            *value /= count;
293        }
294        return Ok(normalize(pooled));
295    }
296    Err(EmbedError::Output(format!(
297        "expected [1,{native_dim}] or [1,tokens,{native_dim}], got shape {shape:?}"
298    )))
299}
300
301fn normalize(mut values: Vec<f32>) -> Vec<f32> {
302    let norm = values.iter().map(|value| value * value).sum::<f32>().sqrt();
303    if norm > 0.0 {
304        for value in &mut values {
305            *value /= norm;
306        }
307    }
308    values
309}
310
311pub fn truncate_and_normalize(native: &[f32], dimension: usize) -> Vec<f32> {
312    let mut values = native[..native.len().min(dimension)].to_vec();
313    values.resize(dimension, 0.0);
314    normalize(values)
315}
rightkit_embed/embedder.rs

rightkit_embed/
embedder.rs
