Update icd9_ui.py

icd9_ui.py

@@ -1,73 +1,428 @@
-# 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"
-#     )
+# # 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()
+# #     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)
+# #     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
+# # model_path = "./clinical_longformer"
+# # 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',
+# #     '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)
+    
+# #     # 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")
+# # 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 and Descriptions")
+# #         for code, description in predictions[0]:
+# #             st.write(f"- {code}: {description}")
+# #     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
+# # model_path = "./clinical_longformer"
+# # 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
+# # icd9_descriptions = dict(zip(icd9_desc_df['ICD9_CODE'].str.replace('.', ''), icd9_desc_df['LONG_TITLE']))  # Remove decimals for matching
+
+# # # Load the ICD-9 to ICD-10 mapping
+# # icd9_to_icd10 = {}
+# # with open("2015_I9gem.txt", "r") as file:
+# #     for line in file:
+# #         parts = line.strip().split()
+# #         if len(parts) == 3:
+# #             icd9, icd10, _ = parts
+# #             icd9_to_icd10[icd9] = icd10
+
+# # # 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 and mapping to ICD-10
+# # 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)
+
+# #     # Fetch descriptions and map to ICD-10 codes
+# #     predictions_with_desc = []
+# #     for codes in predicted_icd9:
+# #         code_with_desc = []
+# #         for code in codes:
+# #             icd9_stripped = code.replace('.', '')
+# #             icd10_code = icd9_to_icd10.get(icd9_stripped, "Mapping not found")
+# #             icd9_desc = icd9_descriptions.get(icd9_stripped, "Description not found")
+# #             code_with_desc.append((code, icd9_desc, icd10_code))
+# #         predictions_with_desc.append(code_with_desc)
+
+# #     return predictions_with_desc
+
+# # # Streamlit UI
+# # st.title("ICD-9 to ICD-10 Code Prediction")
+# # st.sidebar.header("Model Options")
+# # 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 and ICD-10 Codes with Descriptions")
+# #         for icd9_code, description, icd10_code in predictions[0]:
+# #             st.write(f"- ICD-9: {icd9_code} ({description}) -> ICD-10: {icd10_code}")
+# #     else:
+# #         st.error("Please enter a medical summary.")
+
+# # import os
+# # import torch
+# # import pandas as pd
+# # import streamlit as st
+# # from PIL import Image
+# # from transformers import LongformerTokenizer, LongformerForSequenceClassification
+# # from phi.agent import Agent
+# # from phi.model.google import Gemini
+# # from phi.tools.duckduckgo import DuckDuckGo
+
+# # # Load the fine-tuned ICD-9 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
+
+# # # 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',
+# #     '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 ICD-9 predictions
+# # def predict_icd9(texts, tokenizer, model, threshold=0.5):
+#     # inputs = tokenizer(
+#     #     texts,
+#     #     padding="max_length",
+#     #     truncation=True,
+#     #     max_length=512,
+#     #     return_tensors="pt"
+#     # )
     
-#     return predicted_icd9
+#     # 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)
+    
+#     # 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")
-# 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.")
+# # st.title("Medical Diagnosis Assistant")
+# # option = st.selectbox(
+# #     "Choose Diagnosis Method",
+# #     ("ICD-9 Code Prediction", "Medical Image Analysis")
+# # )
+
+# # # ICD-9 Code Prediction
+# # if option == "ICD-9 Code Prediction":
+# #     st.write("### Enter Medical Summary")
+# #     input_text = st.text_area("Medical Summary", placeholder="Enter clinical notes here...")
+
+# #     threshold = st.slider("Prediction Threshold", 0.0, 1.0, 0.5, 0.01)
+
+# #     if st.button("Predict ICD-9 Codes"):
+# #         if input_text.strip():
+# #             predictions = predict_icd9([input_text], tokenizer, model, threshold)
+# #             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.")
+
+# # Medical Image Analysis
+# # elif option == "Medical Image Analysis":
+# #     if "GOOGLE_API_KEY" not in st.session_state:
+# #         st.warning("Please enter your Google API Key in the sidebar to continue")
+# #     else:
+# #         medical_agent = Agent(
+# #             model=Gemini(
+# #                 api_key=st.session_state.GOOGLE_API_KEY,
+# #                 id="gemini-2.0-flash-exp"
+# #             ),
+# #             tools=[DuckDuckGo()],
+# #             markdown=True
+# #         )
+
+# #         query = """
+# #         You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
+
+# #         ### 1. Image Type & Region
+# #         - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
+# #         - Identify the patient's anatomical region and positioning
+# #         - Comment on image quality and technical adequacy
+
+# #         ### 2. Key Findings
+# #         - List primary observations systematically
+# #         - Note any abnormalities in the patient's imaging with precise descriptions
+# #         - Include measurements and densities where relevant
+# #         - Describe location, size, shape, and characteristics
+# #         - Rate severity: Normal/Mild/Moderate/Severe
+
+# #         ### 3. Diagnostic Assessment
+# #         - Provide primary diagnosis with confidence level
+# #         - List differential diagnoses in order of likelihood
+# #         - Support each diagnosis with observed evidence from the patient's imaging
+#         # - Note any critical or urgent findings
+
+#         # ### 4. Patient-Friendly Explanation
+#         # - Explain the findings in simple, clear language that the patient can understand
+#         # - Avoid medical jargon or provide clear definitions
+#         # - Include visual analogies if helpful
+#         # - Address common patient concerns related to these findings
+
+#         # ### 5. Research Context
+#         # - Use the DuckDuckGo search tool to find recent medical literature about similar cases
+#         # - Provide a list of relevant medical links
+#         # - Include key references to support your analysis
+#         # """
+
+#         # upload_container = st.container()
+#         # image_container = st.container()
+#         # analysis_container = st.container()
+
+#         # with upload_container:
+#         #     uploaded_file = st.file_uploader(
+#         #         "Upload Medical Image",
+#         #         type=["jpg", "jpeg", "png", "dicom"],
+#         #         help="Supported formats: JPG, JPEG, PNG, DICOM"
+#         #     )
+
+#         # if uploaded_file is not None:
+#         #     with image_container:
+#         #         col1, col2, col3 = st.columns([1, 2, 1])
+#         #         with col2:
+#         #             image = Image.open(uploaded_file)
+#         #             width, height = image.size
+#         #             aspect_ratio = width / height
+#         #             new_width = 500
+#         #             new_height = int(new_width / aspect_ratio)
+#         #             resized_image = image.resize((new_width, new_height))
+                    
+#         #             st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
+                    
+#         #             analyze_button = st.button("🔍 Analyze Image")
 
+#         #     with analysis_container:
+#         #         if analyze_button:
+#         #             image_path = "temp_medical_image.png"
+#         #             with open(image_path, "wb") as f:
+#         #                 f.write(uploaded_file.getbuffer())
+                    
+#         #             with st.spinner("🔄 Analyzing image... Please wait."):
+#         #                 try:
+#         #                     response = medical_agent.run(query, images=[image_path])
+#         #                     st.markdown("### 📋 Analysis Results")
+#         #                     st.markdown(response.content)
+#         #                 except Exception as e:
+#         #                     st.error(f"Analysis error: {e}")
+#         #                 finally:
+#         #                     if os.path.exists(image_path):
+#         #                         os.remove(image_path)
+#         # else:
+#         #     st.info("👆 Please upload a medical image to begin analysis")
+
+# import os
 # import torch
 # import pandas as pd
 # import streamlit as st
+# from PIL import Image
 # from transformers import LongformerTokenizer, LongformerForSequenceClassification
+# from phi.agent import Agent
+# from phi.model.google import Gemini
+# from phi.tools.duckduckgo import DuckDuckGo
+
+# # Sidebar for Google API Key input
+# st.sidebar.title("Settings")
+# st.sidebar.write("Enter your Google API Key below for the Medical Image Analysis feature.")
+# api_key = st.sidebar.text_input("Google API Key", type="password")
 
-# # Load the fine-tuned model and tokenizer
+# if api_key:
+#     st.session_state["GOOGLE_API_KEY"] = api_key
+# else:
+#     st.session_state.pop("GOOGLE_API_KEY", None)
+
+# # Load the fine-tuned ICD-9 model and tokenizer
 # model_path = "./clinical_longformer"
 # tokenizer = LongformerTokenizer.from_pretrained(model_path)
 # model = LongformerForSequenceClassification.from_pretrained(model_path)
@@ -88,7 +443,7 @@
 #     '995.92', 'V15.82', 'V45.81', 'V45.82', 'V58.61'
 # ]
 
-# # Function for making predictions
+# # Function for making ICD-9 predictions
 # def predict_icd9(texts, tokenizer, model, threshold=0.5):
 #     inputs = tokenizer(
 #         texts,
@@ -112,7 +467,6 @@
 #         codes = [icd9_columns[i] for i, val in enumerate(pred) if val == 1]
 #         predicted_icd9.append(codes)
     
-#     # 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]
@@ -121,173 +475,6 @@
 #     return predictions_with_desc
 
 # # Streamlit UI
-# st.title("ICD-9 Code Prediction")
-# st.sidebar.header("Model Options")
-# 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 and Descriptions")
-#         for code, description in predictions[0]:
-#             st.write(f"- {code}: {description}")
-#     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
-# model_path = "./clinical_longformer"
-# 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
-# icd9_descriptions = dict(zip(icd9_desc_df['ICD9_CODE'].str.replace('.', ''), icd9_desc_df['LONG_TITLE']))  # Remove decimals for matching
-
-# # Load the ICD-9 to ICD-10 mapping
-# icd9_to_icd10 = {}
-# with open("2015_I9gem.txt", "r") as file:
-#     for line in file:
-#         parts = line.strip().split()
-#         if len(parts) == 3:
-#             icd9, icd10, _ = parts
-#             icd9_to_icd10[icd9] = icd10
-
-# # 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 and mapping to ICD-10
-# 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)
-
-#     # Fetch descriptions and map to ICD-10 codes
-#     predictions_with_desc = []
-#     for codes in predicted_icd9:
-#         code_with_desc = []
-#         for code in codes:
-#             icd9_stripped = code.replace('.', '')
-#             icd10_code = icd9_to_icd10.get(icd9_stripped, "Mapping not found")
-#             icd9_desc = icd9_descriptions.get(icd9_stripped, "Description not found")
-#             code_with_desc.append((code, icd9_desc, icd10_code))
-#         predictions_with_desc.append(code_with_desc)
-
-#     return predictions_with_desc
-
-# # Streamlit UI
-# st.title("ICD-9 to ICD-10 Code Prediction")
-# st.sidebar.header("Model Options")
-# 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 and ICD-10 Codes with Descriptions")
-#         for icd9_code, description, icd10_code in predictions[0]:
-#             st.write(f"- ICD-9: {icd9_code} ({description}) -> ICD-10: {icd10_code}")
-#     else:
-#         st.error("Please enter a medical summary.")
-
-# import os
-# import torch
-# import pandas as pd
-# import streamlit as st
-# from PIL import Image
-# from transformers import LongformerTokenizer, LongformerForSequenceClassification
-# from phi.agent import Agent
-# from phi.model.google import Gemini
-# from phi.tools.duckduckgo import DuckDuckGo
-
-# # Load the fine-tuned ICD-9 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
-
-# # 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',
-#     '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 ICD-9 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)
-    
-    # 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("Medical Diagnosis Assistant")
 # option = st.selectbox(
 #     "Choose Diagnosis Method",
@@ -310,14 +497,14 @@
 #         else:
 #             st.error("Please enter a medical summary.")
 
-# Medical Image Analysis
+# # Medical Image Analysis
 # elif option == "Medical Image Analysis":
 #     if "GOOGLE_API_KEY" not in st.session_state:
 #         st.warning("Please enter your Google API Key in the sidebar to continue")
 #     else:
 #         medical_agent = Agent(
 #             model=Gemini(
-#                 api_key=st.session_state.GOOGLE_API_KEY,
+#                 api_key=st.session_state["GOOGLE_API_KEY"],
 #                 id="gemini-2.0-flash-exp"
 #             ),
 #             tools=[DuckDuckGo()],
@@ -326,81 +513,76 @@
 
 #         query = """
 #         You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
-
 #         ### 1. Image Type & Region
 #         - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
 #         - Identify the patient's anatomical region and positioning
 #         - Comment on image quality and technical adequacy
-
 #         ### 2. Key Findings
 #         - List primary observations systematically
 #         - Note any abnormalities in the patient's imaging with precise descriptions
 #         - Include measurements and densities where relevant
 #         - Describe location, size, shape, and characteristics
 #         - Rate severity: Normal/Mild/Moderate/Severe
-
 #         ### 3. Diagnostic Assessment
 #         - Provide primary diagnosis with confidence level
 #         - List differential diagnoses in order of likelihood
 #         - Support each diagnosis with observed evidence from the patient's imaging
-        # - Note any critical or urgent findings
-
-        # ### 4. Patient-Friendly Explanation
-        # - Explain the findings in simple, clear language that the patient can understand
-        # - Avoid medical jargon or provide clear definitions
-        # - Include visual analogies if helpful
-        # - Address common patient concerns related to these findings
-
-        # ### 5. Research Context
-        # - Use the DuckDuckGo search tool to find recent medical literature about similar cases
-        # - Provide a list of relevant medical links
-        # - Include key references to support your analysis
-        # """
-
-        # upload_container = st.container()
-        # image_container = st.container()
-        # analysis_container = st.container()
-
-        # with upload_container:
-        #     uploaded_file = st.file_uploader(
-        #         "Upload Medical Image",
-        #         type=["jpg", "jpeg", "png", "dicom"],
-        #         help="Supported formats: JPG, JPEG, PNG, DICOM"
-        #     )
-
-        # if uploaded_file is not None:
-        #     with image_container:
-        #         col1, col2, col3 = st.columns([1, 2, 1])
-        #         with col2:
-        #             image = Image.open(uploaded_file)
-        #             width, height = image.size
-        #             aspect_ratio = width / height
-        #             new_width = 500
-        #             new_height = int(new_width / aspect_ratio)
-        #             resized_image = image.resize((new_width, new_height))
+#         - Note any critical or urgent findings
+#         ### 4. Patient-Friendly Explanation
+#         - Explain the findings in simple, clear language that the patient can understand
+#         - Avoid medical jargon or provide clear definitions
+#         - Include visual analogies if helpful
+#         - Address common patient concerns related to these findings
+#         ### 5. Research Context
+#         - Use the DuckDuckGo search tool to find recent medical literature about similar cases
+#         - Provide a list of relevant medical links
+#         - Include key references to support your analysis
+#         """
+
+#         upload_container = st.container()
+#         image_container = st.container()
+#         analysis_container = st.container()
+
+#         with upload_container:
+#             uploaded_file = st.file_uploader(
+#                 "Upload Medical Image",
+#                 type=["jpg", "jpeg", "png", "dicom"],
+#                 help="Supported formats: JPG, JPEG, PNG, DICOM"
+#             )
+
+#         if uploaded_file is not None:
+#             with image_container:
+#                 col1, col2, col3 = st.columns([1, 2, 1])
+#                 with col2:
+#                     image = Image.open(uploaded_file)
+#                     width, height = image.size
+#                     aspect_ratio = width / height
+#                     new_width = 500
+#                     new_height = int(new_width / aspect_ratio)
+#                     resized_image = image.resize((new_width, new_height))
                     
-        #             st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
+#                     st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
                     
-        #             analyze_button = st.button("🔍 Analyze Image")
+#                     analyze_button = st.button("🔍 Analyze Image")
 
-        #     with analysis_container:
-        #         if analyze_button:
-        #             image_path = "temp_medical_image.png"
-        #             with open(image_path, "wb") as f:
-        #                 f.write(uploaded_file.getbuffer())
+#             with analysis_container:
+#                 if analyze_button:
+#                     image_path = "temp_medical_image.png"
+#                     with open(image_path, "wb") as f:
+#                         f.write(uploaded_file.getbuffer())
                     
-        #             with st.spinner("🔄 Analyzing image... Please wait."):
-        #                 try:
-        #                     response = medical_agent.run(query, images=[image_path])
-        #                     st.markdown("### 📋 Analysis Results")
-        #                     st.markdown(response.content)
-        #                 except Exception as e:
-        #                     st.error(f"Analysis error: {e}")
-        #                 finally:
-        #                     if os.path.exists(image_path):
-        #                         os.remove(image_path)
-        # else:
-        #     st.info("👆 Please upload a medical image to begin analysis")
+#                     with st.spinner("🔄 Analyzing image... Please wait."):
+#                         try:
+#                             response = medical_agent.run(query, images=[image_path])
+#                             st.markdown("### 📋 Analysis Results")
+#                             st.markdown(response.content)
+#                         except Exception as e:
+#                             st.error(f"Analysis error: {e}")
+#                         finally:
+#                             if os.path.exists(image_path):
+#                                 os.remove(image_path)
+#         else:
+#             st.info("👆 Please upload a medical image to begin analysis")
 
 import os
 import torch
@@ -412,16 +594,6 @@ from phi.agent import Agent
 from phi.model.google import Gemini
 from phi.tools.duckduckgo import DuckDuckGo
 
-# Sidebar for Google API Key input
-st.sidebar.title("Settings")
-st.sidebar.write("Enter your Google API Key below for the Medical Image Analysis feature.")
-api_key = st.sidebar.text_input("Google API Key", type="password")
-
-if api_key:
-    st.session_state["GOOGLE_API_KEY"] = api_key
-else:
-    st.session_state.pop("GOOGLE_API_KEY", None)
-
 # Load the fine-tuned ICD-9 model and tokenizer
 model_path = "./clinical_longformer"
 tokenizer = LongformerTokenizer.from_pretrained(model_path)
@@ -429,7 +601,7 @@ 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 = pd.read_csv("D_ICD_DIAGNOSES.csv")
 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
 
@@ -452,7 +624,6 @@ def predict_icd9(texts, tokenizer, model, threshold=0.5):
         max_length=512,
         return_tensors="pt"
     )
-    
     with torch.no_grad():
         outputs = model(
             input_ids=inputs["input_ids"],
@@ -474,6 +645,9 @@ def predict_icd9(texts, tokenizer, model, threshold=0.5):
     
     return predictions_with_desc
 
+# Define the API key directly in the code
+GOOGLE_API_KEY = "AIzaSyA24A6egT3L0NAKkkw9QHjfoizp7cJUTaA"
+
 # Streamlit UI
 st.title("Medical Diagnosis Assistant")
 option = st.selectbox(
@@ -499,90 +673,86 @@ if option == "ICD-9 Code Prediction":
 
 # Medical Image Analysis
 elif option == "Medical Image Analysis":
-    if "GOOGLE_API_KEY" not in st.session_state:
-        st.warning("Please enter your Google API Key in the sidebar to continue")
-    else:
-        medical_agent = Agent(
-            model=Gemini(
-                api_key=st.session_state["GOOGLE_API_KEY"],
-                id="gemini-2.0-flash-exp"
-            ),
-            tools=[DuckDuckGo()],
-            markdown=True
-        )
-
-        query = """
-        You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
-        ### 1. Image Type & Region
-        - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
-        - Identify the patient's anatomical region and positioning
-        - Comment on image quality and technical adequacy
-        ### 2. Key Findings
-        - List primary observations systematically
-        - Note any abnormalities in the patient's imaging with precise descriptions
-        - Include measurements and densities where relevant
-        - Describe location, size, shape, and characteristics
-        - Rate severity: Normal/Mild/Moderate/Severe
-        ### 3. Diagnostic Assessment
-        - Provide primary diagnosis with confidence level
-        - List differential diagnoses in order of likelihood
-        - Support each diagnosis with observed evidence from the patient's imaging
-        - Note any critical or urgent findings
-        ### 4. Patient-Friendly Explanation
-        - Explain the findings in simple, clear language that the patient can understand
-        - Avoid medical jargon or provide clear definitions
-        - Include visual analogies if helpful
-        - Address common patient concerns related to these findings
-        ### 5. Research Context
-        - Use the DuckDuckGo search tool to find recent medical literature about similar cases
-        - Provide a list of relevant medical links
-        - Include key references to support your analysis
-        """
-
-        upload_container = st.container()
-        image_container = st.container()
-        analysis_container = st.container()
-
-        with upload_container:
-            uploaded_file = st.file_uploader(
-                "Upload Medical Image",
-                type=["jpg", "jpeg", "png", "dicom"],
-                help="Supported formats: JPG, JPEG, PNG, DICOM"
-            )
-
-        if uploaded_file is not None:
-            with image_container:
-                col1, col2, col3 = st.columns([1, 2, 1])
-                with col2:
-                    image = Image.open(uploaded_file)
-                    width, height = image.size
-                    aspect_ratio = width / height
-                    new_width = 500
-                    new_height = int(new_width / aspect_ratio)
-                    resized_image = image.resize((new_width, new_height))
-                    
-                    st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
-                    
-                    analyze_button = st.button("🔍 Analyze Image")
+    medical_agent = Agent(
+        model=Gemini(
+            api_key=GOOGLE_API_KEY,
+            id="gemini-2.0-flash-exp"
+        ),
+        tools=[DuckDuckGo()],
+        markdown=True
+    )
 
-            with analysis_container:
-                if analyze_button:
-                    image_path = "temp_medical_image.png"
-                    with open(image_path, "wb") as f:
-                        f.write(uploaded_file.getbuffer())
-                    
-                    with st.spinner("🔄 Analyzing image... Please wait."):
-                        try:
-                            response = medical_agent.run(query, images=[image_path])
-                            st.markdown("### 📋 Analysis Results")
-                            st.markdown(response.content)
-                        except Exception as e:
-                            st.error(f"Analysis error: {e}")
-                        finally:
-                            if os.path.exists(image_path):
-                                os.remove(image_path)
-        else:
-            st.info("👆 Please upload a medical image to begin analysis")
+    query = """
+    You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the patient's medical image and structure your response as follows:
+    ### 1. Image Type & Region
+    - Specify imaging modality (X-ray/MRI/CT/Ultrasound/etc.)
+    - Identify the patient's anatomical region and positioning
+    - Comment on image quality and technical adequacy
+    ### 2. Key Findings
+    - List primary observations systematically
+    - Note any abnormalities in the patient's imaging with precise descriptions
+    - Include measurements and densities where relevant
+    - Describe location, size, shape, and characteristics
+    - Rate severity: Normal/Mild/Moderate/Severe
+    ### 3. Diagnostic Assessment
+    - Provide primary diagnosis with confidence level
+    - List differential diagnoses in order of likelihood
+    - Support each diagnosis with observed evidence from the patient's imaging
+    - Note any critical or urgent findings
+    ### 4. Patient-Friendly Explanation
+    - Explain the findings in simple, clear language that the patient can understand
+    - Avoid medical jargon or provide clear definitions
+    - Include visual analogies if helpful
+    - Address common patient concerns related to these findings
+    ### 5. Research Context
+    - Use the DuckDuckGo search tool to find recent medical literature about similar cases
+    - Provide a list of relevant medical links
+    - Include key references to support your analysis
+    """
+
+    upload_container = st.container()
+    image_container = st.container()
+    analysis_container = st.container()
+
+    with upload_container:
+        uploaded_file = st.file_uploader(
+            "Upload Medical Image",
+            type=["jpg", "jpeg", "png", "dicom"],
+            help="Supported formats: JPG, JPEG, PNG, DICOM"
+        )
 
+    if uploaded_file is not None:
+        with image_container:
+            col1, col2, col3 = st.columns([1, 2, 1])
+            with col2:
+                image = Image.open(uploaded_file)
+                width, height = image.size
+                aspect_ratio = width / height
+                new_width = 500
+                new_height = int(new_width / aspect_ratio)
+                resized_image = image.resize((new_width, new_height))
+                
+                st.image(resized_image, caption="Uploaded Medical Image", use_container_width=True)
+                
+                analyze_button = st.button("🔍 Analyze Image")
+
+        with analysis_container:
+            if analyze_button:
+                image_path = "temp_medical_image.png"
+                with open(image_path, "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                
+                with st.spinner("🔄 Analyzing image... Please wait."):
+                    try:
+                        response = medical_agent.run(query, images=[image_path])
+                        st.markdown("### 📋 Analysis Results")
+                        st.markdown(response.content)
+                    except Exception as e:
+                        st.error(f"Analysis error: {e}")
+                    finally:
+                        if os.path.exists(image_path):
+                            os.remove(image_path)
+    else:
+        st.info("👆 Please upload a medical image to begin analysis")