Skip to main content

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}