Add PIE-Med Demo

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.csv filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+streamlit_results/

README.md

@@ -11,4 +11,82 @@ license: cc-by-sa-4.0
 short_description: '🩺 PIE-Med: Predicting, Interpreting and Explaining Medical'
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🩺 *PIE*-Med: *P*redicting, *I*nterpreting and *E*xplaining Medical Recommendations 
+
+Welcome to the repository for **PIE-Med**, a cutting-edge system designed to enhance medical decision-making through the integration of Graph Neural Networks (GNNs), Explainable AI (XAI) techniques, and Large Language Models (LLMs).
+
+## 🎥 Demo (or GIF)
+[Watch our demo](https://drive.google.com/file/d/1e9VXslnBzOOp5QHh4GTrT-La1PdKxhzS/preview) to see PIE-Med in action and learn how it can transform healthcare recommendations!
+
+## 📊 Data Source
+We use the **[MIMIC-III](https://mimic.physionet.org/)** dataset, a freely accessible critical care database containing de-identified health information, including vital signs, laboratory test results, medications, and more. You can find more details about the dataset here:
+
+## 🛠 Technologies Used
+- **Python**: Core programming language
+- **Pandas**: Data manipulation and analysis ([Pandas Documentation](https://pandas.pydata.org/))
+- **PyHealth**: Medical data preprocessing ([PyHealth Documentation](https://pyhealth.readthedocs.io/en/latest/))
+- **PyTorch Geometric**: Building and training GNNs ([PyTorch Geometric Documentation](https://pytorch-geometric.readthedocs.io/en/latest/))
+- **Integrated Gradients & GNNExplainer**: Interpretability techniques ([PyTorch Geometric Documentation](https://pytorch-geometric.readthedocs.io/en/latest/))
+- **Streamlit**: User interface development ([Streamlit Documentation](https://streamlit.io/))
+- **Py AutoGen Multi-Agent Conversation Framework**: Multi-agent collaboration and explanation ([Py AutoGen Documentation](https://microsoft.github.io/autogen/))
+
+The PIE-Med system's computational requirements depend on the configuration used. For resource-limited environments, the light configuration with an Intel i7 CPU and 16GB RAM offers a basic but functional setup, suitable for testing on small datasets. However, more demanding tasks, such as working with larger datasets or leveraging advanced machine learning techniques (e.g., Graph Neural Networks), benefit from cloud setups like the complete configuration, which includes a GPU (NVIDIA Tesla T4). In resource-constrained contexts, optimizing models and reducing dataset size would be crucial to ensure feasible performance.
+
+## 🔬 Methodological Workflow
+PIE-Med follows a comprehensive Predict→Interpret→Explain (PIE) paradigm:
+
+1. **Prediction Phase**: We construct a heterogeneous patient graph from MIMIC-III data and apply GNNs to generate personalized medical recommendations.
+2. **Interpretation Phase**: Integrated Gradients and GNNExplainer techniques are used to provide insights into the GNN's decision-making process.
+3. **Explanation Phase**: A collaborative ensemble of LLM agents analyzes the model's outputs and generates comprehensive, understandable explanations.
+
+![image](https://github.com/picuslab/PIE-Med/blob/main/PIE-Med.png)
+
+## 🌟 Key Features
+- **Integration of GNNs and LLMs**: Combining structured machine learning with natural language processing for robust recommendations.
+- **Enhanced Interpretability**: Using XAI techniques to make the decision-making process transparent.
+- **Collaborative Explanation**: Multi-agent LLMs provide detailed and understandable recommendations.
+
+## 🚀 Getting Started
+Follow these steps to set up and run PIE-Med on your local machine:
+
+### Prerequisites
+Ensure you have the following installed:
+- Python 3.7+
+
+### Installation
+1. **Clone the repository**:
+    ```bash
+    git clone https://github.com/picuslab/PIE-Med.git
+    cd PIE-Med
+    ```
+
+2. **Create a virtual environment**:
+    ```bash
+    python -m venv venv
+    source venv/bin/activate  # On Windows use `venv\Scripts\activate`
+    ```
+
+3. **Install the required packages**:
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+### Running the Application
+1. **Run the Streamlit application**:
+    ```bash
+    streamlit run dashboard.py
+    ```
+
+   Open your web browser and go to `http://localhost:8501` to interact with the application.
+
+## 📈 Conclusions
+PIE-Med showcases the potential of combining GNNs, XAI, and LLMs to improve medical recommendations, enhancing both accuracy and interpretability. Our system effectively separates prediction from explanation, reducing biases and enhancing decision quality.
+
+## ⚖ Ethical considerations
+
+**PIE-Med** aims to support medical decision-making, but is not a substitute for professional medical advice. Users should confirm recommendations with authorised healthcare providers, as limitations of AI may affect accuracy. The system ensures transparency through interpretability techniques, but all results should be considered complementary to expert advice. **⚠️ Please note that the following repository is only a DEMO, with anonymised data used for illustrative purposes only**. 
+
+## 🙏 Acknowledgments
+We extend our gratitude to the creators of the MIMIC-III database, the developers of the Python libraries used, and our research team for their contributions to this project.
+
+👨‍💻 This project was developed by Antonio Romano, Giuseppe Riccio, Marco Postiglione and Vincenzo Moscato

chat_utils.py

@@ -0,0 +1,120 @@
+import streamlit as st
+import autogen
+
+from autogen import OpenAIWrapper, AssistantAgent, UserProxyAgent
+from typing import List
+
+
+def setting_config(model_name: str) -> tuple[List[dict], dict]:
+    if model_name.startswith("gpt"):
+        config_list = [
+            {
+                "model": model_name, 
+                "base_url": "https://api.openai.com/v1", 
+                "api_key": st.secrets.openai, 
+            }
+        ]
+    else:
+        config_list = [
+            {
+                "model": model_name, # "google/gemma-2-9b-it",
+                "base_url": "https://integrate.api.nvidia.com/v1", 
+                "api_key": st.secrets.nvidia, 
+                "max_tokens": 1000, 
+            }
+        ]
+
+    llm_config={
+        "timeout": 500, 
+        "seed": 42, 
+        "config_list": config_list, 
+        "temperature": 0.5, 
+    }
+
+    return config_list, llm_config
+
+
+class TrackableUserProxyAgent(UserProxyAgent):
+    t = 0
+    def _process_received_message(self, message, sender, silent):
+        global t
+        with st.chat_message(sender.name, avatar="streamlit_images/{}.png".format(self.t)):
+            st.write(f"**{message['name'].replace('_',' ')}**: {message['content']}")
+            self.t += 1
+            if self.t == 4:
+                self.t = 0
+        st.divider()
+        
+        return super()._process_received_message(message, sender, silent)
+
+
+def doctor_recruiter(prompt_recruiter_doctors: str, model_name: str) -> str:
+    config_list, llm_config = setting_config(model_name) 
+    client = OpenAIWrapper(api_key=config_list[0]['api_key'], config_list=config_list)
+    response = client.create(messages=[{"role": "user", "content": prompt_recruiter_doctors + "\nReturn ONLY JSON file (don't use Markdown tags like: ```json)."}], 
+                            temperature=0.3, 
+                            seed=42, 
+                            model=config_list[0]['model'])
+    text = client.extract_text_or_completion_object(response)
+
+    return text
+
+
+def doctor_discussion(doctor_description: str, prompt_internist_doctor: str, model_name: str) -> str:
+    config_list, llm_config = setting_config(model_name)
+    doctor = OpenAIWrapper(api_key=config_list[0]['api_key'], config_list=config_list)
+    response = doctor.create(messages=[
+                                        {"role": "assistant", 
+                                        "content": doctor_description},
+                                        {"role": "user", 
+                                        "content": prompt_internist_doctor + "\nDon't use Markdown tags."}
+                                        ], 
+                            temperature=0.3, 
+                            model=config_list[0]['model'])
+    text = doctor.extract_text_or_completion_object(response)
+
+    return text
+
+
+def multiagent_doctors(json_data: dict, model_name: str) -> List[AssistantAgent]:
+    config_list, llm_config = setting_config(model_name)
+    doc = []
+    for i in range(len(json_data['doctors'])):
+        doc.append(AssistantAgent(
+            name=json_data['doctors'][i]['role'].replace(" ", "_"), 
+            llm_config=llm_config, 
+            system_message="As a " + json_data['doctors'][i]['role'].replace(" ", "_") + ". Discuss with other medical experts in the team to help the INTERNIST DOCTOR make a final decision. Avoid postponing further examinations and repeating opinions given in the analysis, but explain in a logical and concise manner why you are making this final decision."))
+
+    return doc
+
+
+def care_discussion_start(doc: List[AssistantAgent], prompt_reunion: str, internist_sys_message: str, model_name: str) -> autogen.GroupChatManager:
+    config_list, llm_config = setting_config(model_name)
+    doc.append(TrackableUserProxyAgent(
+        name="internist_doctor", 
+        human_input_mode="NEVER", 
+        max_consecutive_auto_reply=1, 
+        is_termination_msg=lambda x: x.get("content", "").rstrip().endswith(("JUSTIFIABLE", "UNJUSTIFIABLE")), 
+        code_execution_config=False, 
+        llm_config=llm_config, 
+        system_message=internist_sys_message 
+    ))
+
+    groupchat = autogen.GroupChat(agents=doc, 
+                                messages=[], 
+                                max_round=(len(doc)+1), 
+                                speaker_selection_method="round_robin", 
+                                role_for_select_speaker_messages='user', 
+                                )
+
+    manager = autogen.GroupChatManager(groupchat=groupchat, 
+                                    llm_config=llm_config, 
+                                    max_consecutive_auto_reply=1 
+                                    )
+
+    doc[-1].initiate_chat(
+        manager,
+        message=prompt_reunion,
+    )
+
+    return manager

css/style.css

@@ -0,0 +1,34 @@
+<style>
+    div[data-testid="stToolbar"] {
+        visibility: hidden;
+        height: 0%;
+        position: fixed;
+    }
+
+    div[data-testid="stDecoration"] {
+        visibility: hidden;
+        height: 0%;
+        position: fixed;
+    }
+
+    div[data-testid="stStatusWidget"] {
+        visibility: hidden;
+        height: 0%;
+        position: fixed;
+    }
+
+    #MainMenu {
+        visibility: hidden;
+        height: 0%;
+    }
+
+    header {
+        visibility: hidden;
+        height: 0%;
+    }
+
+    footer {
+        visibility: hidden;
+        height: 0%;
+    }
+</style>

dashboard.py

@@ -0,0 +1,470 @@
+from chat_utils import *
+from model_utils import *
+
+import json
+import shutil
+from faker import Faker
+
+
+PATH_MED  = "model/medication_recommendation/best.ckpt"
+PATH_DIAG = "model/diagnosis_prediction/best.ckpt"
+#shutil.rmtree(".cache/", ignore_errors=True)
+
+
+def main():
+
+    # ---- SETTINGS PAGE ----
+    st.set_page_config(page_title="PIE-Med - Dashboard", page_icon="🩺", layout="wide")
+
+    # with open('css/style.css') as f:
+    #     hide_streamlit_style = f.read()
+    # st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
+    # ---- SESSION STATE ----
+    if 'patient' not in st.session_state:
+        st.session_state.patient = None
+    if 'name' not in st.session_state:
+        st.session_state.name = None
+    if 'lastname' not in st.session_state:
+        st.session_state.lastname = None
+    if 'gender_sign' not in st.session_state:
+        st.session_state.gender_sign = None
+
+    # ---- SIDE BAR ----
+    # st.sidebar.image(".\streamlit_images\logo_icon.png")
+    # st.sidebar.divider()
+
+    # ---- MAIN PAGE ----
+    st.title(":rainbow[PIE-Med]")
+    st.markdown("Welcome to PIE-Med 🩺!")
+
+    desc = st.empty()
+    desc1 = st.empty()
+    desc.caption("**PIE-Med** 🩺, a cutting-edge system designed to enhance medical decision-making through \
+                    the integration of **Graph Neural Networks (GNNs)** ⚙️, **eXplainable AI (XAI)** ❓ techniques, \
+                    and **Large Language Models (LLMs)** 🧠.")
+    desc1.caption("**⏳ WAIT MINUTES FOR THE LOADING OF THE MODELS AND THE DATASET**")
+
+    model_med_ig, model_med_gnn, model_diag_ig, model_diag_gnn, \
+        dataset, mimic3sample_med, mimic3sample_diag = load_gnn()
+    checkpoint_MED = torch.load(PATH_MED)
+    checkpoint_DIAG = torch.load(PATH_DIAG)
+
+    desc1.empty()
+
+    fake = Faker()
+
+    selected_patient = None
+    if selected_patient is None:
+        placeholder2 = st.empty()
+        with placeholder2.expander("⚠️ **Before using the framework, read the disclaimer for the use of Framework**"):
+            disclaimer = f"""
+
+            The use of our Healthcare framework based on MIMIC III (https://physionet.org/content/mimiciii/1.4/) is subject to the terms and warnings as follows:
+
+            **Research and Decision Support Purpose:** Our framework has been developed primarily for research and decision support in the healthcare context. The information and recommendations generated should not replace the professional judgment of qualified healthcare practitioners but may be utilized as support for the final decision by the doctor or the directly involved party.
+
+            **Data Origin:** The processed healthcare data originates from the MIMIC III database and undergoes enrichment and modeling through the application of Heterogeneous Graph Neural Network. It is important to note that the original data may contain variations and limitations, and the accuracy of the processed information depends on the quality of the input data.
+
+            **Medical Recommendations:** The drug and diagnosis recommendations generated by the framework are hypothetical and based on Graph Neural Network learning models. These should not be considered definitive prescriptions, and the final decision regarding patient treatment should be made by a qualified medical professional.
+
+            **Human Readable Explanations:** The embedded explainability system in the framework utilizes graph explainability models and Large Language Models (LLM) to generate understandable explanations for end-users, such as physicians. However, these explanations are interpretations of the model results and may not fully reflect the complexity of medical reasoning.
+
+            **Framework Limitations:** Our framework has intrinsic limitations, including those related to the quality of input data, the characteristics of the machine learning model, and the dynamics of the healthcare context. Users are encouraged to exercise caution in interpreting the provided information.
+
+            **User Responsibility:** Users accessing and utilizing our framework are responsible for the accurate interpretation of the provided information and for making appropriate decisions based on their clinical judgment. The creators assume no responsibility for any consequences arising from improper use or misinterpretation of the information generated by the framework.
+
+            By using our healthcare data processing framework, the user agrees to comply with these conditions. The continuous evolution of the fields of medicine and technology may necessitate periodic updates to this disclaimer.
+
+            """
+
+            st.subheader("Disclaimer")
+            st.info(disclaimer)
+            agree = st.checkbox("I accept and have read the disclaimer!")
+            placeholder1 = st.empty()
+            placeholder1.warning("You must accept the disclaimer to use the framework!", icon="⚠️")
+
+            if not(agree):
+                st.stop()
+
+            placeholder1.empty()
+            placeholder2.info("You can now use the framework! 🎉 Please select the task and select a patient! 🩺")
+            task = st.sidebar.selectbox(label='Select __task__: ', index=None, placeholder="Select type of task", options=['medications', 'diagnosis'])
+
+            if task is None:
+                st.stop()
+            elif task == "medications":
+                mimic3sample = mimic3sample_med
+            elif task == "diagnosis":
+                mimic3sample = mimic3sample_diag
+
+            mimic_df = pd.DataFrame(mimic3sample.samples)
+
+            selected_patient = st.sidebar.selectbox(label='Select __patient__ n°: ', index=None, placeholder="Select a patient", options=mimic_df['patient_id'].unique())
+            while selected_patient is None:
+                st.stop()
+
+    desc.empty()
+    placeholder2.empty()
+
+    patient_dict = dataset.patients
+    patient_info = patient_dict[selected_patient]
+    gender = patient_info.gender
+
+    if selected_patient != st.session_state.patient:
+        if gender == "M":
+            first_name = fake.first_name_male()
+            last_name = fake.last_name_male()
+            gender_sign = "male_sign"
+        elif gender == "F":
+            first_name = fake.first_name_female()
+            last_name = fake.last_name_female()
+            gender_sign = "female_sign"
+        else:
+            first_name = "Name"
+            last_name = "Unknown"
+
+        st.session_state.patient = selected_patient
+        st.session_state.name = ":blue[" + first_name + "]"
+        st.session_state.lastname = last_name
+        st.session_state.gender_sign = gender_sign
+
+    patient = st.session_state.patient
+    name = st.session_state.name
+    lastname = st.session_state.lastname
+    gender_sign = st.session_state.gender_sign
+
+    mimic_df_patient = mimic_df[mimic_df['patient_id'] == selected_patient] # select all the rows with the selected patient
+
+    for i in range(len(mimic_df_patient)):
+        if i == len(mimic_df_patient) - 1:
+            last_visit = mimic_df_patient.iloc[[i]]
+
+    # ---- Patient info ----
+    # st.subheader(":blue[DASHBOARD OF] ")
+    st.warning("🚨 **NOTE** 🚨: The patient's name, shown below, was randomly generated for demonstration purposes.")
+    st.title("{} {} :{}:".format(name, lastname, gender_sign))
+    st.caption("Patient n°: {}  -  Gender: {}  -  Ethnicity: {}".format(patient, patient_info.gender, patient_info.ethnicity))
+
+    l1, r1 = st.columns([0.44, 0.56])
+
+    with l1:
+        st.subheader("📋 Medical history")
+        # st.caption("The following table shows the *complete* medical history of the patient n°: **{}**.".format(patient))
+
+        visit = st.selectbox(label='🏥 __Hospital admission__ n°: ', options=mimic_df_patient['visit_id'].unique())
+        if visit:
+            mimic_df_patient_visit = mimic_df_patient[mimic_df_patient['visit_id'] == visit] # select all the rows with the selected visit
+            if task == "medications":
+                mimic_df_patient_visit_filtered = mimic_df_patient_visit.drop(columns=['visit_id', 'patient_id', 'drugs_hist'])
+            elif task == "diagnosis":
+                mimic_df_patient_visit_filtered = mimic_df_patient_visit.drop(columns=['visit_id', 'patient_id'])
+
+            atc = InnerMap.load("ATC")
+            icd9 = InnerMap.load("ICD9CM")
+            icd9_proc = InnerMap.load("ICD9PROC")
+
+            for column in mimic_df_patient_visit_filtered.columns:
+                with st.expander("{}".format(column)):
+                    try:
+                        if column == "medications":
+                            if task == "medications":
+                                med_history = [[med, atc.lookup(med)] for med in mimic_df_patient_visit_filtered[column].explode() if med]
+                            elif task == "diagnosis":
+                                med_history = [[med, atc.lookup(med)] for med in (mimic_df_patient_visit_filtered[column].explode()).explode() if med]
+                            st.dataframe(med_history, hide_index=True, column_config={"0": "ATC", "1": "Description"})
+                        elif column == "diagnosis":
+                            if task == "medications":
+                                col_history = [[idx, icd9.lookup(idx)] for idx in (mimic_df_patient_visit_filtered[column].explode()).explode() if idx]
+                            elif task == "diagnosis":
+                                col_history = [[idx+'0', icd9.lookup(idx+'0')] if idx.startswith('E') else [idx, icd9.lookup(idx)] for idx in mimic_df_patient_visit_filtered[column].explode() if idx]
+                            st.dataframe(col_history, hide_index=True, column_config={"0": "ICD9", "1": "Description"})
+                        elif column == "symptoms":
+                            col_history = [[idx, icd9.lookup(idx)] for idx in (mimic_df_patient_visit_filtered[column].explode()).explode() if idx]
+                            st.dataframe(col_history, hide_index=True, column_config={"0": "ICD9", "1": "Description"})
+                        elif column == "procedures":
+                            col_history = [[idx, icd9_proc.lookup(idx)] for idx in (mimic_df_patient_visit_filtered[column].explode()).explode() if idx]
+                            st.dataframe(col_history, hide_index=True, column_config={"0": "ICD9", "1": "Description"})
+                    except:
+                        st.write("No data available for this column.")
+
+        st.subheader(f"🧾 Recommended _{task}_")
+        st.caption(f"""The following {task} are recommended for the patient during the **hospital admission n°: \
+                   {format(last_visit['visit_id'].item())}**. \n The recommendations are based on the \
+                    output probabilities generated by the **GNN (_Graph Neural Network_)** model.""")
+
+        if task == "medications":
+            model_med_ig.load_state_dict(checkpoint_MED)
+            model_med_gnn.load_state_dict(checkpoint_MED)
+            model = model_med_ig
+        elif task == "diagnosis":
+            model_diag_ig.load_state_dict(checkpoint_DIAG)
+            model_diag_gnn.load_state_dict(checkpoint_DIAG)
+            model = model_diag_ig
+
+        # ---- Output model ----
+        model.eval()
+        output = model(last_visit['patient_id'],
+                    last_visit['visit_id'],
+                    last_visit['diagnosis'],
+                    last_visit['procedures'],
+                    last_visit['symptoms'],
+                    last_visit['medications'])
+
+        list_output, list_indices = get_list_output(output['y_prob'], last_visit, task, mimic3sample)
+        list_output = [[idx, item] for idx, item in zip(*list_indices, *list_output) if item]
+        st.dataframe(list_output, column_config={"0": "ID", "1": f"Recommended {task}"}, height=None, width=None)
+
+    with r1:
+        st.subheader(f"""🗣 *Why* did the model recommend these {task}?""")
+        r1l1, r1c1, r1r1 = st.columns(3)
+        with r1l1:
+            visualization = st.radio("Visualization", options=["Explainable", "Interpretable"], horizontal=True)
+        with r1c1:
+            algorithm = st.radio("Algorithm", options=["IG", "GNNExplainer"], horizontal=True)
+        with r1r1:
+            threshold = st.slider("Threshold", min_value=10, max_value=50, value=15, step=5, format=None, key=None)
+
+        if task == "medications" and algorithm == "IG":
+            model = model_med_ig
+        elif task == "medications" and algorithm == "GNNExplainer":
+            model = model_med_gnn
+        elif task == "diagnosis" and algorithm == "IG":
+            model = model_diag_ig
+        elif task == "diagnosis" and algorithm == "GNNExplainer":
+            model = model_diag_gnn
+
+        st.caption(f"""The graph shown as follows provides an interpretation of the model's decision making process on the recommended \
+                    *{task}* for the patient during the **hospital admission n°: {format(last_visit['visit_id'].item())}**. \
+                    \n\n The interpretability is based on the **{algorithm} (_{task}_)** algorithm.""")
+        options = [item[1] for item in list_output if item]
+        selected_label = st.selectbox(f'Select the {task} to explain', index=None, 
+                                        placeholder=f"Choice a {task} from Recommended {task} ranking to explain", 
+                                        options=options)
+
+        if selected_label is None:
+            st.stop()
+
+        selected_idx = [item[0] for item in list_output if item[1] == selected_label]
+
+        st.caption("Legend of the graph:")
+        col1, col2, col3, col4, col5, col6, col7, col8 = st.columns([0.1, 0.3, 0.1, 0.3, 0.1, 0.3, 0.1, 0.3])
+
+        with col1:
+            st.markdown(
+                """
+                <style>
+                #square1 {
+                    width: 20px;
+                    height: 20px;
+                    background: #20b2aa;
+                    border-radius: 3px;
+                }
+                </style>
+                <div id="square1"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+            st.markdown(
+                """
+                <style>
+                #square2 {
+                    width: 20px;
+                    height: 20px;
+                    background: #fa8072;
+                    border-radius: 3px;
+                    margin-top: 20px;
+                }
+                </style>
+                <div id="square2"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+        with col2:
+            st.caption("Patient")
+
+            st.caption("Visit")
+
+        with col3:
+            st.markdown(
+                """
+                <style>
+                #square3 {
+                    width: 20px;
+                    height: 20px;
+                    background: #cd853f;
+                    border-radius: 3px;
+                }
+                </style>
+                <div id="square3"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+            st.markdown(
+                """
+                <style>
+                #square4 {
+                    width: 20px;
+                    height: 20px;
+                    background: #da70d6;
+                    border-radius: 3px;
+                    margin-top: 20px;
+                }
+                </style>
+                <div id="square4"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+        with col4:
+            st.caption("Diagnosis")
+            st.caption("Procedures")
+
+        with col5:
+            st.markdown(
+                """
+                <style>
+                #square5 {
+                    width: 20px;
+                    height: 20px;
+                    background: #98fb98;
+                    border-radius: 3px;
+                }
+                </style>
+                <div id="square5"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+        with col6:
+            st.caption("Symptoms")
+
+        with col7:
+            st.markdown(
+                """
+                <style>
+                #square6 {
+                    width: 20px;
+                    height: 20px;
+                    background: #87ceeb;
+                    border-radius: 3px;
+                }
+                </style>
+                <div id="square6"></div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+        with col8:
+            st.caption("Medications")
+
+        explain_sample = {}
+        for visit_sample in mimic3sample.samples:
+            if visit_sample['patient_id'] == patient and visit_sample['visit_id'] == last_visit['visit_id'].item():
+                if visit_sample.get('drugs_hist') != None:
+                    del visit_sample['drugs_hist']
+                explain_sample['test'] = visit_sample
+
+        model.eval()
+        explain_dataset = SampleEHRDataset(list(explain_sample.values()), code_vocs="ATC")
+        explainability(model, explain_dataset, selected_idx[0], visualization, algorithm, task, threshold)
+
+
+    ####################### CARE AI module ##################################
+    st.header('🩺🧠 Medical Agents Evaluation')
+    st.caption("The section shown as follows is dedicated to the Explainability module, which is responsible for generating the analysis of the doctors' proposals and the collaborative discussion between the medical team members for the final decision on the patient's treatment.")
+
+    model_name = st.selectbox("Select the LLM model", options=["gpt-4o-mini", "gpt-3.5-turbo", "google/gemma-2-9b-it", "meta/llama3-8b-instruct", "mistralai/mistral-7b-instruct-v0.2"])
+
+    explanation = st.button("Generate explanation")
+    if not(explanation):
+        st.stop()
+
+    col1, col2 = st.columns([0.5, 0.6], gap="large")
+
+    with col1:
+        with open("streamlit_results/medical_scenario.txt", "r") as f:
+            medical_scenario = f.read()
+        st.subheader("📄 Medical Scenario")
+        st.caption(f"The scenario shown as follows for the patient in the **hospital admission n°: {format(last_visit['visit_id'].item())}** is provided by the medical team.")
+        st.markdown('###')
+        with st.expander("👁️ Read the medical scenario", expanded=True):
+            container = st.container(height=1145)
+            container.write(medical_scenario)
+
+    with col2:
+        st.subheader("👨‍⚕️🔎 Doctor Recruiter")
+        st.caption("The doctor recruiter is responsible for recruiting the medical team to help the internist doctor make a final decision on the patient's during the collaborative discussion.")
+        with st.status("Recruiting doctor...", expanded=False) as status:
+            with open("streamlit_results/prompt_recruiter_doctors.txt", "r") as f:
+                prompt_recruiter_doctors = f.read()
+            text = doctor_recruiter(prompt_recruiter_doctors, model_name)
+            if model_name == "meta/llama3-8b-instruct":
+                text[0] = text[0].split("Here is the JSON file:\n\n")[1]
+            json_data = json.loads(str(text[0]))
+            with open("streamlit_results/recruited_doctors.json", "w") as f:
+                json.dump(text[0], f, indent=4)
+
+            for i, doctor in enumerate(json_data['doctors']):
+                role = f"""**🥼 {doctor['role'].replace("_", " ")}**"""
+                st.markdown(role)
+                st.write(doctor['description'])
+                if i != len(json_data['doctors'])-1:
+                    st.divider()
+
+            status.update(label="Doctor recruited!", state="complete", expanded=True)
+        st.button('Rerun')
+
+    st.subheader("Analysis Proposition")
+    with st.spinner("Doctors are thinking..."):
+        with open("streamlit_results/prompt_internist_doctor.txt", "r") as f:
+            prompt_internist_doctor = f.read()
+
+        prompt_reunion = f"""Based on your assessment and the medical team's recommendations regarding {task} during the patient visit:\n"""
+        prompt_reunion += f"""Confront with your medical colleagues, highlighting relevant aspects related to the patient's condition and the {task}. Underline the crucial elements that influence your decision on its justification or unjustification in 30 words.\n"""
+        prompt_reunion += f"""\nAnalysis of doctors' proposals\n\n"""
+
+        for i in range(len(json_data['doctors'])):
+            with st.status(f"The 👨‍⚕️ {json_data['doctors'][i]['role'].replace('_', ' ')} is analysing ...", expanded=False) as status_doc:
+                with st.chat_message(name="user", avatar="streamlit_images/{}.png".format(i)):
+                    analysis = """"""
+                    analysis += f"""**Doctor**: {json_data['doctors'][i]['role'].replace(" ", "_")}\n\n"""
+                    text = doctor_discussion(json_data['doctors'][i]['role'], prompt_internist_doctor, model_name)
+                    analysis += "**Analysis**: " + text[0]
+                    st.markdown(f"**Analysis**: {text[0]}")
+                    status_doc.update(label="The 👨‍⚕️ {} analysed!".format(json_data['doctors'][i]['role'].replace('_', ' ')), state="complete", expanded=True)
+                    prompt_reunion += f"""{analysis}"""
+                    prompt_reunion += f"\n--------------------------------------------------\n\n"
+
+    image, text = st.columns([0.2, 0.8])
+    with image: 
+        st.image("streamlit_images/collaborative.png")
+    with text:
+        st.subheader('Discussion')
+
+    st.caption("The discussion shown as follows is based on the **Large Language Model** (LLM) **chosen**. The LLM is responsible for generating the discussion between the medical team members for the final decision on the patient's treatment.")
+    with st.spinner("Doctors are discussing..."):
+        internist_sys_message = f"""As an INTERNIST DOCTOR, you have the task of globally evaluating and managing the patient's health and pathology.\n"""
+        internist_sys_message += f"""In the light of the entire discussion, you must provide a final schematic report to the doctor based on the recommendation and the doctors' opinions."""
+
+        doc = multiagent_doctors(json_data, model_name)
+        manager = care_discussion_start(doc, prompt_reunion, internist_sys_message, model_name)
+
+        with st.chat_message(name="user", avatar="streamlit_images/internist.png"):
+            internist = list(manager.chat_messages.values())
+            internist_opinion = internist[0][6]['content']
+            st.write(f"**{internist[0][6]['name'].replace('_',' ')}**: {internist_opinion}")
+
+        # Add a download button:
+        st.download_button(
+            label="Download PDF", 
+            data=gen_pdf(patient, name, lastname, last_visit['visit_id'].item(), list_output, medical_scenario, internist_opinion), 
+            file_name=f"Medical_Report_Patient_{patient}.pdf", 
+            mime="application/pdf", 
+        )
+
+
+if __name__ == "__main__":
+    main()

model/diagnosis_prediction/best.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5be49c0eb4f53152c76d86203a492b5baea5f81081d710fd2671e051bae2db3e
+size 6773953

model/diagnosis_prediction/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf971bb5a3ac59951dc94d3fb4e82375d7f5f715988fc700d9644d9f1016359b
+size 6773953

model/medication_recommendation/best.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19ed90e26be751ea72cffd0da6562b799f8ec63e2e34721414aa333b4929992b
+size 7673537

model/medication_recommendation/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:602ccd50078f050c87fc20fb0caecbc836f03dd1e74e7d9586dcd4ff7edf0230
+size 7673537

model_utils.py

@@ -0,0 +1,212 @@
+import streamlit as st
+import streamlit.components.v1 as components
+import numpy as np
+import pandas as pd
+import torch
+
+from typing import Tuple, List
+from fpdf import FPDF
+
+from pyhealth.medcode import InnerMap
+from pyhealth.datasets import MIMIC3Dataset, SampleEHRDataset
+from pyhealth.tasks import medication_recommendation_mimic3_fn, diagnosis_prediction_mimic3_fn
+from pyhealth.models import GNN
+from pyhealth.explainer import HeteroGraphExplainer
+
+
+@st.cache_resource(hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
+def load_gnn() -> Tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module, torch.nn.Module,
+                        MIMIC3Dataset, SampleEHRDataset, SampleEHRDataset]:
+    dataset = MIMIC3Dataset(
+        root=st.secrets.s3,
+        tables=["DIAGNOSES_ICD","PROCEDURES_ICD","PRESCRIPTIONS","NOTEEVENTS_ICD"],
+        code_mapping={"NDC": ("ATC", {"target_kwargs": {"level": 4}})},
+    )
+
+    mimic3sample_med = dataset.set_task(task_fn=medication_recommendation_mimic3_fn)
+    mimic3sample_diag = dataset.set_task(task_fn=diagnosis_prediction_mimic3_fn)
+
+    model_med_ig = GNN(
+        dataset=mimic3sample_med,
+        convlayer="GraphConv",
+        feature_keys=["procedures", "diagnosis", "symptoms"],
+        label_key="medications",
+        k=0,
+        embedding_dim=128,
+        hidden_channels=128
+    )
+
+    model_med_gnn = GNN(
+        dataset=mimic3sample_med,
+        convlayer="GraphConv",
+        feature_keys=["procedures", "diagnosis", "symptoms"],
+        label_key="medications",
+        k=0,
+        embedding_dim=128,
+        hidden_channels=128
+    )
+
+    model_diag_ig = GNN(
+        dataset=mimic3sample_diag,
+        convlayer="GraphConv",
+        feature_keys=["procedures", "medications", "symptoms"],
+        label_key="diagnosis",
+        k=0,
+        embedding_dim=128,
+        hidden_channels=128
+    )
+
+    model_diag_gnn = GNN(
+        dataset=mimic3sample_diag,
+        convlayer="GraphConv",
+        feature_keys=["procedures", "medications", "symptoms"],
+        label_key="diagnosis",
+        k=0,
+        embedding_dim=128,
+        hidden_channels=128
+    )
+
+    return model_med_ig, model_med_gnn, model_diag_ig, model_diag_gnn, dataset, mimic3sample_med, mimic3sample_diag
+
+
+@st.cache_data(hash_funcs={torch.Tensor: lambda _: None})
+def get_list_output(y_prob: torch.Tensor, last_visit: pd.DataFrame, task: str, _mimic3sample: SampleEHRDataset, 
+                   top_k: int = 10) -> List[str]:
+    sorted_indices = []
+    for i in range(len(y_prob)):
+        top_indices = np.argsort(-y_prob[i, :])[:top_k]
+        sorted_indices.append(top_indices)
+
+    list_output = []
+
+    # get the list of all labels in the dataset
+    if task == "medications":
+        list_labels = _mimic3sample.get_all_tokens('medications')
+        atc = InnerMap.load("ATC")
+    elif task == "diagnosis":
+        list_labels = _mimic3sample.get_all_tokens('diagnosis')
+        icd9 = InnerMap.load("ICD9CM")
+
+    sorted_indices = list(sorted_indices)
+    # iterate over the top indexes for each sample in test_ds
+    for (i, sample), top in zip(last_visit.iterrows(), sorted_indices):
+        # create an empty list to store the recommended medications for this sample
+        sample_list_output = []
+
+        # iterate over the top indexes for this sample
+        for k in top:
+            # append the medication at the i-th index to the recommended medications list for this sample
+            if task == "medications":
+                sample_list_output.append(atc.lookup(list_labels[k]))
+            elif task == "diagnosis":
+                if list_labels[k].startswith("E"):
+                    list_labels[k] = list_labels[k] + "0"
+                sample_list_output.append(icd9.lookup(list_labels[k]))
+
+        # append the recommended medications for this sample to the recommended medications list
+        list_output.append(sample_list_output)
+
+    return list_output, sorted_indices
+
+
+def explainability(model: GNN, explain_dataset: SampleEHRDataset, selected_idx: int, 
+                   visualization: str, algorithm: str, task: str, threshold: int):
+    explainer = HeteroGraphExplainer(
+        algorithm=algorithm,
+        dataset=explain_dataset,
+        model=model,
+        label_key=task,
+        threshold_value=threshold,
+        top_k=threshold,
+        feat_size=128,
+        root="./streamlit_results/",
+    )
+
+    if task == "medications":
+        visit_drug = explainer.subgraph['visit', 'medication'].edge_index
+        visit_drug = visit_drug.T
+
+        n = 0
+        for vis_drug in visit_drug:
+            vis_drug = np.array(vis_drug)
+            if vis_drug[1] == selected_idx:
+                break
+            n += 1
+    elif task == "diagnosis":
+        visit_diag = explainer.subgraph['visit', 'diagnosis'].edge_index
+        visit_diag = visit_diag.T
+
+        n = 0
+        for vis_diag in visit_diag:
+            vis_diag = np.array(vis_diag)
+            if vis_diag[1] == selected_idx:
+                break
+            n += 1
+
+    explainer.explain(n=n)
+    if visualization == "Explainable":
+        explainer.explain_graph(k=0, human_readable=True, dashboard=True)
+    else:
+        explainer.explain_graph(k=0, human_readable=False, dashboard=True)
+
+    explainer.explain_results(n=n)
+    explainer.explain_results(n=n, doctor_type="Internist_Doctor")
+
+    HtmlFile = open("streamlit_results/explain_graph.html", 'r', encoding='utf-8')
+    source_code = HtmlFile.read() 
+    components.html(source_code, height=520)
+
+
+def gen_pdf(patient, name, lastname, visit, list_output, medical_scenario, internist_scenario):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.add_font("OpenSans", style="", fname="font/OpenSans.ttf")
+    pdf.add_font("OpenSans", style="B", fname="font/OpenSans-Bold.ttf")
+
+    # Title
+    pdf.set_font("OpenSans", 'B', 14)
+    pdf.cell(0, 10, 'Patient Medical Report', 0, 1, 'C', markdown=True)
+    pdf.ln(5)
+
+    # Patient Info
+    pdf.set_font("OpenSans", 'B', 10)
+    pdf.cell(0, 10, 'Patient Information', 0, 1, 'L', markdown=True)
+    pdf.set_font("OpenSans", '', 8)
+    pdf.cell(0, 3, f"Patient ID: **{patient}** - Name: **{name.split('[')[1].split(']')[0]}** Surname: **{lastname}** - Hospital admission n°: **{visit}**", 0, 1, 'L', markdown=True)
+    pdf.ln(5)
+
+    # Left column (Medical Scenario)
+    left_x = 10
+    right_x = 110
+    col_width = 90
+
+    # Right column (Recommendations)
+    pdf.set_xy(right_x, pdf.get_y())
+    pdf.set_font("OpenSans", 'B', 10)
+    pdf.cell(col_width - 20, 10, 'Recommendations', 0, 1, 'L')
+    pdf.set_xy(right_x, pdf.get_y())
+    pdf.set_font("OpenSans", '', 8)
+    for i, output in enumerate(list_output):
+        tensor_value = output[0].item()  # Convert tensor to number
+        recommendation = output[1]
+        pdf.set_xy(right_x, pdf.get_y())
+        pdf.cell(col_width - 20, 3, f"Medication {i+1}: {tensor_value}, {recommendation}", 0, 1, 'L')
+
+
+    # Medical Scenario
+    pdf.set_xy(left_x, pdf.get_y() - 40)
+    pdf.set_font("OpenSans", 'B', 10)
+    pdf.cell(col_width, 10, 'Medical Scenario', 0, 1, 'L', markdown=True)
+    pdf.set_xy(left_x, pdf.get_y())
+    pdf.set_font("OpenSans", '', 8)
+    pdf.multi_cell(col_width, 3, medical_scenario, 0, 'L', markdown=True)
+
+    # internist_scenario
+    pdf.set_xy(left_x, pdf.get_y())
+    pdf.set_font("OpenSans", 'B', 10)
+    pdf.cell(0, 10, 'Internist Scenario', 0, 1, 'L', markdown=True)
+    pdf.set_font("OpenSans", '', 8)
+    pdf.multi_cell(0, 3, internist_scenario, 0, 'L', markdown=True)
+    pdf.ln(5)
+
+    return bytes(pdf.output())

static-kg/ANAT_DIAG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eca33af4263adec32d995f90b0cfee696d3f1fd30046037e12aeb9111d075fab
+size 42197

static-kg/ATC3/DRUG_DIAG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92062a024d9249185b691d3aa2ce582f775dd2fc60a38c181ca0ca1dc7d66adc
+size 43241

static-kg/ATC3/PC_DRUG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e961e24caa0f2ab7b551e9f543b3bdc6d67185c8f72aec0150f473583a9c927
+size 24293

static-kg/ATC3/SYMP_DRUG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd000e18a437afad1380208b7b5a1c64973e7a662c402e0e98a5ec6ba990c9d2
+size 147177

static-kg/DIAG_SYMP.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdc8140412f83c08fc551f8d22ae3512f51032267920c5b5cd1a94e5a369b954
+size 16005

static-kg/DRUG_DIAG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2792c214c680bc127f4e7157c51ecaadcd4e2687bcad0c46edbe9f9ed0b45d23
+size 45636

static-kg/PC_DRUG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:928ecd97a131b5f7232530ad7dec556e0b13a9c3934489106140b2635e410521
+size 25264

static-kg/SYMP_DRUG.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:937c413c1a2128c09acdee88c83c625690890eb137eaf4d7cd57b805ce690b1d
+size 158485