Update icd9_ui.py

icd9_ui.py

@@ -1,5 +1,70 @@
-import streamlit as st
+# import streamlit as st
+# import torch
+# from transformers import LongformerTokenizer, LongformerForSequenceClassification
+
+# # Load the fine-tuned model and tokenizer
+# model_path = "./clinical_longformer"
+# tokenizer = LongformerTokenizer.from_pretrained(model_path)
+# model = LongformerForSequenceClassification.from_pretrained(model_path)
+# model.eval()  # Set the model to evaluation mode
+
+# # ICD-9 code columns used during training
+# icd9_columns = [
+#     '038.9', '244.9', '250.00', '272.0', '272.4', '276.1', '276.2', '285.1', '285.9',
+#     '287.5', '305.1', '311', '36.15', '37.22', '37.23', '38.91', '38.93', '39.61',
+#     '39.95', '401.9', '403.90', '410.71', '412', '414.01', '424.0', '427.31', '428.0',
+#     '486', '496', '507.0', '511.9', '518.81', '530.81', '584.9', '585.9', '599.0',
+#     '88.56', '88.72', '93.90', '96.04', '96.6', '96.71', '96.72', '99.04', '99.15',
+#     '995.92', 'V15.82', 'V45.81', 'V45.82', 'V58.61'
+# ]
+
+# # Function for making predictions
+# def predict_icd9(texts, tokenizer, model, threshold=0.5):
+#     inputs = tokenizer(
+#         texts,
+#         padding="max_length",
+#         truncation=True,
+#         max_length=512,
+#         return_tensors="pt"
+#     )
+    
+#     with torch.no_grad():
+#         outputs = model(
+#             input_ids=inputs["input_ids"],
+#             attention_mask=inputs["attention_mask"]
+#         )
+#         logits = outputs.logits
+#         probabilities = torch.sigmoid(logits)
+#         predictions = (probabilities > threshold).int()
+    
+#     predicted_icd9 = []
+#     for pred in predictions:
+#         codes = [icd9_columns[i] for i, val in enumerate(pred) if val == 1]
+#         predicted_icd9.append(codes)
+    
+#     return predicted_icd9
+
+# # Streamlit UI
+# st.title("ICD-9 Code Prediction")
+# st.sidebar.header("Model Options")
+# model_option = st.sidebar.selectbox("Select Model", [ "ClinicalLongformer"])
+# threshold = st.sidebar.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
+
+# st.write("### Enter Medical Summary")
+# input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes here...")
+
+# if st.button("Predict"):
+#     if input_text.strip():
+#         predictions = predict_icd9([input_text], tokenizer, model, threshold)
+#         st.write("### Predicted ICD-9 Codes")
+#         for code in predictions[0]:
+#             st.write(f"- {code}")
+#     else:
+#         st.error("Please enter a medical summary.")
+
 import torch
+import pandas as pd
+import streamlit as st
 from transformers import LongformerTokenizer, LongformerForSequenceClassification
 
 # Load the fine-tuned model and tokenizer
@@ -8,6 +73,11 @@ tokenizer = LongformerTokenizer.from_pretrained(model_path)
 model = LongformerForSequenceClassification.from_pretrained(model_path)
 model.eval()  # Set the model to evaluation mode
 
+# Load the ICD-9 descriptions from CSV into a dictionary
+icd9_desc_df = pd.read_csv("D_ICD_DIAGNOSES.csv")  # Adjust the path to your CSV file
+icd9_desc_df['ICD9_CODE'] = icd9_desc_df['ICD9_CODE'].astype(str)  # Ensure ICD9_CODE is string type for matching
+icd9_descriptions = dict(zip(icd9_desc_df['ICD9_CODE'].str.replace('.', ''), icd9_desc_df['LONG_TITLE']))  # Remove decimals in ICD9 code for matching
+
 # ICD-9 code columns used during training
 icd9_columns = [
     '038.9', '244.9', '250.00', '272.0', '272.4', '276.1', '276.2', '285.1', '285.9',
@@ -42,12 +112,17 @@ def predict_icd9(texts, tokenizer, model, threshold=0.5):
         codes = [icd9_columns[i] for i, val in enumerate(pred) if val == 1]
         predicted_icd9.append(codes)
     
-    return predicted_icd9
+    # Fetch descriptions for the predicted ICD-9 codes from the pre-loaded descriptions
+    predictions_with_desc = []
+    for codes in predicted_icd9:
+        code_with_desc = [(code, icd9_descriptions.get(code.replace('.', ''), "Description not found")) for code in codes]
+        predictions_with_desc.append(code_with_desc)
+    
+    return predictions_with_desc
 
 # Streamlit UI
 st.title("ICD-9 Code Prediction")
 st.sidebar.header("Model Options")
-model_option = st.sidebar.selectbox("Select Model", [ "ClinicalLongformer"])
 threshold = st.sidebar.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
 
 st.write("### Enter Medical Summary")
@@ -56,8 +131,9 @@ input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes h
 if st.button("Predict"):
     if input_text.strip():
         predictions = predict_icd9([input_text], tokenizer, model, threshold)
-        st.write("### Predicted ICD-9 Codes")
-        for code in predictions[0]:
-            st.write(f"- {code}")
+        st.write("### Predicted ICD-9 Codes and Descriptions")
+        for code, description in predictions[0]:
+            st.write(f"- {code}: {description}")
     else:
         st.error("Please enter a medical summary.")
+