Skip to main content

spatialrust_semantic/
embedding.rs

1//! Dense embedding vectors.
2
3use crate::{SemanticError, SemanticResult};
4
5/// Dense float embedding.
6#[derive(Clone, Debug, PartialEq)]
7pub struct Embedding {
8    values: Vec<f32>,
9}
10
11impl Embedding {
12    /// Creates an embedding after rejecting empty/non-finite vectors.
13    pub fn try_new(values: Vec<f32>) -> SemanticResult<Self> {
14        if values.is_empty() {
15            return Err(SemanticError::InvalidConfiguration("embedding must be non-empty".into()));
16        }
17        if values.iter().any(|v| !v.is_finite()) {
18            return Err(SemanticError::InvalidConfiguration("embedding must be finite".into()));
19        }
20        Ok(Self { values })
21    }
22
23    /// Returns dimensionality.
24    #[must_use]
25    pub fn dim(&self) -> usize {
26        self.values.len()
27    }
28
29    /// Borrows values.
30    #[must_use]
31    pub fn as_slice(&self) -> &[f32] {
32        &self.values
33    }
34}
35
36/// Cosine similarity in `[-1, 1]`.
37pub fn cosine_similarity(a: &Embedding, b: &Embedding) -> SemanticResult<f32> {
38    if a.dim() != b.dim() {
39        return Err(SemanticError::InvalidConfiguration("embedding dims must match".into()));
40    }
41    let mut dot = 0.0_f32;
42    let mut na = 0.0_f32;
43    let mut nb = 0.0_f32;
44    for (x, y) in a.as_slice().iter().zip(b.as_slice()) {
45        dot += x * y;
46        na += x * x;
47        nb += y * y;
48    }
49    let denom = na.sqrt() * nb.sqrt();
50    if denom < 1e-12 {
51        return Err(SemanticError::InvalidConfiguration("zero-norm embedding".into()));
52    }
53    Ok(dot / denom)
54}
55
56#[cfg(test)]
57mod tests {
58    use super::{cosine_similarity, Embedding};
59
60    #[test]
61    fn identical_vectors_have_unit_cosine() {
62        let a = Embedding::try_new(vec![1.0, 0.0, 0.0]).unwrap();
63        let b = Embedding::try_new(vec![2.0, 0.0, 0.0]).unwrap();
64        assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-5);
65    }
66}