1use std::path::Path;
2use std::sync::{Mutex, MutexGuard};
3
4use fastembed::{
5 InitOptionsUserDefined, OutputKey, Pooling, QuantizationMode, TextEmbedding, TokenizerFiles,
6 UserDefinedEmbeddingModel,
7};
8
9use crate::device::{self, EmbedDevice};
10use crate::error::EmbedError;
11
12const MAX_LENGTH: usize = 128;
13const QUERY_PREFIX: &str = "task: search result | query: ";
14const DOC_PREFIX: &str = "title: none | text: ";
15const ONNX_FILE: &str = "model_q4.onnx";
16const ONNX_DATA_FILE: &str = "model_q4.onnx_data";
17const DETAIL_CAP_BYTES: usize = 4096;
18static WEBGPU_SESSION: Mutex<()> = Mutex::new(());
19
20pub struct Embedder {
21 inner: Option<TextEmbedding>,
22 device: EmbedDevice,
23 fall_off_reason: Option<String>,
24}
25
26impl Embedder {
27 pub fn load(model_dir: impl AsRef<Path>, device: EmbedDevice) -> Result<Self, EmbedError> {
28 let model_dir = model_dir.as_ref();
29 match Self::load_on(model_dir, device) {
30 Err(EmbedError::GpuUnavailable { detail }) if device != EmbedDevice::Cpu => {
31 let mut fallen = Self::load_on(model_dir, EmbedDevice::Cpu)?;
32 fallen.fall_off_reason = Some(detail);
33 Ok(fallen)
34 }
35 outcome => outcome,
36 }
37 }
38
39 pub fn device_in_use(&self) -> EmbedDevice {
40 self.device
41 }
42
43 pub fn fall_off_reason(&self) -> Option<&str> {
44 self.fall_off_reason.as_deref()
45 }
46
47 fn load_on(model_dir: &Path, device: EmbedDevice) -> Result<Self, EmbedError> {
48 let _webgpu = webgpu_session(device);
49 let onnx_file = read_required(&model_dir.join("onnx").join(ONNX_FILE))?;
50 let onnx_data = read_required(&model_dir.join("onnx").join(ONNX_DATA_FILE))?;
51 let tokenizer_file = read_required(&model_dir.join("tokenizer.json"))?;
52 let config_file = read_required(&model_dir.join("config.json"))?;
53 let special_tokens_map_file = read_required(&model_dir.join("special_tokens_map.json"))?;
54 let tokenizer_config_file = read_required(&model_dir.join("tokenizer_config.json"))?;
55
56 ensure_provider_available(device)?;
57
58 let mut model = UserDefinedEmbeddingModel::new(
59 onnx_file,
60 TokenizerFiles {
61 tokenizer_file,
62 config_file,
63 special_tokens_map_file,
64 tokenizer_config_file,
65 },
66 )
67 .with_pooling(Pooling::Mean)
68 .with_quantization(QuantizationMode::None)
69 .with_external_initializer(ONNX_DATA_FILE.to_string(), onnx_data);
70 model.output_key = Some(OutputKey::ByName("sentence_embedding"));
71
72 let execution_providers = providers(device)?;
73 let options = InitOptionsUserDefined::new()
74 .with_execution_providers(execution_providers)
75 .with_max_length(MAX_LENGTH);
76
77 let inner = TextEmbedding::try_new_from_user_defined(model, options).map_err(|e| {
78 EmbedError::LoadFailed {
79 detail: tail_capped(e.to_string()),
80 }
81 })?;
82
83 Ok(Self {
84 inner: Some(inner),
85 device,
86 fall_off_reason: None,
87 })
88 }
89
90 pub fn embed_query(&mut self, text: &str) -> Result<Vec<f32>, EmbedError> {
91 let _webgpu = webgpu_session(self.device);
92 let mut out = self
93 .inner
94 .as_mut()
95 .expect("embedder session exists until drop")
96 .embed(vec![prefixed_query(text)], None)
97 .map_err(inference_err)?;
98 let vec = out.pop().ok_or_else(|| EmbedError::Inference {
99 detail: "embed_query produced no output".to_string(),
100 })?;
101 Ok(normalize(&vec))
102 }
103
104 pub fn embed_documents(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbedError> {
105 if texts.is_empty() {
106 return Ok(Vec::new());
107 }
108 let _webgpu = webgpu_session(self.device);
109 let prefixed: Vec<String> = texts.iter().map(|t| prefixed_document(t)).collect();
110 let out = self
111 .inner
112 .as_mut()
113 .expect("embedder session exists until drop")
114 .embed(prefixed, None)
115 .map_err(inference_err)?;
116 Ok(out.iter().map(|v| normalize(v)).collect())
117 }
118}
119
120impl Drop for Embedder {
121 fn drop(&mut self) {
122 let _webgpu = webgpu_session(self.device);
123 drop(self.inner.take());
124 }
125}
126
127fn webgpu_session(device: EmbedDevice) -> Option<MutexGuard<'static, ()>> {
128 (device == EmbedDevice::WebGpu).then(|| {
129 WEBGPU_SESSION
130 .lock()
131 .unwrap_or_else(|poisoned| poisoned.into_inner())
132 })
133}
134
135fn prefixed_query(text: &str) -> String {
136 format!("{QUERY_PREFIX}{text}")
137}
138
139fn prefixed_document(text: &str) -> String {
140 format!("{DOC_PREFIX}{text}")
141}
142
143fn read_required(path: &Path) -> Result<Vec<u8>, EmbedError> {
144 std::fs::read(path).map_err(|_| EmbedError::ModelMissing {
145 path: path.to_path_buf(),
146 })
147}
148
149fn providers(device: EmbedDevice) -> Result<Vec<ort::ep::ExecutionProviderDispatch>, EmbedError> {
150 let compiled = EmbedDevice::compiled();
151 #[cfg(debug_assertions)]
152 if device == EmbedDevice::Broken {
153 return Err(EmbedError::GpuUnavailable {
154 detail: "the test door: this card's provider always fails to build".to_string(),
155 });
156 }
157 if device != EmbedDevice::Cpu && device != compiled {
158 return Err(EmbedError::GpuUnavailable {
159 detail: format!(
160 "{} requested but this build carries {}",
161 device.label(),
162 compiled.label()
163 ),
164 });
165 }
166 Ok(vec![device::provider(device)])
167}
168
169fn ensure_provider_available(device: EmbedDevice) -> Result<(), EmbedError> {
170 let builder = ort::session::Session::builder().map_err(|e| EmbedError::GpuUnavailable {
171 detail: e.to_string(),
172 })?;
173 builder
174 .with_execution_providers(providers(device)?)
175 .map_err(|e| EmbedError::GpuUnavailable {
176 detail: e.to_string(),
177 })?;
178 Ok(())
179}
180
181fn tail_capped(detail: String) -> String {
182 if detail.len() <= DETAIL_CAP_BYTES {
183 return detail;
184 }
185 let mut start = detail.len() - DETAIL_CAP_BYTES;
186 while start < detail.len() && !detail.is_char_boundary(start) {
187 start += 1;
188 }
189 detail[start..].to_string()
190}
191
192fn inference_err(e: fastembed::Error) -> EmbedError {
193 EmbedError::Inference {
194 detail: e.to_string(),
195 }
196}
197
198fn normalize(v: &[f32]) -> Vec<f32> {
199 let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
200 if n > 0.0 {
201 v.iter().map(|x| x / n).collect()
202 } else {
203 v.to_vec()
204 }
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210
211 #[test]
212 fn query_prefix_is_baked_in_verbatim() {
213 assert_eq!(
214 prefixed_query("boss fight"),
215 "task: search result | query: boss fight"
216 );
217 }
218
219 #[test]
220 fn document_prefix_is_baked_in_verbatim() {
221 assert_eq!(
222 prefixed_document("boss fight"),
223 "title: none | text: boss fight"
224 );
225 }
226
227 #[test]
228 fn query_and_document_prefixes_diverge() {
229 assert_ne!(prefixed_query("same text"), prefixed_document("same text"));
230 }
231
232 #[test]
233 fn normalize_produces_unit_vectors() {
234 let v = vec![3.0_f32, 4.0];
235 let n = normalize(&v);
236 let mag = n.iter().map(|x| x * x).sum::<f32>().sqrt();
237 assert!((mag - 1.0).abs() < 1e-6);
238 assert_eq!(n, vec![0.6, 0.8]);
239 }
240
241 #[test]
242 fn normalize_zero_vector_is_safe() {
243 let v = vec![0.0_f32; 768];
244 let n = normalize(&v);
245 assert_eq!(n, v);
246 assert!(n.iter().all(|x| x.is_finite()));
247 }
248
249 #[test]
250 fn tail_capped_leaves_short_strings_untouched() {
251 let short = "backend refused: bad graph".to_string();
252 assert_eq!(tail_capped(short.clone()), short);
253 }
254
255 #[test]
256 fn tail_capped_bounds_long_strings_to_the_cap() {
257 let long = "x".repeat(DETAIL_CAP_BYTES * 3);
258 let capped = tail_capped(long);
259 assert_eq!(capped.len(), DETAIL_CAP_BYTES);
260 }
261
262 #[test]
263 fn tail_capped_keeps_the_tail_not_the_head() {
264 let long = format!("{}TAIL-MARKER", "head".repeat(DETAIL_CAP_BYTES * 2));
265 let capped = tail_capped(long.clone());
266 assert!(capped.ends_with("TAIL-MARKER"));
267 assert_eq!(capped, long[long.len() - DETAIL_CAP_BYTES..]);
268 }
269
270 #[test]
271 fn tail_capped_respects_utf8_boundaries() {
272 let mut long = "a".repeat(DETAIL_CAP_BYTES - 1);
273 long.push('δΈ');
274 long.push_str(&"b".repeat(10));
275 let capped = tail_capped(long);
276 assert!(std::str::from_utf8(capped.as_bytes()).is_ok());
277 assert!(capped.len() <= DETAIL_CAP_BYTES);
278 }
279}