new: improved cosine similarity by parallelizing it with rayon

src/agent/rag/metrics.rs

@@ -1,10 +1,14 @@
+#[cfg(feature = "rayon")]
+use rayon::prelude::*;
+
 use super::Embeddings;
 
 /// Cosine distance between two vectors
 ///
 /// When the features distances lengths don't match, the longer feature vector is truncated to
 /// shorter one when the distance is calculated
 ///  
+#[cfg(not(feature = "rayon"))]
 #[inline]
 pub fn cosine(vec_a: &Embeddings, vec_b: &Embeddings) -> f64 {
     assert_eq!(vec_a.len(), vec_b.len());
@@ -23,3 +27,19 @@ pub fn cosine(vec_a: &Embeddings, vec_b: &Embeddings) -> f64 {
 
     1.0 - (a_dot_b / (a_mag.sqrt() * b_mag.sqrt()))
 }
+
+#[cfg(feature = "rayon")]
+#[inline]
+pub fn cosine(vec_a: &Embeddings, vec_b: &Embeddings) -> f64 {
+    assert_eq!(vec_a.len(), vec_b.len());
+
+    let dot_product: f64 = vec_a
+        .par_iter()
+        .zip(vec_b.par_iter())
+        .map(|(a, b)| a * b)
+        .sum();
+    let magnitude1: f64 = vec_a.par_iter().map(|a| a * a).sum::<f64>().sqrt();
+    let magnitude2: f64 = vec_b.par_iter().map(|b| b * b).sum::<f64>().sqrt();
+
+    1.0 - dot_product / (magnitude1 * magnitude2)
+}

src/agent/rag/naive.rs

@@ -77,11 +77,12 @@ impl VectorStore for NaiveVectorStore {
         let start = Instant::now();
         let doc_name = document.name.to_string();
         let embeddings = self.embedder.embeddings(&document.data).await?;
+        let size = embeddings.len();
 
         self.documents.insert(doc_name.to_string(), document);
         self.embeddings.insert(doc_name, embeddings);
 
-        println!(" done in {:?}", start.elapsed());
+        println!(" time={:?} embedding_size={}", start.elapsed(), size);
 
         Ok(())
     }