seyk_embed/
embedding.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{Error, Result};
4
5/// A nonzero embedding width, independent of any model's supported dimensions.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7#[serde(into = "usize", try_from = "usize")]
8pub struct Dimension(usize);
9
10impl Dimension {
11    pub const D128: Self = Self(128);
12    pub const D256: Self = Self(256);
13    pub const D512: Self = Self(512);
14    pub const D768: Self = Self(768);
15    /// Common Matryoshka widths. Individual backends validate their own support.
16    pub const ALL: [Self; 4] = [Self::D128, Self::D256, Self::D512, Self::D768];
17    pub fn get(self) -> usize {
18        self.0
19    }
20}
21
22/// The historical default width; backends may choose a different default.
23impl Default for Dimension {
24    fn default() -> Self {
25        Self::D256
26    }
27}
28
29impl From<Dimension> for usize {
30    fn from(d: Dimension) -> usize {
31        d.get()
32    }
33}
34
35impl TryFrom<usize> for Dimension {
36    type Error = Error;
37
38    fn try_from(value: usize) -> Result<Self> {
39        if value == 0 {
40            Err(Error::Config("embedding dimension must be nonzero".into()))
41        } else {
42            Ok(Self(value))
43        }
44    }
45}
46
47impl std::fmt::Display for Dimension {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        write!(f, "{}", self.get())
50    }
51}
52
53impl std::str::FromStr for Dimension {
54    type Err = Error;
55
56    fn from_str(s: &str) -> Result<Self> {
57        let n: usize = s
58            .parse()
59            .map_err(|_| Error::Config(format!("invalid dimension `{s}`")))?;
60        Dimension::try_from(n)
61    }
62}
63
64/// A vector with two explicit construction policies.
65///
66/// [`Embedding::from_raw`] truncates and L2-normalises model output;
67/// [`Embedding::from_preserved`] validates nonempty, finite input and preserves
68/// every component's bits without truncation or normalization. A zero vector
69/// remains zero and reports `normalized() == false`.
70/// Derived deserialization retains its existing, unvalidated behavior.
71#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
72pub struct Embedding(Vec<f32>);
73
74impl Embedding {
75    /// Truncate a raw model output to `dim` leading components, then L2-normalise.
76    ///
77    /// Normalising after truncation matters: a slice of a unit vector is not
78    /// unit length, and skipping this silently degrades ranking.
79    pub fn from_raw(mut raw: Vec<f32>, dim: Dimension) -> Result<Self> {
80        let d = dim.get();
81        if raw.len() < d {
82            return Err(Error::DimensionMismatch {
83                expected: d,
84                actual: raw.len(),
85            });
86        }
87        raw.truncate(d);
88        l2_normalize(&mut raw);
89        if raw.iter().any(|x| !x.is_finite()) {
90            return Err(Error::Embedding(
91                "model produced a non-finite embedding".into(),
92            ));
93        }
94        Ok(Self(raw))
95    }
96
97    /// Restore a nonempty, finite vector without truncating or normalizing it.
98    /// Every component's bits are preserved, including signed zero and subnormals.
99    pub fn from_preserved(values: Vec<f32>) -> Result<Self> {
100        if values.is_empty() {
101            return Err(Error::Embedding(
102                "cannot preserve an empty embedding".into(),
103            ));
104        }
105        if values.iter().any(|value| !value.is_finite()) {
106            return Err(Error::Embedding(
107                "cannot preserve a non-finite embedding".into(),
108            ));
109        }
110        Ok(Self(values))
111    }
112
113    pub fn as_slice(&self) -> &[f32] {
114        &self.0
115    }
116
117    pub fn dim(&self) -> usize {
118        self.0.len()
119    }
120
121    /// Width derived from the vector, never duplicated mutable metadata.
122    pub fn dimension(&self) -> usize {
123        self.dim()
124    }
125
126    /// Whether this vector has unit L2 norm (within 1e-5).
127    pub fn normalized(&self) -> bool {
128        let norm = self.0.iter().map(|v| v * v).sum::<f32>().sqrt();
129        norm.is_finite() && (norm - 1.0).abs() < 1e-5
130    }
131
132    pub fn into_vec(self) -> Vec<f32> {
133        self.0
134    }
135
136    /// Raw dot product, equal to cosine similarity only when both embeddings
137    /// are L2-normalised. This method does not normalise either operand;
138    /// deserialised, unnormalised vectors therefore give non-cosine values.
139    pub fn cosine(&self, other: &Embedding) -> f32 {
140        dot(&self.0, &other.0)
141    }
142}
143
144// Preserve the historical model-output arithmetic and cached vector bytes.
145pub(crate) fn l2_normalize(v: &mut [f32]) {
146    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
147    if norm > 0.0 {
148        v.iter_mut().for_each(|x| *x /= norm);
149    }
150}
151
152/// Dot product with eight independent accumulators.
153///
154/// A plain `iter().sum()` is a strict sequential float reduction that LLVM may
155/// not reorder, so it does not vectorise. Separate lanes let it use SIMD; this
156/// was 5-6x faster on 256- and 768-d vectors (benches/similarity.rs, dot_kernel).
157#[inline]
158pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
159    debug_assert_eq!(a.len(), b.len());
160    let ((ca, ra), (cb, rb)) = (a.as_chunks::<8>(), b.as_chunks::<8>());
161    let tail: f32 = ra.iter().zip(rb).map(|(x, y)| x * y).sum();
162    let mut acc = [0.0f32; 8];
163    for (x, y) in ca.iter().zip(cb) {
164        for i in 0..8 {
165            acc[i] += x[i] * y[i];
166        }
167    }
168    acc.iter().sum::<f32>() + tail
169}
170
171/// Dot product. Unequal dimensions are errors, including in release builds.
172pub fn dot_product(a: &[f32], b: &[f32]) -> Result<f32> {
173    if a.len() != b.len() {
174        return Err(Error::DimensionMismatch {
175            expected: a.len(),
176            actual: b.len(),
177        });
178    }
179    let value = dot(a, b);
180    if !value.is_finite() {
181        return Err(Error::Embedding("non-finite dot product".into()));
182    }
183    Ok(value)
184}
185
186/// Cosine similarity, including for vectors deserialized without normalization.
187/// Zero vectors have no cosine and return an error.
188pub fn cosine_similarity(a: &Embedding, b: &Embedding) -> Result<f32> {
189    let product = dot_product(a.as_slice(), b.as_slice())?;
190    let norm_a = a.as_slice().iter().map(|v| v * v).sum::<f32>().sqrt();
191    let norm_b = b.as_slice().iter().map(|v| v * v).sum::<f32>().sqrt();
192    let denominator = norm_a * norm_b;
193    if denominator == 0.0 || !denominator.is_finite() {
194        return Err(Error::Embedding(
195            "cosine requires finite nonzero norms".into(),
196        ));
197    }
198    Ok(product / denominator)
199}
200
201/// L2-normalize in place using f64 accumulation to avoid f32 overflow/underflow.
202/// Zero (including empty) and non-finite inputs fail without modifying the input.
203pub fn normalize(values: &mut [f32]) -> Result<()> {
204    if values.iter().any(|v| !v.is_finite()) {
205        return Err(Error::Embedding(
206            "cannot normalize non-finite values".into(),
207        ));
208    }
209    let norm = values
210        .iter()
211        .map(|&v| f64::from(v).powi(2))
212        .sum::<f64>()
213        .sqrt();
214    if norm == 0.0 {
215        return Err(Error::Embedding("cannot normalize a zero vector".into()));
216    }
217    for value in values {
218        *value = (f64::from(*value) / norm) as f32;
219    }
220    Ok(())
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn preserved_vectors_round_trip_exact_bits() {
229        for values in [
230            vec![0.0, -0.0],
231            vec![
232                f32::from_bits(1),
233                -f32::from_bits(1),
234                f32::from_bits(0x007fffff),
235            ],
236            vec![3.0, -4.0, f32::MIN_POSITIVE, f32::MIN, f32::MAX],
237            vec![2.0; 769],
238        ] {
239            let bits: Vec<_> = values.iter().map(|v| v.to_bits()).collect();
240            let embedding = Embedding::from_preserved(values).unwrap();
241            assert_eq!(embedding.dim(), bits.len());
242            assert_eq!(
243                embedding
244                    .as_slice()
245                    .iter()
246                    .map(|v| v.to_bits())
247                    .collect::<Vec<_>>(),
248                bits
249            );
250            assert_eq!(
251                embedding
252                    .into_vec()
253                    .iter()
254                    .map(|v| v.to_bits())
255                    .collect::<Vec<_>>(),
256                bits
257            );
258        }
259    }
260
261    #[test]
262    fn preserved_vectors_reject_empty_and_non_finite_input() {
263        assert!(matches!(
264            Embedding::from_preserved(vec![]),
265            Err(Error::Embedding(_))
266        ));
267        for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
268            assert!(matches!(
269                Embedding::from_preserved(vec![1.0, value]),
270                Err(Error::Embedding(_))
271            ));
272        }
273    }
274
275    #[test]
276    fn truncation_renormalizes() {
277        let raw: Vec<f32> = (0..768).map(|i| (i as f32).sin()).collect();
278        let e = Embedding::from_raw(raw, Dimension::D128).unwrap();
279        assert_eq!(e.dim(), 128);
280        let norm: f32 = e.as_slice().iter().map(|x| x * x).sum::<f32>().sqrt();
281        assert!((norm - 1.0).abs() < 1e-5);
282    }
283
284    #[test]
285    fn rejects_short_vectors() {
286        assert!(Embedding::from_raw(vec![1.0; 100], Dimension::D128).is_err());
287    }
288
289    #[test]
290    fn dot_matches_naive_sum() {
291        for n in [1, 7, 8, 9, 128, 131, 768] {
292            let a: Vec<f32> = (0..n).map(|i| (i as f32 * 0.37).sin()).collect();
293            let b: Vec<f32> = (0..n).map(|i| (i as f32 * 0.11).cos()).collect();
294            let naive: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
295            assert!((dot(&a, &b) - naive).abs() < 1e-4, "n={n}");
296        }
297    }
298
299    #[test]
300    fn dimension_parse() {
301        assert_eq!("512".parse::<Dimension>().unwrap(), Dimension::D512);
302        assert_eq!("300".parse::<Dimension>().unwrap().get(), 300);
303        assert!("0".parse::<Dimension>().is_err());
304    }
305}
seyk_embed/embedding.rs

seyk_embed/
embedding.rs
