Initialize with patient information

app.py

@@ -5,10 +5,54 @@ import torch
 import time
 import gradio as gr
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, LlamaTokenizer, TextIteratorStreamer
 import threading
 import queue
 
+# Globals
+current_model = None
+current_tokenizer = None
+
+# Curated models
+model_choices = [
+    "meta-llama/Llama-3.2-3B-Instruct",
+    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
+    "google/gemma-7b-it",
+    "mistralai/Mistral-Nemo-Instruct-FP8-2407"
+]
+
+# Example patient database
+patient_db = {
+    "001 - John Doe": {
+        "name": "John Doe",
+        "age": "45",
+        "id": "001",
+        "notes": "History of chest pain and hypertension. No prior surgeries."
+    },
+    "002 - Maria Sanchez": {
+        "name": "Maria Sanchez",
+        "age": "62",
+        "id": "002",
+        "notes": "Suspected pulmonary embolism. Shortness of breath, tachycardia."
+    },
+    "003 - Ahmed Al-Farsi": {
+        "name": "Ahmed Al-Farsi",
+        "age": "29",
+        "id": "003",
+        "notes": "Persistent migraines. MRI scheduled for brain imaging."
+    },
+    "004 - Lin Wei": {
+        "name": "Lin Wei",
+        "age": "51",
+        "id": "004",
+        "notes": "Annual screening. Family history of breast cancer."
+    }
+}
+
+# Store conversations per patient
+patient_conversations = {}
+
+
 class RichTextStreamer(TextIteratorStreamer):
     def __init__(self, tokenizer, prompt_len=0, **kwargs):
         super().__init__(tokenizer, **kwargs)
@@ -169,6 +213,10 @@ def chat_with_model(messages):
 
         messages[-1]["content"] = output_text
 
+        current_id = patient_id.value
+        if current_id:
+            patient_conversations[current_id] = messages
+
         yield messages
 
     if in_think:
@@ -182,11 +230,7 @@ def chat_with_model(messages):
 
 
 
-# Globals
-current_model = None
-current_tokenizer = None
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, LlamaTokenizer
 
 def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     global current_model, current_tokenizer
@@ -198,7 +242,7 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     progress(0.2, desc="Loading tokenizer...")
 
     # Default
-    current_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
+    current_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code= True, use_auth_token=token)
 
     progress(0.5, desc="Loading model...")
     current_model = AutoModelForCausalLM.from_pretrained(
@@ -225,50 +269,39 @@ def format_prompt(messages):
     return prompt
 
 def add_user_message(user_input, history):
-    return "", history + [{"role": "user", "content": user_input}]
+    current_id = patient_id.value
+    if current_id:
+        conversation = patient_conversations.get(current_id, [])
+        conversation.append({"role": "user", "content": user_input})
+        patient_conversations[current_id] = conversation
+    return "", patient_conversations[current_id]
 
-# Curated models
-model_choices = [
-    "meta-llama/Llama-3.2-3B-Instruct",
-    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
-    "google/gemma-7b",
-    "mistralai/Mistral-Nemo-Instruct-FP8-2407"
-]
 
-# Example patient database
-patient_db = {
-    "001 - John Doe": {
-        "name": "John Doe",
-        "age": "45",
-        "id": "001",
-        "notes": "History of chest pain and hypertension. No prior surgeries."
-    },
-    "002 - Maria Sanchez": {
-        "name": "Maria Sanchez",
-        "age": "62",
-        "id": "002",
-        "notes": "Suspected pulmonary embolism. Shortness of breath, tachycardia."
-    },
-    "003 - Ahmed Al-Farsi": {
-        "name": "Ahmed Al-Farsi",
-        "age": "29",
-        "id": "003",
-        "notes": "Persistent migraines. MRI scheduled for brain imaging."
-    },
-    "004 - Lin Wei": {
-        "name": "Lin Wei",
-        "age": "51",
-        "id": "004",
-        "notes": "Annual screening. Family history of breast cancer."
-    }
-}
 
 def autofill_patient(patient_key):
     if patient_key in patient_db:
         info = patient_db[patient_key]
+
+        # Init conversation if not existing
+        if info["id"] not in patient_conversations:
+            welcome_message = (
+                "**Welcome to the Radiologist's Companion!**\n\n"
+                "You can ask me about the patient's medical history or available imaging data.\n"
+                "- I can summarize key details from the EHR.\n"
+                "- I can tell you which medical images are available.\n"
+                "- If you'd like an organ segmentation (e.g. spleen, liver, kidney_left, colon, femur_right) on an abdominal CT scan, just ask!\n\n"
+                "**Example Requests:**\n"
+                "- \"What do we know about this patient?\"\n"
+                "- \"Which images are available for this patient?\"\n"
+                "- \"Can you segment the spleen from the CT scan?\"\n"
+            )
+
+            patient_conversations[info["id"]] = [{"role": "assistant", "content": welcome_message}]
+
         return info["name"], info["age"], info["id"], info["notes"]
     return "", "", "", ""
 
+
 with gr.Blocks(css=".gradio-container {height: 100vh; overflow: hidden;}") as demo:
     gr.Markdown("<h2 style='text-align: center;'>Radiologist's Companion</h2>")
 
@@ -311,6 +344,25 @@ with gr.Blocks(css=".gradio-container {height: 100vh; overflow: hidden;}") as de
         outputs=[patient_name, patient_age, patient_id, patient_notes]
     )
 
+    # After patient selected, load their conversation into chatbot
+    def load_patient_conversation(patient_key):
+        if patient_key in patient_db:
+            patient_id = patient_db[patient_key]["id"]
+            history = patient_conversations.get(patient_id, [])
+            return history
+        return []
+
+    patient_selector.change(
+        autofill_patient,
+        inputs=[patient_selector],
+        outputs=[patient_name, patient_age, patient_id, patient_notes]
+    ).then(
+        load_patient_conversation,
+        inputs=[patient_selector],
+        outputs=[chatbot]
+    )
+
+
     # Load on launch
     demo.load(fn=load_model_on_selection, inputs=default_model, outputs=model_status)