deduplicate w.r.t. mean pooled vector

Author: ggordonhall

server/bleep/src/semantic.rs

@@ -30,6 +30,7 @@ pub use schema::{Embedding, Payload};
 
 const COLLECTION_NAME: &str = "documents";
 const SCORE_THRESHOLD: f32 = 0.3;
+const EMBEDDING_DIM: usize = 384;
 
 #[derive(Error, Debug)]
 pub enum SemanticError {
@@ -149,7 +150,7 @@ fn collection_config() -> CreateCollection {
         collection_name: COLLECTION_NAME.to_string(),
         vectors_config: Some(VectorsConfig {
             config: Some(vectors_config::Config::Params(VectorParams {
-                size: 384,
+                size: EMBEDDING_DIM as u64,
                 distance: Distance::Cosine.into(),
                 ..Default::default()
             })),
@@ -396,8 +397,9 @@ impl Semantic {
                     .collect::<Vec<_>>()
             })?;
 
-        // TODO: Deduplicate with respect to all vectors
-        let target_vector = vectors.first().unwrap().clone();
+        // deduplicate with mmr with respect to the mean of query vectors
+        // TODO: implement a more robust multi-vector deduplication strategy
+        let target_vector = mean_pool(vectors);
         Ok(deduplicate_snippets(results, target_vector, limit))
     }
 
@@ -629,6 +631,19 @@ fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
     dot(a, b) / (norm(a) * norm(b))
 }
 
+// Calculate the element-wise mean of the embeddings
+fn mean_pool(embeddings: Vec<Vec<f32>>) -> Vec<f32> {
+    let len = embeddings.len() as f32;
+    let mut result = vec![0.0; EMBEDDING_DIM];
+    for embedding in embeddings {
+        for (i, v) in embedding.iter().enumerate() {
+            result[i] += v;
+        }
+    }
+    result.iter_mut().for_each(|v| *v /= len);
+    result
+}
+
 // returns a list of indices to preserve from `snippets`
 //
 // query_embedding: the embedding of the query terms

server/bleep/src/webserver/answer/prompts.rs

@@ -261,7 +261,7 @@ pub fn hypothetical_document_prompt(query: &str) -> String {
     format!(
         r#"Write three code snippets that could hypothetically be returned by a code search engine as the answer to the query: {query}
 
-- All three snippets should be written in any one of these languages: [Rust, Typescript, TSX, YAML]
+- Write these snippets in a variety of programming languages
 - The snippets should not be too similar to one another
 - Each snippet should be between 5 and 10 lines long
 - Surround the snippets in triple backticks