Update app.py

app.py

@@ -160,8 +160,16 @@ def predict_icd(text_input, model_name, label_count):
     label_probabilities = {label_map[i]: prob * 100 for i, prob in enumerate(probabilities[0]) if label_map[i] in predicted_labels}
     y_pred = np.array(predicted_labels)
     y_pred = "\n".join(y_pred)
-    label_dict = label_dict.set_index('icd9_code')['long_title'].to_dict()
-    predicted_labels_with_titles = [(label, label_dict[label]) for label in predicted_labels]
+    label_dict_df = pd.read_csv("data/mimic3/"+label_count+"/labels_dictionary_"+label_count+"_level_1.csv")
+    # label_dict = label_dict.set_index('icd9_code')['long_title'].to_dict()
+    # predicted_labels_with_titles = [(label, label_dict[label]) for label in predicted_labels]
+
+    label_dict_df['icd9_code'] = label_dict_df['icd9_code'].astype(str)
+    label_dict = label_dict_df.set_index('icd9_code')['long_title'].to_dict()
+
+    # Retrieve long titles from label_dict
+    predicted_labels_with_titles = [(label, label_dict.get(label, "Not Found")) for label in predicted_labels]
+
     
     # html_output ="<h2>ICD Codes</h2>"
     # for label in predicted_labels: