llm_kernel/embedding/
qwen3.rs

1//! Qwen3 embedding provider via fastembed-rs candle backend.
2//!
3//! Uses the candle-nn pure Rust inference engine (no ONNX Runtime).
4//! Models are downloaded from HuggingFace on first use.
5//!
6//! Supported repos: `Qwen/Qwen3-Embedding-0.6B`, `Qwen/Qwen3-Embedding-8B`,
7//! `Qwen/Qwen3-VL-Embedding-2B` (text-only mode).
8//!
9//! ```ignore
10//! use llm_kernel::embedding::Qwen3Provider;
11//! use llm_kernel::embedding::EmbeddingProvider;
12//!
13//! let provider = Qwen3Provider::new("Qwen/Qwen3-Embedding-0.6B")?;
14//! let result = provider.embed("hello world")?;
15//! ```
16
17use crate::embedding::types::{EmbeddingProvider, EmbeddingResult};
18use crate::error::{KernelError, Result};
19
20/// Qwen3 embedding provider backed by candle-nn.
21///
22/// Unlike [`FastembedProvider`](super::FastembedProvider) (ONNX), this uses
23/// candle for pure Rust GPU/CPU inference. The `embed()` method takes `&self`,
24/// so no `Mutex` is needed.
25pub struct Qwen3Provider {
26    inner: fastembed::Qwen3TextEmbedding,
27    model_id: String,
28    dim: usize,
29}
30
31/// Default HuggingFace repo for Qwen3-Embedding-0.6B.
32pub const QWEN3_EMBEDDING_0_6B: &str = "Qwen/Qwen3-Embedding-0.6B";
33
34/// Default HuggingFace repo for Qwen3-Embedding-8B.
35pub const QWEN3_EMBEDDING_8B: &str = "Qwen/Qwen3-Embedding-8B";
36
37/// Default HuggingFace repo for Qwen3-VL-Embedding-2B (text-only mode).
38pub const QWEN3_VL_EMBEDDING_2B: &str = "Qwen/Qwen3-VL-Embedding-2B";
39
40/// Default max sequence length for Qwen3 models.
41const DEFAULT_MAX_LENGTH: usize = 512;
42
43impl Qwen3Provider {
44    /// Create a new provider using CPU with F32 precision.
45    ///
46    /// Downloads the model from HuggingFace on first call (cached locally).
47    pub fn new(model_id: &str) -> Result<Self> {
48        Self::with_options(
49            model_id,
50            candle_core::Device::Cpu,
51            candle_core::DType::F32,
52            DEFAULT_MAX_LENGTH,
53        )
54    }
55
56    /// Create with custom device (GPU), dtype, and max sequence length.
57    pub fn with_options(
58        model_id: &str,
59        device: candle_core::Device,
60        dtype: candle_core::DType,
61        max_length: usize,
62    ) -> Result<Self> {
63        let te = fastembed::Qwen3TextEmbedding::from_hf(model_id, &device, dtype, max_length)
64            .map_err(KernelError::embedding)?;
65        let dim = te.config().hidden_size;
66        Ok(Self {
67            inner: te,
68            model_id: model_id.to_string(),
69            dim,
70        })
71    }
72
73    /// The HuggingFace model repo ID.
74    pub fn model_id(&self) -> &str {
75        &self.model_id
76    }
77}
78
79impl EmbeddingProvider for Qwen3Provider {
80    fn dim(&self) -> usize {
81        self.dim
82    }
83
84    fn name(&self) -> &str {
85        &self.model_id
86    }
87
88    fn embed(&self, text: &str) -> Result<EmbeddingResult> {
89        let embeddings = self.inner.embed(&[text]).map_err(KernelError::embedding)?;
90        let vector = embeddings
91            .into_iter()
92            .next()
93            .ok_or_else(|| KernelError::Embedding("empty embedding output".into()))?;
94
95        let preview = if text.len() > 64 {
96            format!("{}…", &text[..64])
97        } else {
98            text.to_string()
99        };
100        Ok(EmbeddingResult {
101            vector,
102            text_preview: preview,
103        })
104    }
105
106    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingResult>> {
107        if texts.is_empty() {
108            return Ok(vec![]);
109        }
110        let embeddings = self.inner.embed(texts).map_err(KernelError::embedding)?;
111        Ok(embeddings
112            .into_iter()
113            .zip(texts.iter())
114            .map(|(vector, &text)| {
115                let preview = if text.len() > 64 {
116                    format!("{}…", &text[..64])
117                } else {
118                    text.to_string()
119                };
120                EmbeddingResult {
121                    vector,
122                    text_preview: preview,
123                }
124            })
125            .collect())
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn model_id_constants() {
135        assert_eq!(QWEN3_EMBEDDING_0_6B, "Qwen/Qwen3-Embedding-0.6B");
136        assert_eq!(QWEN3_EMBEDDING_8B, "Qwen/Qwen3-Embedding-8B");
137        assert_eq!(QWEN3_VL_EMBEDDING_2B, "Qwen/Qwen3-VL-Embedding-2B");
138    }
139
140    #[test]
141    #[ignore = "requires model download"]
142    fn embed_with_qwen3_0_6b() {
143        let provider = Qwen3Provider::new(QWEN3_EMBEDDING_0_6B).unwrap();
144        let result = provider.embed("hello world").unwrap();
145        // Qwen3-Embedding-0.6B has hidden_size that the config reports
146        assert!(!result.vector.is_empty());
147        assert_eq!(result.vector.len(), provider.dim());
148    }
149
150    #[test]
151    #[ignore = "requires model download"]
152    fn embed_batch_with_qwen3() {
153        let provider = Qwen3Provider::new(QWEN3_EMBEDDING_0_6B).unwrap();
154        let results = provider
155            .embed_batch(&["hello", "world", "foo bar"])
156            .unwrap();
157        assert_eq!(results.len(), 3);
158        for r in &results {
159            assert_eq!(r.vector.len(), provider.dim());
160        }
161    }
162}
llm_kernel/embedding/qwen3.rs

llm_kernel/embedding/
qwen3.rs
