skm_embed/
bge.rs

1//! BGE-M3 embedding provider via fastembed.
2//!
3//! BGE-M3: 568M params, 100+ languages, 1024-dim.
4//! Chinese + English first-class support.
5
6use std::num::NonZeroUsize;
7use std::sync::{Arc, Mutex};
8
9use async_trait::async_trait;
10use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
11use lru::LruCache;
12use tokio::task::spawn_blocking;
13
14use crate::embedding::Embedding;
15use crate::error::EmbedError;
16use crate::provider::EmbeddingProvider;
17
18/// BGE-M3 embedding provider.
19///
20/// Uses fastembed for ONNX inference.
21/// 1024-dimensional embeddings, supports 100+ languages.
22pub struct BgeM3Provider {
23    /// The fastembed model (Arc for thread safety).
24    model: Arc<Mutex<TextEmbedding>>,
25
26    /// LRU cache for embeddings.
27    cache: Mutex<LruCache<u64, Vec<f32>>>,
28}
29
30impl BgeM3Provider {
31    /// Create a new BGE-M3 provider with default settings.
32    pub fn new() -> Result<Self, EmbedError> {
33        Self::with_cache_size(1000)
34    }
35
36    /// Create a new BGE-M3 provider with custom cache size.
37    pub fn with_cache_size(cache_size: usize) -> Result<Self, EmbedError> {
38        let model = TextEmbedding::try_new(InitOptions::new(EmbeddingModel::BGELargeENV15))
39            .map_err(|e| EmbedError::ModelInit(e.to_string()))?;
40
41        let cache = Mutex::new(LruCache::new(
42            NonZeroUsize::new(cache_size).unwrap_or(NonZeroUsize::new(1).unwrap()),
43        ));
44
45        Ok(Self { 
46            model: Arc::new(Mutex::new(model)), 
47            cache,
48        })
49    }
50
51    /// Check cache for an embedding.
52    fn get_cached(&self, text_hash: u64) -> Option<Vec<f32>> {
53        self.cache.lock().ok()?.get(&text_hash).cloned()
54    }
55
56    /// Store embedding in cache.
57    fn set_cached(&self, text_hash: u64, vector: Vec<f32>) {
58        if let Ok(mut cache) = self.cache.lock() {
59            cache.put(text_hash, vector);
60        }
61    }
62}
63
64#[async_trait]
65impl EmbeddingProvider for BgeM3Provider {
66    async fn embed(&self, texts: &[&str]) -> Result<Vec<Embedding>, EmbedError> {
67        if texts.is_empty() {
68            return Ok(Vec::new());
69        }
70
71        // Check for empty texts
72        for text in texts {
73            if text.is_empty() {
74                return Err(EmbedError::EmptyInput);
75            }
76        }
77
78        // Compute hashes and check cache
79        let hashes: Vec<u64> = texts
80            .iter()
81            .map(|t| xxhash_rust::xxh64::xxh64(t.as_bytes(), 0))
82            .collect();
83
84        let mut results = vec![None; texts.len()];
85        let mut to_embed: Vec<(usize, String)> = Vec::new();
86
87        for (i, (text, hash)) in texts.iter().zip(hashes.iter()).enumerate() {
88            if let Some(cached) = self.get_cached(*hash) {
89                results[i] = Some(Embedding::from_normalized(cached, *hash));
90            } else {
91                to_embed.push((i, text.to_string()));
92            }
93        }
94
95        // Embed uncached texts
96        if !to_embed.is_empty() {
97            let texts_to_embed: Vec<String> = to_embed.iter().map(|(_, t)| t.clone()).collect();
98
99            // Run embedding in blocking thread
100            let model = Arc::clone(&self.model);
101            let embeddings = spawn_blocking(move || {
102                let guard = model.lock().map_err(|e| EmbedError::Embedding(e.to_string()))?;
103                guard.embed(texts_to_embed, None)
104                    .map_err(|e| EmbedError::Embedding(e.to_string()))
105            })
106            .await
107            .map_err(|e| EmbedError::Embedding(e.to_string()))??;
108
109            for ((idx, _), vector) in to_embed.iter().zip(embeddings.into_iter()) {
110                let hash = hashes[*idx];
111                self.set_cached(hash, vector.clone());
112                results[*idx] = Some(Embedding::new(vector, hash));
113            }
114        }
115
116        Ok(results.into_iter().map(|r| r.unwrap()).collect())
117    }
118
119    fn dimensions(&self) -> usize {
120        1024 // BGE-Large uses 1024 dimensions
121    }
122
123    fn model_id(&self) -> &str {
124        "bge-large-en-v1.5"
125    }
126
127    fn max_batch_size(&self) -> usize {
128        32
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    // Tests require model download, skip in CI
135    // Run manually with: cargo test --features embed-bge-m3 -- --ignored
136
137    use super::*;
138
139    #[tokio::test]
140    #[ignore = "requires model download"]
141    async fn test_bge_m3_embed() {
142        let provider = BgeM3Provider::new().unwrap();
143        let embedding = provider.embed_one("Hello, world!").await.unwrap();
144
145        assert_eq!(embedding.dimensions(), 1024);
146        assert!(embedding.is_normalized());
147    }
148
149    #[tokio::test]
150    #[ignore = "requires model download"]
151    async fn test_bge_m3_batch() {
152        let provider = BgeM3Provider::new().unwrap();
153        let texts = vec!["Hello", "World", "Test"];
154        let embeddings = provider.embed(&texts).await.unwrap();
155
156        assert_eq!(embeddings.len(), 3);
157    }
158
159    #[tokio::test]
160    #[ignore = "requires model download"]
161    async fn test_bge_m3_cache() {
162        let provider = BgeM3Provider::new().unwrap();
163
164        let e1 = provider.embed_one("test text").await.unwrap();
165        let e2 = provider.embed_one("test text").await.unwrap();
166
167        // Same text should have same hash
168        assert_eq!(e1.text_hash, e2.text_hash);
169    }
170
171    #[tokio::test]
172    #[ignore = "requires model download"]
173    async fn test_bge_m3_chinese() {
174        let provider = BgeM3Provider::new().unwrap();
175        let embedding = provider.embed_one("你好世界").await.unwrap();
176
177        assert_eq!(embedding.dimensions(), 1024);
178    }
179}
skm_embed/bge.rs

skm_embed/
bge.rs
