skm_embed/
embedding.rs

1//! Embedding vector type with similarity operations.
2
3use serde::{Deserialize, Serialize};
4
5use crate::simd;
6
7/// A single embedding vector with metadata.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct Embedding {
10    /// The embedding vector (normalized to unit length).
11    pub vector: Vec<f32>,
12
13    /// Hash of the original text for cache lookups.
14    pub text_hash: u64,
15}
16
17impl Embedding {
18    /// Create a new embedding from a vector.
19    /// Automatically normalizes the vector to unit length.
20    pub fn new(mut vector: Vec<f32>, text_hash: u64) -> Self {
21        simd::normalize(&mut vector);
22        Self { vector, text_hash }
23    }
24
25    /// Create an embedding without normalizing (assumes already normalized).
26    pub fn from_normalized(vector: Vec<f32>, text_hash: u64) -> Self {
27        Self { vector, text_hash }
28    }
29
30    /// Get the dimensionality of this embedding.
31    pub fn dimensions(&self) -> usize {
32        self.vector.len()
33    }
34
35    /// Cosine similarity with another embedding.
36    /// For normalized vectors, this equals the dot product.
37    pub fn cosine_similarity(&self, other: &Embedding) -> f32 {
38        // For pre-normalized vectors, cosine similarity = dot product
39        simd::dot_product(&self.vector, &other.vector)
40    }
41
42    /// Dot product (for pre-normalized vectors, equals cosine similarity).
43    pub fn dot_product(&self, other: &Embedding) -> f32 {
44        simd::dot_product(&self.vector, &other.vector)
45    }
46
47    /// Euclidean distance to another embedding.
48    pub fn euclidean_distance(&self, other: &Embedding) -> f32 {
49        let mut sum = 0.0f32;
50        for (a, b) in self.vector.iter().zip(other.vector.iter()) {
51            let diff = a - b;
52            sum += diff * diff;
53        }
54        sum.sqrt()
55    }
56
57    /// L2 norm of the vector.
58    pub fn norm(&self) -> f32 {
59        simd::dot_product(&self.vector, &self.vector).sqrt()
60    }
61
62    /// Check if this embedding is normalized (unit length).
63    pub fn is_normalized(&self) -> bool {
64        let norm = self.norm();
65        (norm - 1.0).abs() < 1e-5
66    }
67}
68
69impl PartialEq for Embedding {
70    fn eq(&self, other: &Self) -> bool {
71        self.text_hash == other.text_hash && self.vector == other.vector
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    fn make_embedding(values: &[f32]) -> Embedding {
80        Embedding::new(values.to_vec(), 0)
81    }
82
83    #[test]
84    fn test_embedding_normalization() {
85        let embed = make_embedding(&[3.0, 4.0]);
86        assert!(embed.is_normalized());
87        // 3-4-5 triangle, normalized: (0.6, 0.8)
88        assert!((embed.vector[0] - 0.6).abs() < 1e-5);
89        assert!((embed.vector[1] - 0.8).abs() < 1e-5);
90    }
91
92    #[test]
93    fn test_cosine_similarity_identical() {
94        let embed = make_embedding(&[1.0, 0.0, 0.0]);
95        let similarity = embed.cosine_similarity(&embed);
96        assert!((similarity - 1.0).abs() < 1e-5);
97    }
98
99    #[test]
100    fn test_cosine_similarity_orthogonal() {
101        let a = make_embedding(&[1.0, 0.0]);
102        let b = make_embedding(&[0.0, 1.0]);
103        let similarity = a.cosine_similarity(&b);
104        assert!(similarity.abs() < 1e-5);
105    }
106
107    #[test]
108    fn test_cosine_similarity_opposite() {
109        let a = make_embedding(&[1.0, 0.0]);
110        let b = make_embedding(&[-1.0, 0.0]);
111        let similarity = a.cosine_similarity(&b);
112        assert!((similarity + 1.0).abs() < 1e-5);
113    }
114
115    #[test]
116    fn test_euclidean_distance() {
117        let a = make_embedding(&[1.0, 0.0]);
118        let b = make_embedding(&[0.0, 1.0]);
119        let dist = a.euclidean_distance(&b);
120        // For normalized vectors, opposite corners: sqrt(2) * something
121        assert!(dist > 0.0);
122    }
123
124    #[test]
125    fn test_from_normalized() {
126        let embed = Embedding::from_normalized(vec![0.6, 0.8], 123);
127        assert!(embed.is_normalized());
128        assert_eq!(embed.text_hash, 123);
129    }
130
131    #[test]
132    fn test_dimensions() {
133        let embed = make_embedding(&[1.0, 2.0, 3.0, 4.0]);
134        assert_eq!(embed.dimensions(), 4);
135    }
136}
skm_embed/embedding.rs

skm_embed/
embedding.rs
