llm_kernel/embedding/
nomic_moe.rs

1//! Nomic V2 MoE 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//! `nomic-ai/nomic-embed-text-v2-moe` — 475M total / 305M active params,
7//! 8 experts with top-2 routing, hidden_size=768.
8//!
9//! ```ignore
10//! use llm_kernel::embedding::NomicMoeProvider;
11//! use llm_kernel::embedding::EmbeddingProvider;
12//!
13//! let provider = NomicMoeProvider::new()?;
14//! let result = provider.embed("hello world")?;
15//! ```
16
17use crate::embedding::types::{EmbeddingProvider, EmbeddingResult};
18use crate::error::{KernelError, Result};
19
20/// Nomic V2 MoE 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 NomicMoeProvider {
26    inner: fastembed::NomicV2MoeTextEmbedding,
27    model_id: String,
28    dim: usize,
29}
30
31/// Default HuggingFace repo for nomic-embed-text-v2-moe.
32pub const NOMIC_EMBED_TEXT_V2_MOE: &str = "nomic-ai/nomic-embed-text-v2-moe";
33
34/// Default max sequence length for Nomic V2 MoE.
35const DEFAULT_MAX_LENGTH: usize = 512;
36
37impl NomicMoeProvider {
38    /// Create a new provider using CPU with F32 precision.
39    ///
40    /// Downloads the model from HuggingFace on first call (cached locally).
41    pub fn new() -> Result<Self> {
42        Self::with_options(
43            NOMIC_EMBED_TEXT_V2_MOE,
44            candle_core::Device::Cpu,
45            candle_core::DType::F32,
46            DEFAULT_MAX_LENGTH,
47        )
48    }
49
50    /// Create with custom repo, device, dtype, and max sequence length.
51    pub fn with_options(
52        model_id: &str,
53        device: candle_core::Device,
54        dtype: candle_core::DType,
55        max_length: usize,
56    ) -> Result<Self> {
57        let te = fastembed::NomicV2MoeTextEmbedding::from_hf(model_id, &device, dtype, max_length)
58            .map_err(KernelError::embedding)?;
59        let dim = te.config().hidden_size;
60        Ok(Self {
61            inner: te,
62            model_id: model_id.to_string(),
63            dim,
64        })
65    }
66
67    /// The HuggingFace model repo ID.
68    pub fn model_id(&self) -> &str {
69        &self.model_id
70    }
71}
72
73impl EmbeddingProvider for NomicMoeProvider {
74    fn dim(&self) -> usize {
75        self.dim
76    }
77
78    fn name(&self) -> &str {
79        &self.model_id
80    }
81
82    fn embed(&self, text: &str) -> Result<EmbeddingResult> {
83        let embeddings = self.inner.embed(&[text]).map_err(KernelError::embedding)?;
84        let vector = embeddings
85            .into_iter()
86            .next()
87            .ok_or_else(|| KernelError::Embedding("empty embedding output".into()))?;
88
89        let preview = if text.len() > 64 {
90            format!("{}…", &text[..64])
91        } else {
92            text.to_string()
93        };
94        Ok(EmbeddingResult {
95            vector,
96            text_preview: preview,
97        })
98    }
99
100    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingResult>> {
101        if texts.is_empty() {
102            return Ok(vec![]);
103        }
104        let embeddings = self.inner.embed(texts).map_err(KernelError::embedding)?;
105        Ok(embeddings
106            .into_iter()
107            .zip(texts.iter())
108            .map(|(vector, &text)| {
109                let preview = if text.len() > 64 {
110                    format!("{}…", &text[..64])
111                } else {
112                    text.to_string()
113                };
114                EmbeddingResult {
115                    vector,
116                    text_preview: preview,
117                }
118            })
119            .collect())
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn model_id_constant() {
129        assert_eq!(NOMIC_EMBED_TEXT_V2_MOE, "nomic-ai/nomic-embed-text-v2-moe");
130    }
131
132    #[test]
133    #[ignore = "requires model download"]
134    fn embed_with_nomic_moe() {
135        let provider = NomicMoeProvider::new().unwrap();
136        let result = provider.embed("hello world").unwrap();
137        // nomic-embed-text-v2-moe hidden_size = 768
138        assert_eq!(result.vector.len(), 768);
139        assert_eq!(result.vector.len(), provider.dim());
140    }
141
142    #[test]
143    #[ignore = "requires model download"]
144    fn embed_batch_with_nomic_moe() {
145        let provider = NomicMoeProvider::new().unwrap();
146        let results = provider
147            .embed_batch(&["hello", "world", "foo bar"])
148            .unwrap();
149        assert_eq!(results.len(), 3);
150        for r in &results {
151            assert_eq!(r.vector.len(), 768);
152        }
153    }
154}
llm_kernel/embedding/nomic_moe.rs

llm_kernel/embedding/
nomic_moe.rs
