clawft_kernel/
embedding_e5.rs

1//! e5-small-v2 ONNX embedding provider (WEFT-640).
2//!
3//! Thin asymmetric wrapper over the BERT-family
4//! [`OnnxEmbeddingProvider`](crate::embedding_onnx::OnnxEmbeddingProvider):
5//! same WordPiece tokenizer, masked-mean pool, and L2 norm — the *only*
6//! behavioural difference is the e5 training convention of prepending
7//! `passage: ` on documents and `query: ` on queries.
8//!
9//! - **Dim**: 384 (matches RVF/HNSW default and `weave.toml` target).
10//! - **License**: MIT weights (`intfloat/e5-small-v2`); community ONNX exports.
11//! - **Staging**: `.weftos/models/e5-small-v2/model.onnx` (+ sibling `vocab.txt`).
12//! - **CI**: without the `onnx-embeddings` feature or staged artifacts the
13//!   provider degrades to the ONNX hash fallback (deterministic, no download).
14//!
15//! Selection: [`select_embedding_provider`](crate::embedding::select_embedding_provider)
16//! promotes e5 when its ONNX file is present and the runtime loads, **below**
17//! project-canonical Qwen3 (ADR-059) and **above** legacy MiniLM.
18
19use std::path::{Path, PathBuf};
20
21use async_trait::async_trait;
22
23use crate::embedding::{EmbeddingError, EmbeddingProvider};
24use crate::embedding_onnx::OnnxEmbeddingProvider;
25
26/// Canonical model identifier (embedder_id / consistency contract).
27pub const E5_MODEL_NAME: &str = "e5-small-v2";
28
29/// Native output dimensionality (matches RVF / HNSW default).
30pub const E5_DIMS: usize = 384;
31
32/// Max tokens for e5-small-v2 (BERT 512 context; leave room for specials).
33pub const E5_MAX_TOKENS: usize = 512;
34
35/// Document / passage prefix required by e5 training.
36pub const PASSAGE_PREFIX: &str = "passage: ";
37
38/// Query prefix required by e5 training.
39pub const QUERY_PREFIX: &str = "query: ";
40
41/// e5-small-v2 provider — 384-d BERT encoder with asymmetric query/passage prefixes.
42pub struct E5EmbeddingProvider {
43    inner: OnnxEmbeddingProvider,
44    /// Stable name exposed via [`EmbeddingProvider::model_name`].
45    model_name: String,
46}
47
48impl E5EmbeddingProvider {
49    /// Create a provider pointing at an e5 ONNX model file.
50    ///
51    /// Degrades to the ONNX hash fallback when the feature is off, the file is
52    /// missing, or session load fails — never panics, never requires the model
53    /// to construct (CI-safe).
54    pub fn new(model_path: impl Into<PathBuf>) -> Self {
55        let model_path = model_path.into();
56        let inner =
57            OnnxEmbeddingProvider::with_config(model_path, E5_DIMS, E5_MAX_TOKENS);
58        let model_name = if inner.is_runtime_available() {
59            E5_MODEL_NAME.to_string()
60        } else {
61            format!("{E5_MODEL_NAME}-hash-fallback")
62        };
63        Self { inner, model_name }
64    }
65
66    /// Whether real ONNX inference is active (vs. hash fallback).
67    pub fn is_runtime_available(&self) -> bool {
68        self.inner.is_runtime_available()
69    }
70
71    /// Path to the configured ONNX model file.
72    pub fn model_path(&self) -> &Path {
73        self.inner.model_path()
74    }
75
76    /// Prefix a string for the **passage** (document) lane.
77    pub fn prefix_passage(text: &str) -> String {
78        if text.starts_with(PASSAGE_PREFIX) || text.starts_with(QUERY_PREFIX) {
79            text.to_string()
80        } else {
81            format!("{PASSAGE_PREFIX}{text}")
82        }
83    }
84
85    /// Prefix a string for the **query** lane.
86    pub fn prefix_query(text: &str) -> String {
87        if text.starts_with(QUERY_PREFIX) || text.starts_with(PASSAGE_PREFIX) {
88            text.to_string()
89        } else {
90            format!("{QUERY_PREFIX}{text}")
91        }
92    }
93}
94
95#[async_trait]
96impl EmbeddingProvider for E5EmbeddingProvider {
97    async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
98        // Document / passage lane.
99        self.inner.embed(&Self::prefix_passage(text)).await
100    }
101
102    async fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
103        let mut results = Vec::with_capacity(texts.len());
104        for text in texts {
105            results.push(self.embed(text).await?);
106        }
107        Ok(results)
108    }
109
110    async fn embed_query(&self, query: &str) -> Result<Vec<f32>, EmbeddingError> {
111        // Asymmetric query lane — required for e5 retrieval quality.
112        self.inner.embed(&Self::prefix_query(query)).await
113    }
114
115    fn dimensions(&self) -> usize {
116        self.inner.dimensions()
117    }
118
119    fn model_name(&self) -> &str {
120        &self.model_name
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn prefix_passage_is_idempotent() {
130        let once = E5EmbeddingProvider::prefix_passage("hello");
131        assert_eq!(once, "passage: hello");
132        assert_eq!(E5EmbeddingProvider::prefix_passage(&once), once);
133    }
134
135    #[test]
136    fn prefix_query_is_idempotent() {
137        let once = E5EmbeddingProvider::prefix_query("find deploys");
138        assert_eq!(once, "query: find deploys");
139        assert_eq!(E5EmbeddingProvider::prefix_query(&once), once);
140    }
141
142    #[test]
143    fn degraded_provider_reports_hash_fallback_name() {
144        // No real model file — hash fallback, CI-safe.
145        let p = E5EmbeddingProvider::new("/nonexistent/e5-small-v2/model.onnx");
146        assert!(!p.is_runtime_available());
147        assert_eq!(p.model_name(), "e5-small-v2-hash-fallback");
148        assert_eq!(p.dimensions(), E5_DIMS);
149    }
150
151    #[tokio::test]
152    async fn embed_and_embed_query_differ_under_hash_fallback() {
153        // Prefixes change the hash input → different vectors. Proves the
154        // asymmetric lane is wired even without ONNX weights.
155        let p = E5EmbeddingProvider::new("/nonexistent/e5-small-v2/model.onnx");
156        let doc = p.embed("the deploy OOM'd on model load").await.unwrap();
157        let qry = p.embed_query("the deploy OOM'd on model load").await.unwrap();
158        assert_eq!(doc.len(), E5_DIMS);
159        assert_eq!(qry.len(), E5_DIMS);
160        assert_ne!(
161            doc, qry,
162            "query: vs passage: prefixes must change the hash-fallback vector"
163        );
164    }
165
166    #[tokio::test]
167    async fn embed_is_deterministic() {
168        let p = E5EmbeddingProvider::new("/nonexistent/e5-small-v2/model.onnx");
169        let a = p.embed("stable input").await.unwrap();
170        let b = p.embed("stable input").await.unwrap();
171        assert_eq!(a, b);
172    }
173}
clawft_kernel/embedding_e5.rs

clawft_kernel/
embedding_e5.rs
