clawft_kernel/
embedding_qwen3.rs

1//! Qwen3-Embedding-0.6B provider (ADR-059) — the preferred ECC embedder.
2//!
3//! Unlike the BERT [`OnnxEmbeddingProvider`](crate::embedding_onnx::OnnxEmbeddingProvider),
4//! Qwen3-Embedding is a **causal decoder-with-KV-cache** ONNX export. This
5//! provider therefore:
6//!
7//! - feeds `input_ids` / `attention_mask` / `position_ids` (`0..seq_len`) **plus
8//!   empty, dtype-matched `past_key_values.*` tensors** for every layer (shape
9//!   `[1, n_kv_heads, 0, head_dim]`, derived from the graph — never hardcoded);
10//! - pools the **last non-pad token** hidden state (not mean/CLS);
11//! - applies **MRL**: slice to the first [`QWEN3_DIMS`] dims, *then* L2-renormalize;
12//! - uses the HF `tokenizers` crate (Qwen3 BPE from `tokenizer.json`);
13//! - exposes [`embed_query`](Qwen3EmbeddingProvider::embed_query) for the
14//!   asymmetric query lane (prepends the Qwen3 instruction prefix); the
15//!   [`EmbeddingProvider::embed`] trait method is document mode.
16//!
17//! Without the `onnx-embeddings` feature, the model file, or the tokenizer, the
18//! provider degrades gracefully to a [`MockEmbeddingProvider`] so builds and
19//! tests never require the ~1.2 GB artifact. The BERT provider stays intact as a
20//! lower-priority fallback (see [`select_embedding_provider`](crate::embedding::select_embedding_provider)).
21
22use std::path::{Path, PathBuf};
23#[cfg(feature = "onnx-embeddings")]
24use std::sync::{Arc, Mutex};
25
26use async_trait::async_trait;
27
28use crate::embedding::{EmbeddingError, EmbeddingProvider, MockEmbeddingProvider};
29
30/// MRL target dimensionality — the empirical sweet spot (ADR-059, recall@10 0.983).
31pub const QWEN3_DIMS: usize = 512;
32
33/// Canonical model identifier (pins the consistency contract; ADR-059).
34pub const QWEN3_MODEL_NAME: &str = "Qwen3-Embedding-0.6B";
35
36/// Qwen3 asymmetric **query** instruction prefix. Documents are embedded raw;
37/// queries are prefixed with this (`Instruct: {task}\nQuery: {q}` format). The
38/// task description is the code-search instruction pinned by the lab probe set
39/// (ADR-059 consistency contract) so the lab and prod share one vector space;
40/// `embed_query` appends the raw query directly after the trailing space.
41pub const QUERY_INSTRUCTION: &str = "Instruct: Given a code-search query, retrieve the relevant code, log, diff, traceback or config chunk.\nQuery: ";
42
43// ---------------------------------------------------------------------------
44// Pure helpers (always compiled + unit-tested; only wired into production
45// inference under `onnx-embeddings`, hence the conditional dead-code allow).
46// ---------------------------------------------------------------------------
47
48/// In-place L2 normalisation.
49#[cfg_attr(not(feature = "onnx-embeddings"), allow(dead_code))]
50fn l2_normalize(vec: &mut [f32]) {
51    let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
52    if norm > 0.0 {
53        vec.iter_mut().for_each(|x| *x /= norm);
54    }
55}
56
57/// Last-token pooling: return the hidden state at the **last non-pad position**.
58///
59/// `hidden` is the flattened `(seq, hidden_dim)` last-hidden-state for a single
60/// sequence (batch 1). The last attended position is located via
61/// `attention_mask`; if the mask is empty/all-zero the final row is used.
62#[cfg_attr(not(feature = "onnx-embeddings"), allow(dead_code))]
63fn last_token_pool(
64    hidden: &[f32],
65    seq: usize,
66    hidden_dim: usize,
67    attention_mask: &[i64],
68) -> Vec<f32> {
69    let last = attention_mask
70        .iter()
71        .take(seq)
72        .rposition(|&m| m != 0)
73        .unwrap_or(seq.saturating_sub(1));
74    let start = last * hidden_dim;
75    let end = start + hidden_dim;
76    if end <= hidden.len() {
77        hidden[start..end].to_vec()
78    } else {
79        // Defensive: shape mismatch — fall back to a zero vector.
80        vec![0.0; hidden_dim]
81    }
82}
83
84/// MRL: slice to the first `dims` dimensions **then** L2-renormalize (never the
85/// reverse). Shorter inputs are zero-padded so the output is always `dims` long.
86#[cfg_attr(not(feature = "onnx-embeddings"), allow(dead_code))]
87fn mrl_slice_l2norm(full: &[f32], dims: usize) -> Vec<f32> {
88    let mut v: Vec<f32> = full.iter().take(dims).copied().collect();
89    if v.len() < dims {
90        v.resize(dims, 0.0);
91    }
92    l2_normalize(&mut v);
93    v
94}
95
96// ---------------------------------------------------------------------------
97// KV-cache spec (derived from the ONNX graph, never hardcoded)
98// ---------------------------------------------------------------------------
99
100/// One `past_key_values.*` graph input: its name and the concrete head geometry.
101#[cfg(feature = "onnx-embeddings")]
102#[derive(Debug, Clone)]
103struct KvInput {
104    name: String,
105    n_kv_heads: i64,
106    head_dim: i64,
107    fp16: bool,
108}
109
110// ---------------------------------------------------------------------------
111// Qwen3EmbeddingProvider
112// ---------------------------------------------------------------------------
113
114/// Qwen3-Embedding-0.6B provider — decoder-ONNX, last-token pooling, MRL-512.
115pub struct Qwen3EmbeddingProvider {
116    /// Path to `model_fp16.onnx`.
117    model_path: PathBuf,
118    /// Path to `tokenizer.json`.
119    tokenizer_path: PathBuf,
120    /// Output dimensionality (always [`QWEN3_DIMS`]).
121    dimensions: usize,
122    /// Model name (real name when live, `*-mock-fallback` when degraded).
123    model_name: String,
124    /// Whether real decoder-ONNX inference is active (model + tokenizer loaded).
125    runtime_available: bool,
126    /// Degraded fallback (graceful Mock when no model/tokenizer/feature).
127    fallback: MockEmbeddingProvider,
128    /// Qwen3 BPE tokenizer loaded from `tokenizer.json`.
129    #[cfg(feature = "onnx-embeddings")]
130    tokenizer: Option<tokenizers::Tokenizer>,
131    /// Decoder-ONNX session (`run` needs `&mut`, hence the mutex).
132    #[cfg(feature = "onnx-embeddings")]
133    session: Option<Arc<Mutex<ort::session::Session>>>,
134    /// Empty-KV-cache spec derived from the graph inputs.
135    #[cfg(feature = "onnx-embeddings")]
136    kv_spec: Vec<KvInput>,
137}
138
139impl Qwen3EmbeddingProvider {
140    /// Create a provider pointing at `model_fp16.onnx` + `tokenizer.json`.
141    ///
142    /// Degrades to Mock when the feature is off, files are missing, or loading
143    /// fails — never panics, never requires the model to construct.
144    pub fn new(model_path: impl Into<PathBuf>, tokenizer_path: impl Into<PathBuf>) -> Self {
145        let model_path = model_path.into();
146        let tokenizer_path = tokenizer_path.into();
147
148        #[cfg(feature = "onnx-embeddings")]
149        let tokenizer = Self::try_load_tokenizer(&tokenizer_path);
150        #[cfg(feature = "onnx-embeddings")]
151        let session = Self::try_load_session(&model_path);
152        #[cfg(feature = "onnx-embeddings")]
153        let kv_spec = session
154            .as_ref()
155            .and_then(|s| s.lock().ok().map(|g| Self::derive_kv_spec(&g)))
156            .unwrap_or_default();
157
158        #[cfg(feature = "onnx-embeddings")]
159        let runtime_available = tokenizer.is_some() && session.is_some();
160        #[cfg(not(feature = "onnx-embeddings"))]
161        let runtime_available = false;
162
163        Self {
164            model_name: if runtime_available {
165                QWEN3_MODEL_NAME.to_string()
166            } else {
167                format!("{QWEN3_MODEL_NAME}-mock-fallback")
168            },
169            model_path,
170            tokenizer_path,
171            dimensions: QWEN3_DIMS,
172            runtime_available,
173            fallback: MockEmbeddingProvider::new(QWEN3_DIMS),
174            #[cfg(feature = "onnx-embeddings")]
175            tokenizer,
176            #[cfg(feature = "onnx-embeddings")]
177            session,
178            #[cfg(feature = "onnx-embeddings")]
179            kv_spec,
180        }
181    }
182
183    /// Whether real decoder-ONNX inference is active (vs. Mock fallback).
184    pub fn is_runtime_available(&self) -> bool {
185        self.runtime_available
186    }
187
188    /// Path to the configured ONNX model file.
189    pub fn model_path(&self) -> &Path {
190        &self.model_path
191    }
192
193    /// Path to the configured tokenizer file.
194    pub fn tokenizer_path(&self) -> &Path {
195        &self.tokenizer_path
196    }
197
198    /// Embed a **query** (asymmetric lane): prepends the Qwen3 instruction prefix.
199    pub async fn embed_query(&self, query: &str) -> Result<Vec<f32>, EmbeddingError> {
200        #[cfg(feature = "onnx-embeddings")]
201        if self.runtime_available {
202            let prompted = format!("{QUERY_INSTRUCTION}{query}");
203            return self.qwen3_embed(&prompted);
204        }
205        self.fallback.embed(query).await
206    }
207
208    #[cfg(feature = "onnx-embeddings")]
209    fn try_load_tokenizer(path: &Path) -> Option<tokenizers::Tokenizer> {
210        if !path.exists() {
211            tracing::debug!("Qwen3 tokenizer.json not found at {}", path.display());
212            return None;
213        }
214        match tokenizers::Tokenizer::from_file(path) {
215            Ok(t) => {
216                tracing::info!("Qwen3 tokenizer loaded from {}", path.display());
217                Some(t)
218            }
219            Err(e) => {
220                tracing::warn!("Failed to load Qwen3 tokenizer: {e}");
221                None
222            }
223        }
224    }
225
226    #[cfg(feature = "onnx-embeddings")]
227    fn try_load_session(model_path: &Path) -> Option<Arc<Mutex<ort::session::Session>>> {
228        if !model_path.exists() {
229            tracing::debug!("Qwen3 ONNX model not found at {}", model_path.display());
230            return None;
231        }
232        match ort::session::Session::builder()
233            .and_then(|mut builder| builder.commit_from_file(model_path))
234        {
235            Ok(session) => {
236                tracing::info!("Qwen3 ONNX session loaded from {}", model_path.display());
237                Some(Arc::new(Mutex::new(session)))
238            }
239            Err(e) => {
240                tracing::warn!("Failed to load Qwen3 ONNX session: {e}");
241                None
242            }
243        }
244    }
245
246    /// Derive the empty-KV-cache spec from the graph's `past_key_values.*`
247    /// inputs — layer count, kv-heads, head-dim, and dtype are all read from the
248    /// model, never hardcoded.
249    #[cfg(feature = "onnx-embeddings")]
250    fn derive_kv_spec(session: &ort::session::Session) -> Vec<KvInput> {
251        let mut out = Vec::new();
252        for input in session.inputs() {
253            let name = input.name();
254            if !name.starts_with("past_key_values") && !name.starts_with("past_key") {
255                continue;
256            }
257            if let ort::value::ValueType::Tensor { ty, shape, .. } = input.dtype() {
258                // [batch, n_kv_heads, past_len, head_dim]
259                if shape.len() == 4 {
260                    out.push(KvInput {
261                        name: name.to_string(),
262                        n_kv_heads: shape[1],
263                        head_dim: shape[3],
264                        fp16: matches!(ty, ort::value::TensorElementType::Float16),
265                    });
266                }
267            }
268        }
269        out
270    }
271
272    /// Run real decoder-ONNX inference: tokenize → feed ids/mask/position_ids +
273    /// empty past_key_values → last-token pool → MRL-512 slice + L2-norm.
274    #[cfg(feature = "onnx-embeddings")]
275    fn qwen3_embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
276        use std::borrow::Cow;
277
278        use ort::session::SessionInputValue;
279        use ort::value::Tensor;
280
281        let session = self
282            .session
283            .as_ref()
284            .ok_or_else(|| EmbeddingError::BackendError("Qwen3 session not loaded".to_string()))?;
285        let tokenizer = self.tokenizer.as_ref().ok_or_else(|| {
286            EmbeddingError::BackendError("Qwen3 tokenizer not loaded".to_string())
287        })?;
288
289        let encoding = tokenizer
290            .encode(text, true)
291            .map_err(|e| EmbeddingError::BackendError(format!("tokenize error: {e}")))?;
292        let mut input_ids: Vec<i64> = encoding.get_ids().iter().map(|&u| u as i64).collect();
293        let mut attention_mask: Vec<i64> = encoding
294            .get_attention_mask()
295            .iter()
296            .map(|&u| u as i64)
297            .collect();
298        // Append the Qwen3 EOS (`<|im_end|>`) for last-token pooling if the
299        // tokenizer's post-processor didn't already terminate on it. The fp16
300        // export pools the hidden state at this EOS position; omitting it pools
301        // the wrong token and destroys lab parity (~0.72 cosine without it).
302        if let Some(eos) = tokenizer.token_to_id("<|im_end|>").map(|u| u as i64)
303            && input_ids.last() != Some(&eos)
304        {
305            input_ids.push(eos);
306            attention_mask.push(1);
307        }
308        let seq_len = input_ids.len();
309        if seq_len == 0 {
310            return Ok(vec![0.0; self.dimensions]);
311        }
312        let position_ids: Vec<i64> = (0..seq_len as i64).collect();
313        let shape = vec![1_i64, seq_len as i64];
314
315        let ids_t = Tensor::from_array((shape.clone(), input_ids))
316            .map_err(|e| EmbeddingError::BackendError(format!("input_ids tensor: {e}")))?;
317        let mask_t = Tensor::from_array((shape.clone(), attention_mask.clone()))
318            .map_err(|e| EmbeddingError::BackendError(format!("attention_mask tensor: {e}")))?;
319        let pos_t = Tensor::from_array((shape, position_ids))
320            .map_err(|e| EmbeddingError::BackendError(format!("position_ids tensor: {e}")))?;
321
322        let mut inputs = ort::inputs![
323            "input_ids" => ids_t,
324            "attention_mask" => mask_t,
325            "position_ids" => pos_t,
326        ];
327
328        // Empty, dtype-matched KV cache for every layer (past_len = 0). ort
329        // forbids a 0-dim in the (shape, data) `from_array` path, so allocate a
330        // 0-element tensor via the default CPU allocator instead.
331        let allocator = ort::memory::Allocator::default();
332        for kv in &self.kv_spec {
333            let kv_shape = ort::value::Shape::new([1, kv.n_kv_heads, 0, kv.head_dim]);
334            let value: SessionInputValue<'_> = if kv.fp16 {
335                Tensor::<half::f16>::new(&allocator, kv_shape)
336                    .map_err(|e| EmbeddingError::BackendError(format!("kv tensor: {e}")))?
337                    .into()
338            } else {
339                Tensor::<f32>::new(&allocator, kv_shape)
340                    .map_err(|e| EmbeddingError::BackendError(format!("kv tensor: {e}")))?
341                    .into()
342            };
343            inputs.push((Cow::from(kv.name.clone()), value));
344        }
345
346        let mut guard = session.lock().map_err(|_| {
347            EmbeddingError::BackendError("Qwen3 session mutex poisoned".to_string())
348        })?;
349        let outputs = guard
350            .run(inputs)
351            .map_err(|e| EmbeddingError::BackendError(format!("inference error: {e}")))?;
352
353        let output = match outputs.get("last_hidden_state") {
354            Some(v) => v,
355            None if outputs.len() > 0 => &outputs[0],
356            None => {
357                return Err(EmbeddingError::BackendError("no output tensor".to_string()));
358            }
359        };
360
361        // fp16 model → output is likely fp16; try f32 first, then f16→f32.
362        let (out_shape, hidden): (Vec<i64>, Vec<f32>) = match output.try_extract_tensor::<f32>() {
363            Ok((s, d)) => (s.to_vec(), d.to_vec()),
364            Err(_) => {
365                let (s, d) = output
366                    .try_extract_tensor::<half::f16>()
367                    .map_err(|e| EmbeddingError::BackendError(format!("extract error: {e}")))?;
368                (s.to_vec(), d.iter().map(|h| h.to_f32()).collect())
369            }
370        };
371
372        if out_shape.len() < 3 {
373            return Err(EmbeddingError::BackendError(format!(
374                "unexpected output shape: {out_shape:?}"
375            )));
376        }
377        let seq = out_shape[out_shape.len() - 2] as usize;
378        let hidden_dim = out_shape[out_shape.len() - 1] as usize;
379
380        let pooled = last_token_pool(&hidden, seq, hidden_dim, &attention_mask);
381        Ok(mrl_slice_l2norm(&pooled, self.dimensions))
382    }
383}
384
385#[async_trait]
386impl EmbeddingProvider for Qwen3EmbeddingProvider {
387    async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
388        #[cfg(feature = "onnx-embeddings")]
389        if self.runtime_available {
390            return self.qwen3_embed(text);
391        }
392        self.fallback.embed(text).await
393    }
394
395    async fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
396        let mut results = Vec::with_capacity(texts.len());
397        for text in texts {
398            results.push(self.embed(text).await?);
399        }
400        Ok(results)
401    }
402
403    /// Asymmetric query lane — delegates to the inherent
404    /// [`Self::embed_query`] so trait-object call sites
405    /// (`SessionView::graft_text`) get the instruction prefix.
406    async fn embed_query(&self, query: &str) -> Result<Vec<f32>, EmbeddingError> {
407        // Call the inherent method body directly (avoid infinite recursion on
408        // the trait method of the same name).
409        #[cfg(feature = "onnx-embeddings")]
410        if self.runtime_available {
411            let prompted = format!("{QUERY_INSTRUCTION}{query}");
412            return self.qwen3_embed(&prompted);
413        }
414        self.fallback.embed(query).await
415    }
416
417    fn dimensions(&self) -> usize {
418        self.dimensions
419    }
420
421    fn model_name(&self) -> &str {
422        &self.model_name
423    }
424}
425
426// ── Tests (no model required; split to a sibling file for the <500-line rule) ─
427
428#[cfg(test)]
429#[path = "embedding_qwen3_tests.rs"]
430mod tests;
clawft_kernel/embedding_qwen3.rs

clawft_kernel/
embedding_qwen3.rs
