more comet

src/streamlit_app.py

@@ -103,16 +103,20 @@ def get_chat_completion_openai(client, prompt: str):
     )
 
 
+def get_similar_prompt(query_embedding, text_embeddings):
+    cos_scores = util.pytorch_cos_sim(query_embedding, text_embeddings)[0]
+    top_result = torch.topk(cos_scores, k=1)
+    idx = top_result.indices[0].item()
+    return data[idx]
+
+
 if query:
     with st.spinner("Searching medical cases..."):
         text_embeddings = embed_texts(data[TEXT_COLUMN])
         query_embedding = embed_texts([query])[0]
 
         # Compute similarity
-        cos_scores = util.pytorch_cos_sim(query_embedding, text_embeddings)[0]
-        top_result = torch.topk(cos_scores, k=1)
-        idx = top_result.indices[0].item()
-        selected = data[idx]
+        selected = get_similar_prompt(query_embedding, text_embeddings)
 
         # Show Image
         st.image(selected['image'], caption="Most relevant medical image", use_container_width=True)