1use crate::error::Result;
8
9pub trait EmbeddingProvider: Send {
10 fn id(&self) -> &str;
13 fn dim(&self) -> usize;
14 fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
15}
16
17pub struct HashEmbedder {
22 dim: usize,
23}
24
25impl HashEmbedder {
26 pub fn new(dim: usize) -> Self {
27 Self { dim: dim.max(1) }
28 }
29}
30
31impl EmbeddingProvider for HashEmbedder {
32 fn id(&self) -> &str {
33 "hash-v1"
34 }
35
36 fn dim(&self) -> usize {
37 self.dim
38 }
39
40 fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
41 Ok(texts
42 .iter()
43 .map(|text| {
44 let mut v = vec![0.0f32; self.dim];
45 for token in text.to_lowercase().split_whitespace() {
46 let mut h: u64 = 0xcbf2_9ce4_8422_2325;
48 for b in token.as_bytes() {
49 h ^= u64::from(*b);
50 h = h.wrapping_mul(0x0000_0100_0000_01b3);
51 }
52 let bucket = (h % self.dim as u64) as usize;
53 let sign = if h & (1 << 63) == 0 { 1.0 } else { -1.0 };
55 v[bucket] += sign;
56 }
57 let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
58 if norm > 0.0 {
59 for x in &mut v {
60 *x /= norm;
61 }
62 }
63 v
64 })
65 .collect())
66 }
67}
68
69#[cfg(feature = "local-embed")]
72pub struct OnnxEmbedder {
73 model: std::cell::RefCell<fastembed::TextEmbedding>,
74}
75
76#[cfg(feature = "local-embed")]
77impl OnnxEmbedder {
78 pub fn new(cache_dir: &std::path::Path) -> Result<Self> {
79 let options = fastembed::InitOptions::new(fastembed::EmbeddingModel::BGESmallENV15)
80 .with_cache_dir(cache_dir.to_path_buf())
81 .with_show_download_progress(false);
82 let model = fastembed::TextEmbedding::try_new(options)
83 .map_err(|e| crate::SconeError::Embed(e.to_string()))?;
84 Ok(Self {
85 model: std::cell::RefCell::new(model),
86 })
87 }
88}
89
90#[cfg(feature = "local-embed")]
91impl EmbeddingProvider for OnnxEmbedder {
92 fn id(&self) -> &str {
93 "bge-small-en-v1.5"
94 }
95
96 fn dim(&self) -> usize {
97 384
98 }
99
100 fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
101 self.model
102 .borrow_mut()
103 .embed(texts, None)
104 .map_err(|e| crate::SconeError::Embed(e.to_string()))
105 }
106}