From 941869ad24c437aad096c820f0ade0e05ef3689d Mon Sep 17 00:00:00 2001
From: Sampreeth Sarma <sivauday.sampreeth8@gmail.com>
Date: Fri, 20 Sep 2024 09:49:28 -0700
Subject: [PATCH] fix similarity bug (#63)

---
 envoyfilter/src/stream_context.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/envoyfilter/src/stream_context.rs b/envoyfilter/src/stream_context.rs
index 557f0c90..080dc2fa 100644
--- a/envoyfilter/src/stream_context.rs
+++ b/envoyfilter/src/stream_context.rs
@@ -262,14 +262,25 @@ impl StreamContext {
 
         debug!("zeroshot intent response: {:?}", zeroshot_intent_response);
 
+        let desc_emb_similarity_map: HashMap<String, f64> = callout_context
+            .similarity_scores
+            .clone()
+            .unwrap()
+            .into_iter()
+            .collect();
+
+        let pred_class_desc_emb_similarity = desc_emb_similarity_map
+            .get(&zeroshot_intent_response.predicted_class)
+            .unwrap();
+
         let prompt_target_similarity_score = zeroshot_intent_response.predicted_class_score * 0.7
-            + callout_context.similarity_scores.as_ref().unwrap()[0].1 * 0.3;
+            + pred_class_desc_emb_similarity * 0.3;
 
         debug!(
             "similarity score: {:.3}, intent score: {:.3}, description embedding score: {:.3}",
             prompt_target_similarity_score,
             zeroshot_intent_response.predicted_class_score,
-            callout_context.similarity_scores.as_ref().unwrap()[0].1
+            pred_class_desc_emb_similarity
         );
 
         let prompt_target_name = zeroshot_intent_response.predicted_class.clone();