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 id: String,
75 dim: usize,
76}
77
78#[cfg(feature = "local-embed")]
79impl OnnxEmbedder {
80 pub fn new(cache_dir: &std::path::Path) -> Result<Self> {
81 Self::with_model(cache_dir, "bge-small-en-v1.5")
82 }
83
84 pub fn with_model(cache_dir: &std::path::Path, name: &str) -> Result<Self> {
88 let (model, dim) = match name {
89 "bge-small-en-v1.5" => (fastembed::EmbeddingModel::BGESmallENV15, 384),
90 "bge-base-en-v1.5" => (fastembed::EmbeddingModel::BGEBaseENV15, 768),
91 "nomic-embed-text-v1.5" => (fastembed::EmbeddingModel::NomicEmbedTextV15, 768),
92 other => {
93 return Err(crate::SconeError::InvalidInput(format!(
94 "unknown embed model {other:?}"
95 )));
96 }
97 };
98 let options = fastembed::InitOptions::new(model)
99 .with_cache_dir(cache_dir.to_path_buf())
100 .with_show_download_progress(false);
101 let model = fastembed::TextEmbedding::try_new(options)
102 .map_err(|e| crate::SconeError::Embed(e.to_string()))?;
103 Ok(Self {
104 model: std::cell::RefCell::new(model),
105 id: name.to_owned(),
106 dim,
107 })
108 }
109}
110
111#[cfg(feature = "local-embed")]
112impl EmbeddingProvider for OnnxEmbedder {
113 fn id(&self) -> &str {
114 &self.id
115 }
116
117 fn dim(&self) -> usize {
118 self.dim
119 }
120
121 fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
122 self.model
123 .borrow_mut()
124 .embed(texts, None)
125 .map_err(|e| crate::SconeError::Embed(e.to_string()))
126 }
127}