onnx_embedding_plugin/lib.rs
1//! In-process ONNX embedding provider for the graph-storage gear.
2//!
3//! ADR-0005 makes this the default: a `MiniLM`-class sentence-embedding model
4//! run through ONNX Runtime, in the gear's own process, so a small deployment
5//! needs no inference service to use vector search at all.
6//!
7//! # Artifacts are supplied, never fetched
8//!
9//! The model and tokenizer are read from paths the operator configures. The
10//! crate downloads nothing. That is not caution about the network: the
11//! embedding-space identity has to be *verifiable*, and the only identity a
12//! downloader can offer is the name it asked for. Reading a file lets the
13//! identity be the SHA-256 of the bytes actually loaded, so two deployments
14//! claiming one space either agree on that hash or are visibly different.
15//!
16//! # The runtime is loaded, not linked
17//!
18//! `ort` is pinned with `load-dynamic`, so ONNX Runtime is resolved by
19//! `dlopen` at first use through `ORT_DYLIB_PATH`. Building this crate needs
20//! no runtime headers; running it needs the shared library. See
21//! [`OnnxEmbeddingProvider::load`] for what happens when that path is wrong,
22//! which is worse than an error.
23
24use std::path::{Path, PathBuf};
25use std::sync::Arc;
26use std::time::Duration;
27
28use async_trait::async_trait;
29use aws_lc_rs::digest::{SHA256, digest as sha256};
30use graph_storage_sdk::models::EmbeddingSpaceId;
31use graph_storage_sdk::plugin_api::{
32 EmbedRequest, EmbedResponse, EmbeddingProviderError, EmbeddingProviderV1,
33};
34use ort::session::Session;
35use ort::session::builder::GraphOptimizationLevel;
36use ort::value::Tensor;
37use thiserror::Error;
38use tokenizers::Tokenizer;
39use tokio::sync::Mutex;
40use tracing::warn;
41
42/// How the model turns a sequence of token vectors into one sentence vector.
43///
44/// Part of the embedding-space identity rather than a tuning knob: the same
45/// weights pooled two ways produce vectors of the same width that must never
46/// be compared.
47#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
48pub enum Pooling {
49 /// Average over the tokens the attention mask keeps. What the
50 /// `sentence-transformers` `MiniLM` models are trained with.
51 #[default]
52 Mean,
53 /// The first token's vector.
54 Cls,
55}
56
57impl Pooling {
58 fn as_str(self) -> &'static str {
59 match self {
60 Self::Mean => "mean",
61 Self::Cls => "cls",
62 }
63 }
64}
65
66/// What a deployment declares about its model.
67#[derive(Clone, Debug)]
68pub struct OnnxProviderConfig {
69 pub model_path: PathBuf,
70 pub tokenizer_path: PathBuf,
71 /// Vector width the model emits. Checked against the model's own output
72 /// on the first batch, not taken on trust.
73 pub dimension: u32,
74 pub pooling: Pooling,
75 /// L2-normalize the pooled vector. On for cosine similarity, which is
76 /// what the gear's index serves.
77 pub normalize: bool,
78 /// Longest token sequence handed to the model; longer inputs are
79 /// truncated.
80 pub max_tokens: usize,
81 /// ONNX Runtime intra-op threads. `None` leaves the runtime's default.
82 pub intra_op_threads: Option<usize>,
83}
84
85impl OnnxProviderConfig {
86 /// The `MiniLM-L6-v2` defaults ADR-0005 describes: 384 dimensions, mean
87 /// pooling, L2-normalized.
88 #[must_use]
89 pub fn new(model_path: impl Into<PathBuf>, tokenizer_path: impl Into<PathBuf>) -> Self {
90 Self {
91 model_path: model_path.into(),
92 tokenizer_path: tokenizer_path.into(),
93 dimension: 384,
94 pooling: Pooling::Mean,
95 normalize: true,
96 max_tokens: 256,
97 intra_op_threads: None,
98 }
99 }
100}
101
102#[derive(Debug, Error)]
103#[non_exhaustive]
104pub enum OnnxLoadError {
105 #[error("cannot read {what} at {path}: {source}")]
106 Artifact {
107 what: &'static str,
108 path: PathBuf,
109 source: std::io::Error,
110 },
111 #[error("tokenizer at {path} is not loadable: {reason}")]
112 Tokenizer { path: PathBuf, reason: String },
113 #[error("ONNX session could not be created: {0}")]
114 Session(String),
115 /// The one failure a caller cannot recover from in-process.
116 #[error(
117 "ONNX Runtime did not load within {seconds}s. `ort` 2.0.0-rc.12 hangs \
118 instead of erroring on an unloadable library, so the thread that \
119 tried is abandoned rather than killed: check ORT_DYLIB_PATH and \
120 restart the process"
121 )]
122 RuntimeHung { seconds: u64 },
123}
124
125/// How long to wait for the runtime before deciding it has hung.
126const LOAD_TIMEOUT: Duration = Duration::from_secs(30);
127
128/// How long the boot probe's single inference may take.
129///
130/// Shorter than `LOAD_TIMEOUT` because it measures a different thing: loading
131/// a model reads and plans a graph, running one short input through it should
132/// be milliseconds. The bound exists for the same reason the load one does --
133/// `ort` can hang rather than error -- and a probe that could hang would turn
134/// a check meant to catch a broken session into a gear that never starts.
135const PROBE_TIMEOUT: Duration = Duration::from_secs(10);
136
137/// A `MiniLM`-class sentence-embedding model, in this process.
138pub struct OnnxEmbeddingProvider {
139 /// `ort`'s `Session::run` takes `&mut self`, so inference is serialized
140 /// whatever the sharing. A fair mutex over one session is then the honest
141 /// shape: extra sessions would each hold a resident copy of the weights
142 /// and their own intra-op thread pool, which is the wrong trade for a
143 /// component called once per ingest batch.
144 session: Arc<Mutex<Session>>,
145 tokenizer: Tokenizer,
146 space: EmbeddingSpaceId,
147 config: OnnxProviderConfig,
148 /// What the last real exchange with the session showed, if it failed.
149 ///
150 /// Readiness reports this rather than running inference of its own. A
151 /// probe that embedded something would make an anonymous, scheduled
152 /// endpoint the busiest caller of the model -- the same amplification the
153 /// remote provider is guarded against -- and a probe that only takes the
154 /// lock, which is what this used to do, reports that the weights are
155 /// resident and nothing else.
156 observed: std::sync::Mutex<Option<String>>,
157}
158
159impl OnnxEmbeddingProvider {
160 /// Load the artifacts and open a session.
161 ///
162 /// # The hang this guards against
163 ///
164 /// `ort` 2.0.0-rc.12 **hangs forever instead of erroring** when the
165 /// library at `ORT_DYLIB_PATH` cannot be loaded — measured in this
166 /// repository against a nonexistent path, where neither `Session::builder`
167 /// nor `ort::init_from` returns in 45 seconds. No pre-flight validation
168 /// exists; every entry point funnels through the same lazy init. Since the
169 /// hang cannot be interrupted from inside, the blocked thread is
170 /// **abandoned**: a raw `std::thread` rather than `spawn_blocking`,
171 /// because Tokio joins blocking threads at shutdown and a wedged one would
172 /// hang that too. A caller receiving [`OnnxLoadError::RuntimeHung`] has
173 /// leaked one thread and must terminate the process rather than retry.
174 ///
175 /// See `gears/file-parser/file-parser/src/gear.rs` for the same mitigation
176 /// and the measurements behind it.
177 ///
178 /// # Errors
179 ///
180 /// Unreadable artifacts, an unloadable tokenizer, a session that refuses
181 /// to open, or the runtime hang above.
182 pub async fn load(config: OnnxProviderConfig) -> Result<Self, OnnxLoadError> {
183 let model_digest = file_digest("model", &config.model_path)?;
184 let tokenizer_digest = file_digest("tokenizer", &config.tokenizer_path)?;
185
186 let tokenizer = Tokenizer::from_file(&config.tokenizer_path).map_err(|error| {
187 OnnxLoadError::Tokenizer {
188 path: config.tokenizer_path.clone(),
189 reason: error.to_string(),
190 }
191 })?;
192
193 let session = open_session(&config).await?;
194
195 // The identity is the bytes actually loaded plus how they are used.
196 // A deployment that swaps the file under one configured name gets a
197 // different identity and is caught at boot rather than in ranking.
198 let space = EmbeddingSpaceId::new(
199 artifact_name(&config.model_path, &model_digest),
200 artifact_name(&config.tokenizer_path, &tokenizer_digest),
201 serde_json::json!({
202 "tokenizer": "file",
203 "max_tokens": config.max_tokens,
204 "truncation": "longest_first",
205 }),
206 serde_json::json!({ "strategy": config.pooling.as_str() }),
207 serde_json::json!({ "l2": config.normalize }),
208 config.dimension,
209 );
210
211 let provider = Arc::new(Self {
212 session: Arc::new(Mutex::new(session)),
213 tokenizer,
214 space,
215 config,
216 observed: std::sync::Mutex::new(None),
217 });
218 // One inference before the provider is handed out. A session can load
219 // and still be unable to run -- a runtime built without the execution
220 // provider the graph needs is the usual way -- and every check up to
221 // here would pass: the file hashes, the tokenizer parses, the session
222 // opens. Without this the first evidence arrives when a producer's
223 // ingest fails, long after readiness said the deployment was fine.
224 Self::probe(Arc::clone(&provider)).await?;
225 // The probe's task has ended, so this is the only reference left.
226 Ok(Arc::try_unwrap(provider).unwrap_or_else(|_| {
227 unreachable!("the probe task is the only other holder and it has finished")
228 }))
229 }
230}
231
232fn artifact_name(path: &Path, digest: &str) -> String {
233 let name = path.file_name().map_or_else(
234 || "unnamed".to_owned(),
235 |n| n.to_string_lossy().into_owned(),
236 );
237 format!("{name}@sha256:{digest}")
238}
239
240fn file_digest(what: &'static str, path: &Path) -> Result<String, OnnxLoadError> {
241 let bytes = std::fs::read(path).map_err(|source| OnnxLoadError::Artifact {
242 what,
243 path: path.to_path_buf(),
244 source,
245 })?;
246 Ok(hex::encode(sha256(&SHA256, &bytes)))
247}
248
249/// Open the session on an abandonable thread. See [`OnnxEmbeddingProvider::load`].
250async fn open_session(config: &OnnxProviderConfig) -> Result<Session, OnnxLoadError> {
251 let (tx, rx) = tokio::sync::oneshot::channel();
252 let model_path = config.model_path.clone();
253 let threads = config.intra_op_threads;
254
255 std::thread::Builder::new()
256 .name("graph-storage-onnx-init".to_owned())
257 .spawn(move || {
258 let result = build_session(&model_path, threads);
259 // The receiver is gone when the timeout already fired; that is the
260 // abandoned case, and dropping the session here is correct.
261 drop(tx.send(result));
262 })
263 .map_err(|error| OnnxLoadError::Session(error.to_string()))?;
264
265 match tokio::time::timeout(LOAD_TIMEOUT, rx).await {
266 Ok(Ok(result)) => result,
267 // The sender was dropped without sending: the init thread panicked.
268 Ok(Err(_)) => Err(OnnxLoadError::Session(
269 "the ONNX init thread ended without a result".to_owned(),
270 )),
271 Err(_) => {
272 warn!(
273 path = %config.model_path.display(),
274 "ONNX Runtime did not load in time; leaking the init thread deliberately"
275 );
276 Err(OnnxLoadError::RuntimeHung {
277 seconds: LOAD_TIMEOUT.as_secs(),
278 })
279 }
280 }
281}
282
283fn build_session(
284 model_path: &Path,
285 intra_op_threads: Option<usize>,
286) -> Result<Session, OnnxLoadError> {
287 let mut builder = Session::builder()
288 .map_err(|error| OnnxLoadError::Session(error.to_string()))?
289 .with_optimization_level(GraphOptimizationLevel::Level3)
290 .map_err(|error| OnnxLoadError::Session(error.to_string()))?;
291 if let Some(threads) = intra_op_threads {
292 builder = builder
293 .with_intra_threads(threads)
294 .map_err(|error| OnnxLoadError::Session(error.to_string()))?;
295 }
296 builder
297 .commit_from_file(model_path)
298 .map_err(|error| OnnxLoadError::Session(error.to_string()))
299}
300
301#[async_trait]
302impl EmbeddingProviderV1 for OnnxEmbeddingProvider {
303 fn embedding_space(&self) -> &EmbeddingSpaceId {
304 &self.space
305 }
306
307 fn dimension(&self) -> u32 {
308 self.config.dimension
309 }
310
311 async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, EmbeddingProviderError> {
312 if req.cancel.is_cancelled() {
313 return Err(EmbeddingProviderError::Cancelled);
314 }
315 if req.budget.is_exhausted() {
316 return Err(EmbeddingProviderError::Deadline);
317 }
318 if req.inputs.is_empty() {
319 return Ok(EmbedResponse {
320 vectors: Vec::new(),
321 space: self.space.clone(),
322 });
323 }
324
325 let encoded = self.encode(&req.inputs)?;
326 let mut session = self.session.lock().await;
327 // Checked again after the queue: a batch that waited out its deadline
328 // -- or whose caller gave up -- behind another one should not then
329 // spend CPU on it.
330 if req.budget.is_exhausted() {
331 return Err(EmbeddingProviderError::Deadline);
332 }
333 if req.cancel.is_cancelled() {
334 return Err(EmbeddingProviderError::Cancelled);
335 }
336 // `Session::run` is synchronous CPU work measured in tens to hundreds
337 // of milliseconds. Calling it directly would hold a Tokio worker
338 // thread for that long, starving every other task scheduled on it --
339 // in a gear whose other work is database round trips, that is the
340 // difference between a slow embedding and a stalled request.
341 //
342 // `block_in_place` rather than `spawn_blocking` because the session
343 // guard borrows `self`: moving it into a `'static` task would mean
344 // cloning the `Arc` and re-locking inside, which is the same work
345 // with more moving parts. It needs a multi-threaded runtime, which is
346 // what the gear runs on; a current-thread runtime (some tests) gets
347 // the direct call, which is correct there because there are no other
348 // tasks to starve.
349 let outcome = if tokio::runtime::Handle::try_current().is_ok_and(|handle| {
350 handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread
351 }) {
352 tokio::task::block_in_place(|| self.run(&mut session, &encoded))
353 } else {
354 self.run(&mut session, &encoded)
355 };
356 // Both outcomes are evidence, which is what lets readiness answer
357 // without running inference of its own: a failure here is what a
358 // later probe reports, and a success clears one.
359 self.record(match &outcome {
360 Ok(_) => None,
361 Err(error) => Some(error.to_string()),
362 });
363 let vectors = outcome?;
364
365 Ok(EmbedResponse {
366 vectors,
367 space: self.space.clone(),
368 })
369 }
370
371 async fn health(&self) -> Result<(), EmbeddingProviderError> {
372 // Observed, not polled: `load` proved the session can infer, and
373 // every `embed` since has been evidence of its own. Taking the lock
374 // still matters -- a session whose holder panicked cannot be locked
375 // -- but on its own it only reports that the weights are resident.
376 drop(self.session.lock().await);
377 match self.failure() {
378 Some(reason) => Err(EmbeddingProviderError::Unavailable { reason }),
379 None => Ok(()),
380 }
381 }
382}
383
384impl OnnxEmbeddingProvider {
385 /// Run one inference, so "the session loaded" and "the session works" are
386 /// not the same claim.
387 ///
388 /// The input is a fixed short string: the point is that the graph
389 /// executes end to end and returns a vector of the declared width, not
390 /// what the vector says.
391 async fn probe(provider: Arc<Self>) -> Result<(), OnnxLoadError> {
392 // On a blocking thread and under a bound, for the same two reasons
393 // `open_session` uses them. `Session::run` is synchronous CPU work,
394 // so calling it on a Tokio worker holds that worker for the duration
395 // -- the reason `embed` uses `block_in_place`. And `ort` can hang
396 // rather than error, which is what `LOAD_TIMEOUT` is there for: a
397 // probe added to catch a session that cannot run would otherwise be
398 // able to stop the gear from ever starting, which is worse than the
399 // fault it looks for.
400 let inference = tokio::task::spawn_blocking(move || {
401 let encoded = provider
402 .encode(std::slice::from_ref(&PROBE_INPUT.to_owned()))
403 .map_err(|error| OnnxLoadError::Session(format!("probe tokenization: {error}")))?;
404 // Nobody else holds the session yet; this is load.
405 let mut session = provider.session.blocking_lock();
406 // The width is checked by `run` itself, which answers
407 // `SpaceMismatch` -- a declared width the model does not produce
408 // is exactly what that name is for, and comparing again here
409 // would be a second answer to one question. Reaching it is the
410 // point: this makes the check happen once, before anyone depends
411 // on the provider.
412 provider
413 .run(&mut session, &encoded)
414 .map_err(|error| OnnxLoadError::Session(format!("probe inference: {error}")))?;
415 Ok::<(), OnnxLoadError>(())
416 });
417
418 match tokio::time::timeout(PROBE_TIMEOUT, inference).await {
419 Ok(Ok(result)) => result,
420 Ok(Err(join)) => Err(OnnxLoadError::Session(format!(
421 "the ONNX probe thread ended without a result: {join}"
422 ))),
423 Err(_) => {
424 warn!(
425 seconds = PROBE_TIMEOUT.as_secs(),
426 "the ONNX session loaded but did not answer one inference in time"
427 );
428 Err(OnnxLoadError::RuntimeHung {
429 seconds: PROBE_TIMEOUT.as_secs(),
430 })
431 }
432 }
433 }
434
435 /// The failure the last exchange left behind, if any.
436 fn failure(&self) -> Option<String> {
437 self.observed.lock().map_or(
438 Some("the observation lock is poisoned".to_owned()),
439 |seen| seen.clone(),
440 )
441 }
442
443 /// Record what an exchange showed. `None` clears an earlier failure: a
444 /// session that has just produced a vector is working, whatever it did
445 /// before.
446 fn record(&self, failure: Option<String>) {
447 if let Ok(mut seen) = self.observed.lock() {
448 *seen = failure;
449 }
450 }
451}
452
453/// What the boot probe embeds. Fixed and short: it is asked whether the graph
454/// runs, not what it thinks.
455const PROBE_INPUT: &str = "graph storage readiness probe";
456
457/// One tokenized batch, padded to its own longest sequence.
458struct Encoded {
459 ids: Vec<i64>,
460 mask: Vec<i64>,
461 type_ids: Vec<i64>,
462 rows: usize,
463 columns: usize,
464}
465
466impl OnnxEmbeddingProvider {
467 fn encode(&self, inputs: &[String]) -> Result<Encoded, EmbeddingProviderError> {
468 let encodings = self
469 .tokenizer
470 .encode_batch(inputs.to_vec(), true)
471 .map_err(|error| EmbeddingProviderError::Internal(error.to_string()))?;
472
473 // The tokenizer's own configuration decides the sequence length, and
474 // for the `MiniLM` artifacts this crate targets that is a fixed 128
475 // whatever the text -- measured, not assumed. `max_tokens` is a
476 // ceiling on top of it, so it shortens sequences and never lengthens
477 // them. Either way most positions are padding, which is why the
478 // attention mask below is load-bearing rather than tidy: pooling over
479 // the padding as well makes every text resemble every other one.
480 let columns = encodings
481 .iter()
482 .map(|e| e.get_ids().len().min(self.config.max_tokens))
483 .max()
484 .unwrap_or(1)
485 .max(1);
486 let rows = encodings.len();
487
488 let mut ids = vec![0_i64; rows * columns];
489 let mut mask = vec![0_i64; rows * columns];
490 let type_ids = vec![0_i64; rows * columns];
491 for (row, encoding) in encodings.iter().enumerate() {
492 let take = encoding.get_ids().len().min(columns);
493 for column in 0..take {
494 ids[row * columns + column] = i64::from(encoding.get_ids()[column]);
495 mask[row * columns + column] = i64::from(encoding.get_attention_mask()[column]);
496 }
497 }
498 Ok(Encoded {
499 ids,
500 mask,
501 type_ids,
502 rows,
503 columns,
504 })
505 }
506
507 fn run(
508 &self,
509 session: &mut Session,
510 encoded: &Encoded,
511 ) -> Result<Vec<Vec<f32>>, EmbeddingProviderError> {
512 let internal = |what: String| EmbeddingProviderError::Internal(what);
513 let shape = [encoded.rows, encoded.columns];
514 let tensor = |data: &[i64]| {
515 Tensor::from_array((shape, data.to_vec().into_boxed_slice()))
516 .map_err(|error| internal(error.to_string()))
517 };
518
519 let outputs = session
520 .run(ort::inputs![
521 "input_ids" => tensor(&encoded.ids)?,
522 "attention_mask" => tensor(&encoded.mask)?,
523 "token_type_ids" => tensor(&encoded.type_ids)?,
524 ])
525 .map_err(|error| EmbeddingProviderError::Unavailable {
526 reason: error.to_string(),
527 })?;
528
529 // The first output is the token-level hidden state whatever the export
530 // named it; `MiniLM` exports vary between `last_hidden_state` and
531 // `output_0`, and a provider that insisted on one name would refuse
532 // half the artifacts an operator might reasonably supply.
533 let (_, first) = outputs
534 .iter()
535 .next()
536 .ok_or_else(|| internal("the model produced no output".to_owned()))?;
537 let (out_shape, values) = first
538 .try_extract_tensor::<f32>()
539 .map_err(|error| internal(error.to_string()))?;
540
541 if out_shape.len() != 3 {
542 return Err(internal(format!(
543 "expected a [batch, tokens, hidden] output, got {out_shape:?}"
544 )));
545 }
546 let dim = |axis: usize| {
547 out_shape
548 .get(axis)
549 .copied()
550 .and_then(|value| usize::try_from(value).ok())
551 .ok_or_else(|| internal(format!("the model reported a shape of {out_shape:?}")))
552 };
553 // Every axis, not only the last. `pool` walks the buffer with the
554 // *encoded* batch and token counts, so a model whose output disagrees
555 // with them on either axis would be read at the wrong offsets —
556 // panicking on a short buffer, and silently mixing rows on a long one.
557 // The width is checked below against what the deployment declared.
558 let (batch, tokens, hidden) = (dim(0)?, dim(1)?, dim(2)?);
559 if batch != encoded.rows || tokens != encoded.columns {
560 return Err(internal(format!(
561 "the model answered a [{batch}, {tokens}, _] batch for a [{}, {}, _] input",
562 encoded.rows, encoded.columns
563 )));
564 }
565 // The declared width against the width the model actually emits.
566 // A configuration that names one and loads another would write
567 // vectors nothing can rank, and the column would refuse them anyway.
568 if hidden != self.config.dimension as usize {
569 // The numbers, because `SpaceMismatch` carries none and this is
570 // the only place that knows them. At boot the probe turns this
571 // into a startup failure, and "embedding space mismatch" on its
572 // own leaves an operator to find the model's real width by
573 // reading code -- which is what happened the first time this
574 // fired. Same shape as the remote plugin's.
575 warn!(
576 got = hidden,
577 want = self.config.dimension,
578 model = %self.config.model_path.display(),
579 "the model's output width is not the configured embedding dimension"
580 );
581 return Err(EmbeddingProviderError::SpaceMismatch);
582 }
583
584 Ok(self.pool(values, encoded, hidden))
585 }
586
587 fn pool(&self, values: &[f32], encoded: &Encoded, hidden: usize) -> Vec<Vec<f32>> {
588 let mut out = Vec::with_capacity(encoded.rows);
589 for row in 0..encoded.rows {
590 let base = row * encoded.columns * hidden;
591 let mut pooled = vec![0.0_f64; hidden];
592 match self.config.pooling {
593 Pooling::Cls => {
594 for (lane, slot) in pooled.iter_mut().enumerate() {
595 *slot = f64::from(values[base + lane]);
596 }
597 }
598 Pooling::Mean => {
599 // Masked mean: padding contributes nothing, or a batch's
600 // vectors would depend on the longest text beside them.
601 let mut kept = 0.0_f64;
602 for column in 0..encoded.columns {
603 if encoded.mask[row * encoded.columns + column] == 0 {
604 continue;
605 }
606 kept += 1.0;
607 let offset = base + column * hidden;
608 for (lane, slot) in pooled.iter_mut().enumerate() {
609 *slot += f64::from(values[offset + lane]);
610 }
611 }
612 if kept > 0.0 {
613 for slot in &mut pooled {
614 *slot /= kept;
615 }
616 }
617 }
618 }
619 if self.config.normalize {
620 let norm = pooled.iter().map(|x| x * x).sum::<f64>().sqrt();
621 if norm > 0.0 {
622 for slot in &mut pooled {
623 *slot /= norm;
624 }
625 }
626 }
627 out.push(narrow(&pooled));
628 }
629 out
630 }
631}
632
633/// Narrowing to f32 is the destination, not an accident: pgvector stores
634/// single precision, so a vector that did not round here would round on the
635/// way into the column instead.
636#[expect(
637 clippy::cast_possible_truncation,
638 reason = "see the function's own documentation"
639)]
640fn narrow(values: &[f64]) -> Vec<f32> {
641 values.iter().map(|v| *v as f32).collect()
642}
643
644#[cfg(test)]
645mod tests {
646 use super::*;
647
648 #[test]
649 fn pooling_is_part_of_the_identity() {
650 // Two spaces that differ only in how they pool must not be confused:
651 // same weights, same width, incomparable vectors.
652 let one = EmbeddingSpaceId::new(
653 "m@sha256:a",
654 "t@sha256:b",
655 serde_json::json!({}),
656 serde_json::json!({ "strategy": Pooling::Mean.as_str() }),
657 serde_json::json!({ "l2": true }),
658 384,
659 );
660 let other = EmbeddingSpaceId::new(
661 "m@sha256:a",
662 "t@sha256:b",
663 serde_json::json!({}),
664 serde_json::json!({ "strategy": Pooling::Cls.as_str() }),
665 serde_json::json!({ "l2": true }),
666 384,
667 );
668 assert_ne!(one.identity_hash, other.identity_hash);
669 }
670
671 #[test]
672 fn a_missing_artifact_is_named_in_the_error() {
673 let error = file_digest("model", Path::new("/nonexistent/model.onnx"))
674 .expect_err("a missing file cannot be digested");
675 let rendered = error.to_string();
676 assert!(rendered.contains("model"), "{rendered}");
677 assert!(rendered.contains("/nonexistent/model.onnx"), "{rendered}");
678 }
679
680 #[test]
681 fn the_artifact_name_carries_the_digest_rather_than_the_path() {
682 // The path is a deployment detail; the bytes are the identity. Two
683 // hosts with the same model in different directories must agree.
684 assert_eq!(
685 artifact_name(Path::new("/opt/models/model.onnx"), "abc"),
686 artifact_name(Path::new("/srv/other/model.onnx"), "abc")
687 );
688 }
689}