elph_core/memz/
embed.rs

1//! Optional local embedding backends for [`EmbedFn`].
2//!
3//! Enable with the `fastembed` feature (all-MiniLM-L6-v2 and other ONNX models).
4
5use std::path::PathBuf;
6
7use super::store::EmbedFn;
8
9#[cfg(feature = "fastembed")]
10use std::str::FromStr;
11#[cfg(feature = "fastembed")]
12use std::sync::{Arc, Mutex};
13
14#[cfg(feature = "fastembed")]
15use super::store::EmbedFuture;
16
17/// Default embedding model when none is configured.
18pub const DEFAULT_EMBED_MODEL: &str = "AllMiniLML6V2";
19
20/// Options for the fastembed-backed local embedder.
21#[derive(Debug, Clone)]
22pub struct FastEmbedOptions {
23    /// Use the quantized ONNX variant when available (default: true).
24    pub quantized: bool,
25    /// Model name — fastembed [`EmbeddingModel`] debug name or common HF alias.
26    pub model: Option<String>,
27    /// ONNX/tokenizer cache directory (default: fastembed's `.fastembed_cache`).
28    pub cache_dir: Option<PathBuf>,
29    /// Show Hugging Face download progress (default: true).
30    pub show_download_progress: Option<bool>,
31}
32
33impl Default for FastEmbedOptions {
34    fn default() -> Self {
35        Self {
36            quantized: true,
37            model: None,
38            cache_dir: None,
39            show_download_progress: None,
40        }
41    }
42}
43
44impl FastEmbedOptions {
45    pub fn model(mut self, model: impl Into<String>) -> Self {
46        self.model = Some(model.into());
47        self
48    }
49
50    pub fn cache_dir(mut self, path: impl Into<PathBuf>) -> Self {
51        self.cache_dir = Some(path.into());
52        self
53    }
54
55    pub fn quantized(mut self, quantized: bool) -> Self {
56        self.quantized = quantized;
57        self
58    }
59
60    pub fn show_download_progress(mut self, show: bool) -> Self {
61        self.show_download_progress = Some(show);
62        self
63    }
64}
65
66/// Resolve a user-facing model name to a fastembed [`EmbeddingModel`].
67#[cfg(feature = "fastembed")]
68pub fn resolve_embedding_model(name: &str, quantized: bool) -> Result<fastembed::EmbeddingModel, String> {
69    use fastembed::EmbeddingModel;
70
71    let trimmed = name.trim();
72    let canonical = alias_model_name(trimmed);
73    let mut model = EmbeddingModel::from_str(canonical)?;
74
75    if quantized {
76        model = prefer_quantized_variant(model);
77    }
78
79    Ok(model)
80}
81
82#[cfg(feature = "fastembed")]
83fn alias_model_name(name: &str) -> &str {
84    match name.to_ascii_lowercase().as_str() {
85        "sentence-transformers/all-minilm-l6-v2" | "all-minilm-l6-v2" => "AllMiniLML6V2",
86        "sentence-transformers/all-minilm-l12-v2" | "all-minilm-l12-v2" => "AllMiniLML12V2",
87        "sentence-transformers/all-mpnet-base-v2" => "AllMpnetBaseV2",
88        "baai/bge-small-en-v1.5" | "bge-small-en-v1.5" => "BGESmallENV15",
89        "baai/bge-base-en-v1.5" | "bge-base-en-v1.5" => "BGEBaseENV15",
90        "baai/bge-large-en-v1.5" | "bge-large-en-v1.5" => "BGELargeENV15",
91        "nomic-ai/nomic-embed-text-v1" => "NomicEmbedTextV1",
92        "nomic-ai/nomic-embed-text-v1.5" => "NomicEmbedTextV15",
93        "xenova/all-minilm-l6-v2" => "AllMiniLML6V2Q",
94        _ => name,
95    }
96}
97
98#[cfg(feature = "fastembed")]
99fn prefer_quantized_variant(model: fastembed::EmbeddingModel) -> fastembed::EmbeddingModel {
100    use fastembed::EmbeddingModel;
101    use std::str::FromStr;
102
103    let debug = format!("{model:?}");
104    if debug.ends_with('Q') {
105        return model;
106    }
107
108    let q_name = format!("{debug}Q");
109    EmbeddingModel::from_str(&q_name).unwrap_or(model)
110}
111
112/// Embedding output dimensions for a resolved model.
113#[cfg(feature = "fastembed")]
114pub fn embedding_dims(model: &fastembed::EmbeddingModel) -> u32 {
115    use fastembed::TextEmbedding;
116
117    TextEmbedding::get_model_info(model)
118        .map(|info| info.dim as u32)
119        .unwrap_or(super::util::DEFAULT_EMBEDDING_DIMS)
120}
121
122/// Create a shared local embedder using [fastembed](https://github.com/Anush008/fastembed-rs).
123///
124/// Default model: **AllMiniLML6V2** (quantized by default → `AllMiniLML6V2Q`).
125/// Inference runs on a blocking thread pool; safe to call from async contexts.
126#[cfg(feature = "fastembed")]
127pub fn create_fastembed(options: FastEmbedOptions) -> anyhow::Result<EmbedFn> {
128    use fastembed::{TextEmbedding, TextInitOptions};
129
130    let model_name = options.model.unwrap_or_else(|| DEFAULT_EMBED_MODEL.to_string());
131
132    let embedding_model =
133        resolve_embedding_model(&model_name, options.quantized).map_err(|e| anyhow::anyhow!("{e}"))?;
134
135    let expected_dims = embedding_dims(&embedding_model) as usize;
136
137    let mut init = TextInitOptions::new(embedding_model);
138    if let Some(dir) = options.cache_dir {
139        init = init.with_cache_dir(dir);
140    }
141    if let Some(show) = options.show_download_progress {
142        init = init.with_show_download_progress(show);
143    }
144
145    let model = TextEmbedding::try_new(init)?;
146
147    let shared = Arc::new(Mutex::new(model));
148    Ok(Arc::new(move |text: &str| {
149        let shared = Arc::clone(&shared);
150        let text = text.to_string();
151        Box::pin(async move {
152            let vec = tokio::task::spawn_blocking(move || {
153                let mut model = shared
154                    .lock()
155                    .map_err(|e| anyhow::anyhow!("embedder lock poisoned: {e}"))?;
156                let embeddings = model.embed(vec![text], None)?;
157                embeddings
158                    .into_iter()
159                    .next()
160                    .ok_or_else(|| anyhow::anyhow!("fastembed returned no vectors"))
161            })
162            .await??;
163            if vec.len() != expected_dims {
164                anyhow::bail!("expected {expected_dims}-dim embedding, got {}", vec.len());
165            }
166            Ok(vec)
167        }) as EmbedFuture
168    }))
169}
170
171#[cfg(not(feature = "fastembed"))]
172pub fn create_fastembed(_options: FastEmbedOptions) -> anyhow::Result<EmbedFn> {
173    anyhow::bail!("fastembed embedder requires the `fastembed` feature");
174}
175
176#[cfg(feature = "fastembed")]
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use fastembed::EmbeddingModel;
181
182    #[test]
183    fn resolves_hf_alias() {
184        let m = resolve_embedding_model("sentence-transformers/all-MiniLM-L6-v2", false).unwrap();
185        assert_eq!(m, EmbeddingModel::AllMiniLML6V2);
186    }
187
188    #[test]
189    fn quantized_prefers_q_variant() {
190        let m = resolve_embedding_model("AllMiniLML6V2", true).unwrap();
191        assert_eq!(m, EmbeddingModel::AllMiniLML6V2Q);
192    }
193
194    #[test]
195    fn quantized_skips_already_quantized() {
196        let m = resolve_embedding_model("AllMiniLML6V2Q", true).unwrap();
197        assert_eq!(m, EmbeddingModel::AllMiniLML6V2Q);
198    }
199
200    #[test]
201    fn resolves_bge_alias() {
202        let m = resolve_embedding_model("BAAI/bge-small-en-v1.5", true).unwrap();
203        assert_eq!(m, EmbeddingModel::BGESmallENV15Q);
204    }
205
206    #[test]
207    fn embedding_dims_matches_model() {
208        let m = EmbeddingModel::AllMiniLML6V2;
209        assert_eq!(embedding_dims(&m), 384);
210    }
211}
elph_core/memz/embed.rs

elph_core/memz/
embed.rs
