Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import torch
 from sentence_transformers import SentenceTransformer, util
-#from spellchecker import SpellChecker
 import pickle
 import re
 
@@ -27,19 +26,7 @@ def validate_input(input_string):
     else:
         return False
 
-# Define the function for mapping code
-def mapping_code(user_input, mode):
-    if mode == "CPT_to_SBS":
-        stored_embeddings_cpt = stored_embeddings
-        stored_data_cpt = stored_data_cpt
-        code_column = stored_data["CPT_CODE"]
-        description_column = stored_data["FULL_DESCRIPTION"]
-    elif mode == "SBS_to_CPT":
-        stored_embeddings = stored_embeddings
-        stored_data = stored_data_sbs
-        code_column = stored_data["SBS_code"]
-        description_column = stored_data["Description"]
-
+def cpt_code(user_input):
     emb1 = model.encode(user_input.lower())
     similarities = []
     for sentence in stored_embeddings:
@@ -47,7 +34,7 @@ def mapping_code(user_input, mode):
         similarities.append(similarity)
 
     # Filter results with similarity scores above 0.70
-    result = [(code, desc, sim) for (code, desc, sim) in zip(code_column, description_column, similarities)]
+    result = [(code, desc, sim) for (code, desc, sim) in zip(stored_data["SBS_code"], stored_data["Description"], similarities)]
 
     # Sort results by similarity scores
     result.sort(key=lambda x: x[2], reverse=True)
@@ -65,9 +52,39 @@ def mapping_code(user_input, mode):
 
     return top_5_results
 
-# Streamlit frontend interface
-import streamlit as st
+def sbs_code(user_input):
+    emb1 = model.encode(user_input.lower())
+    similarities = []
+    for sentence in stored_embeddings_cpt:
+        similarity = util.cos_sim(sentence, emb1)
+        similarities.append(similarity)
+
+    # Filter results with similarity scores above 0.70
+    result = [(code, desc, sim) for (code, desc, sim) in zip(stored_data_cpt["CPT_CODE"], stored_data_cpt["FULL_DESCRIPTION"], similarities)]
+
+    # Sort results by similarity scores
+    result.sort(key=lambda x: x[2], reverse=True)
+
+    num_results = min(5, len(result))
 
+    # Return top 5 entries with 'code', 'description', and 'similarity_score'
+    top_5_results = []
+    if num_results > 0:
+        for i in range(num_results):
+            code, description, similarity_score = result[i]
+            top_5_results.append({"Code": code, "Description": description, "Similarity Score": similarity_score})
+    else:
+        top_5_results.append({"Code": "", "Description": "No match", "Similarity Score": 0.0})
+
+    return top_5_results
+
+def mapping_code(user_input, mode):
+    if mode == "CPT_to_SBS":
+        return cpt_code(user_input)
+    elif mode == "SBS_to_CPT":
+        return sbs_code(user_input)
+
+# Streamlit frontend interface
 def main():
     st.title("CPT-SBS Code Mapping")
 
@@ -104,4 +121,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-