Create app.py

app.py

@@ -0,0 +1,283 @@
+# app.py for Gradio App on Hugging Face Spaces
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import random
+import time
+import json
+from typing import List, Dict, Tuple, Optional
+
+# --- Configuration ---
+MODEL_NAME = "Qwen/Qwen3-VL-4B-Instruct-FP8"
+# For Hugging Face Spaces, loading the model once here is common.
+# Performance on CPU will be slower.
+print("Loading model...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16, # Use float16 for FP8 model if supported, or bfloat16
+        device_map="auto", # Automatically map to available devices (CPU or GPU)
+        trust_remote_code=True
+    ).eval() # Set to evaluation mode
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    model = None
+    tokenizer = None
+
+# --- Simulated Cost Database ---
+COSTS = {
+    "ask_question": 10.0,
+    "physical_exam": 25.0,
+    "order_cbc": 50.0,
+    "order_xray": 150.0,
+    "administer_med": 30.0,
+    "end_case": 0.0,
+    "start_case": 0.0 # Cost for starting is typically 0
+}
+
+# --- State Management Class ---
+class MedicalSimulatorState:
+    def __init__(self):
+        self.patient_profile: Optional[Dict] = None
+        self.chat_history: List[Tuple[Optional[str], Optional[str]]] = [] # Gradio chat format: [(user_msg, bot_msg), ...]
+        self.vitals: Dict[str, float] = {"HR": 72.0, "BP_Sys": 120.0, "BP_Dia": 80.0, "Temp": 98.6, "O2_Sat": 98.0}
+        self.total_cost: float = 0.0
+        self.is_case_active: bool = False
+        self.underlying_diagnosis: str = ""
+        self.ordered_tests: Dict[str, str] = {} # e.g., {"cbc": "pending", "xray": "result..."}
+        # Add more state variables as needed
+
+# --- Core AI Interaction Function ---
+def get_ai_response(user_input: str, history: List[Tuple[Optional[str], Optional[str]]], patient_profile: Dict, underlying_diagnosis: str) -> str:
+    if not model or not tokenizer:
+        return "Error: AI model is not loaded."
+
+    # Construct a prompt for the AI based on history and patient profile
+    history_str = "\n".join([f"{'User' if h[0] else 'System/AI'}: {h[0] or h[1]}" for h in history])
+    context = f"Patient Profile: Name: {patient_profile['name']}, Age: {patient_profile['age']}, Gender: {patient_profile['gender']}, Chief Complaint: {patient_profile['chief_complaint']}, History: {patient_profile['history']}, Medications: {patient_profile['medications']}, Allergies: {patient_profile['allergies']}, Social History: {patient_profile['social_history']}, Financial Status: {patient_profile['financial_status']}, Code Status: {patient_profile['code_status']}\nCurrent Chat History:\n{history_str}\nUser Action/Question: {user_input}\n"
+    prompt = f"<|system|>You are an AI patient in a medical simulation. Role-play as the patient described in the profile. Be consistent with their history, demographics, and potential complaints. Respond to the user's input (which could be a question, exam instruction, or order). Your responses should simulate realistic patient dialogue and reactions, including potential anxiety or concerns. Do not reveal the secret diagnosis '{underlying_diagnosis}' directly, but your responses should be consistent with having that condition. Respond as if you are the patient speaking.<|user|>{context}<|assistant|>"
+
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt")
+        if inputs["input_ids"].shape[1] > 32768: # Check for model max length
+             return "Error: Input prompt is too long for the model."
+
+        # Generate response
+        generate_ids = model.generate(
+            inputs.input_ids,
+            max_new_tokens=512, # Limit generated tokens
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+        # Decode the generated response
+        response_text = tokenizer.decode(generate_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response_text.strip()
+
+    except Exception as e:
+        print(f"Error during AI generation: {e}")
+        return f"An error occurred while processing the AI response: {e}"
+
+# --- Tool Functions (Modify State) ---
+def start_case(case_type: str, state: MedicalSimulatorState) -> Tuple[MedicalSimulatorState, List[Tuple[Optional[str], Optional[str]]], str, str]:
+    # --- Generate Patient Profile (Simplified Example) ---
+    names = ["John Smith", "Emily Johnson", "Michael Brown", "Sarah Davis"]
+    chief_complaints = {
+        "General": ["Chest pain", "Shortness of breath", "Abdominal pain"],
+        "Pediatric": ["Fever", "Cough", "Ear ache"],
+        "Psychiatry": ["Feeling anxious", "Difficulty sleeping", "Low mood"],
+        "Dual Diagnosis": ["Chest pain and feels anxious", "Abdominal pain after drinking"]
+    }
+    complaint_options = chief_complaints.get(case_type, chief_complaints["General"])
+    name = random.choice(names)
+    age = random.randint(18, 80) if case_type != "Pediatric" else random.randint(0, 17)
+    gender = random.choice(["Male", "Female"])
+    chief_complaint = random.choice(complaint_options)
+    # Define underlying diagnosis based on complaint or case type
+    diag_map = {
+        "Chest pain": "Acute Myocardial Infarction",
+        "Shortness of breath": "Pneumonia",
+        "Abdominal pain": "Appendicitis",
+        "Fever": "Viral Infection",
+        "Cough": "Bronchitis",
+        "Ear ache": "Otitis Media",
+        "Feeling anxious": "Generalized Anxiety Disorder",
+        "Difficulty sleeping": "Insomnia",
+        "Low mood": "Major Depressive Disorder",
+        "Chest pain and feels anxious": "Acute MI with Anxiety",
+        "Abdominal pain after drinking": "Alcoholic Gastritis with Substance Use Disorder"
+    }
+    underlying_diagnosis = diag_map.get(chief_complaint, "Unknown Condition")
+
+    patient = {
+        "name": name,
+        "age": age,
+        "gender": gender,
+        "chief_complaint": chief_complaint,
+        "history": "Patient history relevant to complaint.",
+        "medications": "Current medications.",
+        "allergies": "Known allergies (e.g., Penicillin).",
+        "social_history": "Social history details.",
+        "financial_status": "Insurance status.",
+        "code_status": "Full Code",
+        "language": "English"
+    }
+
+    # Reset state
+    state.patient_profile = patient
+    state.chat_history = [("System", "New Case Started."), ("AI Patient", f"Hi, I'm {patient['name']}. I've been having {patient['chief_complaint']}.")]
+    state.vitals = {"HR": 72.0, "BP_Sys": 120.0, "BP_Dia": 80.0, "Temp": 98.6, "O2_Sat": 98.0}
+    state.total_cost = 0.0
+    state.is_case_active = True
+    state.underlying_diagnosis = underlying_diagnosis
+    state.ordered_tests = {}
+
+    profile_str = "\n".join([f"{k.replace('_', ' ').title()}: {v}" for k, v in patient.items()])
+    return state, state.chat_history, f"${state.total_cost:.2f}", profile_str
+
+
+def handle_chat(user_input: str, history: List[Tuple[Optional[str], Optional[str]]], state: MedicalSimulatorState) -> Tuple[List[Tuple[Optional[str], Optional[str]]], str]:
+    if not state.is_case_active or not user_input.strip():
+        return history, f"${state.total_cost:.2f}"
+
+    # Add user message to history
+    history.append((user_input, None))
+
+    # Get AI response
+    ai_response = get_ai_response(user_input, history[:-1], state.patient_profile, state.underlying_diagnosis) # Pass history without the user's new message yet
+
+    # Add AI response to history
+    history[-1] = (user_input, ai_response) # Update the last entry with the AI's response
+
+    return history, f"${state.total_cost:.2f}" # Cost doesn't change here, just return current
+
+
+def use_tool(tool_name: str, state: MedicalSimulatorState) -> Tuple[MedicalSimulatorState, List[Tuple[Optional[str], Optional[str]]], str]:
+    if not state.is_case_active:
+        return state, state.chat_history, f"${state.total_cost:.2f}"
+
+    cost = COSTS.get(tool_name, 0.0)
+    state.total_cost += cost
+
+    if tool_name == "ask_question":
+        ai_response = get_ai_response("The user asks a general question to gather more history.", state.chat_history, state.patient_profile, state.underlying_diagnosis)
+        state.chat_history.append(("System", f"[Action: {tool_name}, Cost: ${cost:.2f}]"))
+        state.chat_history.append(("AI Patient", ai_response))
+
+    elif tool_name == "order_cbc":
+        state.chat_history.append(("System", f"[Action: {tool_name}, Cost: ${cost:.2f}]"))
+        state.chat_history.append(("Lab", "CBC Ordered. Result pending..."))
+        # Simulate result appearing after a delay or another action
+        # For now, add a simple result shortly after
+        time.sleep(0.5) # Simulate processing delay
+        state.chat_history.append(("Lab", "CBC Result: WBC slightly elevated, otherwise unremarkable."))
+
+    elif tool_name == "administer_med":
+        med_name = "Medication X" # Simplified, could take input
+        state.chat_history.append(("System", f"[Action: {tool_name} - {med_name}, Cost: ${cost:.2f}]"))
+        # Check for allergies here in a real app
+        state.chat_history.append(("AI Patient", f"Okay, I took the {med_name}."))
+
+    elif tool_name == "physical_exam":
+         state.chat_history.append(("System", f"[Action: {tool_name}, Cost: ${cost:.2f}]"))
+         state.chat_history.append(("System", "Physical Exam Performed. Findings noted."))
+
+    elif tool_name == "order_xray":
+         state.chat_history.append(("System", f"[Action: {tool_name}, Cost: ${cost:.2f}]"))
+         state.chat_history.append(("Imaging", "X-Ray Ordered. Result pending..."))
+         # Placeholder for image result (could be a URL or base64 string)
+         time.sleep(0.5) # Simulate processing delay
+         state.chat_history.append(("Imaging", "Chest X-Ray Result: Normal lung fields, no acute findings. (Placeholder Image)"))
+
+    # Add other tools as needed...
+
+    return state, state.chat_history, f"${state.total_cost:.2f}"
+
+
+def end_case(state: MedicalSimulatorState) -> Tuple[MedicalSimulatorState, List[Tuple[Optional[str], Optional[str]]], str, str]:
+    if not state.is_case_active:
+        # Return current state if no case is active
+        profile_str = "\n".join([f"{k.replace('_', ' ').title()}: {v}" for k, v in (state.patient_profile or {}).items()])
+        return state, state.chat_history, f"${state.total_cost:.2f}", profile_str
+
+    state.chat_history.append(("System", "Case Ended by User."))
+    state.is_case_active = False
+
+    # In a full implementation, trigger the evaluation logic here
+    # evaluation = run_evaluation(state) # Placeholder
+    # state.chat_history.append(("System", f"Evaluation: {evaluation}")) # Add evaluation to chat or separate component
+
+    profile_str = "\n".join([f"{k.replace('_', ' ').title()}: {v}" for k, v in (state.patient_profile or {}).items()])
+    return state, state.chat_history, f"${state.total_cost:.2f}", profile_str
+
+
+# --- Gradio Interface ---
+with gr.Blocks(title="Advanced Medical Simulator") as demo:
+    # State component to hold the simulator state across interactions
+    state = gr.State(lambda: MedicalSimulatorState())
+
+    gr.Markdown("# Advanced Medical Simulator")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            # Chat Interface
+            chatbot = gr.Chatbot(label="Patient Interaction", height=400, bubble_full_width=False)
+            with gr.Row():
+                user_input = gr.Textbox(label="Your Action / Question", placeholder="Type your action or question here...", scale=4)
+                submit_btn = gr.Button("Submit", scale=1)
+
+        with gr.Column(scale=1):
+            # Patient Chart / Info
+            patient_chart = gr.Markdown(label="Patient Chart", value="Click 'Start New Case' to begin.")
+            cost_display = gr.Textbox(label="Total Cost", value="$0.00", interactive=False)
+
+    with gr.Row():
+        # Tool Panel
+        with gr.Column():
+            gr.Markdown("### Tools")
+            with gr.Row():
+                ask_btn = gr.Button("Ask Question ($10)")
+                exam_btn = gr.Button("Physical Exam ($25)")
+            with gr.Row():
+                cbc_btn = gr.Button("Order CBC ($50)")
+                xray_btn = gr.Button("Order X-Ray ($150)")
+            with gr.Row():
+                 med_btn = gr.Button("Administer Med ($30)")
+                 end_btn = gr.Button("End Case", variant="stop") # Red button for ending
+
+    with gr.Row():
+        # Case Controls
+        start_case_btn = gr.Button("Start New Case (General)")
+        case_type_dropdown = gr.Dropdown(["General", "Psychiatry", "Pediatric", "Dual Diagnosis"], label="Case Type", value="General")
+
+    # Event Handling
+    start_case_btn.click(
+        fn=start_case,
+        inputs=[case_type_dropdown, state],
+        outputs=[state, chatbot, cost_display, patient_chart]
+    )
+
+    submit_btn.click(
+        fn=handle_chat,
+        inputs=[user_input, chatbot, state],
+        outputs=[chatbot, cost_display]
+    ).then(
+         fn=lambda: "", # Clear the input textbox after submission
+         inputs=[],
+         outputs=[user_input]
+    )
+
+    ask_btn.click(fn=lambda s: use_tool("ask_question", s), inputs=[state], outputs=[state, chatbot, cost_display])
+    exam_btn.click(fn=lambda s: use_tool("physical_exam", s), inputs=[state], outputs=[state, chatbot, cost_display])
+    cbc_btn.click(fn=lambda s: use_tool("order_cbc", s), inputs=[state], outputs=[state, chatbot, cost_display])
+    xray_btn.click(fn=lambda s: use_tool("order_xray", s), inputs=[state], outputs=[state, chatbot, cost_display])
+    med_btn.click(fn=lambda s: use_tool("administer_med", s), inputs=[state], outputs=[state, chatbot, cost_display])
+    end_btn.click(fn=end_case, inputs=[state], outputs=[state, chatbot, cost_display, patient_chart])
+
+# Launch the app
+# For Hugging Face Spaces, Gradio handles the launch.
+demo.launch()