add language diversity factor to MMR score

Author: ggordonhall

server/bleep/src/semantic.rs

@@ -551,10 +551,12 @@ fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
 pub fn deduplicate_with_mmr(
     query_embedding: &[f32],
     embeddings: &[&[f32]],
+    languages: &[&str],
     lambda: f32,
     k: usize,
 ) -> Vec<usize> {
     let mut idxs = vec![];
+    let mut lang_counts = HashMap::new();
 
     if embeddings.len() < k {
         return (0..embeddings.len()).collect();
@@ -576,14 +578,20 @@ pub fn deduplicate_with_mmr(
                     second_part = cos_sim;
                 }
             }
-            let equation_score = lambda * first_part - (1. - lambda) * second_part;
+            let mut equation_score = lambda * first_part - (1. - lambda) * second_part;
+
+            // score is MMR + (1/4)^n where n is the number of times a language has been selected
+            let count = lang_counts.get(languages[i]).unwrap_or(&0);
+            equation_score += 0.25_f32.powi(*count);
+
             if equation_score > best_score {
                 best_score = equation_score;
                 idx_to_add = Some(i);
             }
         }
         if let Some(i) = idx_to_add {
             idxs.push(i);
+            *lang_counts.entry(languages[i]).or_insert(0) += 1;
         }
     }
     idxs
@@ -601,7 +609,11 @@ pub fn deduplicate_snippets(
             .iter()
             .map(|s| s.embedding.as_deref().unwrap())
             .collect::<Vec<_>>();
-        deduplicate_with_mmr(&query_embedding, &embeddings, lambda, k)
+        let languages = all_snippets
+            .iter()
+            .map(|s| s.lang.as_ref())
+            .collect::<Vec<_>>();
+        deduplicate_with_mmr(&query_embedding, &embeddings, &languages, lambda, k)
     };
 
     info!("preserved idxs after MMR are {:?}", idxs);