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

elph_core/memz/
embed.rs
