Update app.py

app.py

@@ -3,7 +3,6 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-import re
 
 phi4_model_path = "Intelligent-Internet/II-Medical-8B"
 
@@ -16,30 +15,23 @@ phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
 def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
     if not user_message.strip():
         return history_state, history_state
-        
+
     model = phi4_model
     tokenizer = phi4_tokenizer
     start_tag = "<|im_start|>"
     sep_tag = "<|im_sep|>"
     end_tag = "<|im_end|>"
 
-    system_message = """You are a highly knowledgeable and thoughtful AI medical assistant. Your primary role is to assist with diagnostic reasoning by evaluating patient symptoms, medical history, and relevant clinical context. 
+    system_message = """You are a medical assistant AI designed to help diagnose symptoms, explain possible conditions, and recommend next steps. You must be cautious, thorough, and explain medical reasoning step-by-step. Structure your answer in two sections: 
 
-Structure your response into two main sections using the following format: <think> {Thought section} </think> {Solution section}.
+<think> In this section, reason through the symptoms by considering patient history, differential diagnoses, relevant physiological mechanisms, and possible investigations. Explain your thought process step-by-step. </think> 
 
-In the <think> section, use structured clinical reasoning to:
-- Identify possible differential diagnoses based on the given symptoms.
-- Consider risk factors, medical history, duration, and severity of symptoms.
-- Use step-by-step logic to rule in or rule out conditions.
-- Reflect on diagnostic uncertainty and suggest further assessments if needed.
+In the Solution section, summarize your working diagnosis, differential options, and suggest what to do next (e.g., tests, referral, lifestyle changes). Always clarify that this is not a replacement for a licensed medical professional.
 
-In the <solution> section, provide your most likely diagnosis or clinical assessment along with the rationale. Include brief suggestions for potential next steps like labs, imaging, or referrals if appropriate.
+Use LaTeX for any formulas or values (e.g., $\\text{BMI} = \\frac{\\text{weight (kg)}}{\\text{height (m)}^2}$). 
 
-IMPORTANT: When referencing lab values or pathophysiological mechanisms, use LaTeX formatting for clarity. Use $...$ for inline and $$...$$ for block-level expressions.
+Now, analyze the following case:"""
 
-Now, please analyze and respond to the following case:
-"""
-    
     prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
     for message in history_state:
         if message["role"] == "user":
@@ -51,7 +43,6 @@ Now, please analyze and respond to the following case:
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-
     generation_kwargs = {
         "input_ids": inputs["input_ids"],
         "attention_mask": inputs["attention_mask"],
@@ -81,15 +72,14 @@ Now, please analyze and respond to the following case:
 
     yield new_history, new_history
 
-# Updated example cases for medical diagnostics
+
 example_messages = {
-    "Chest Pain": "A 58-year-old man presents with chest pain that started 20 minutes ago while climbing stairs. He describes it as a heavy pressure in the center of his chest, radiating to his left arm. He has a history of hypertension and smoking. What is the likely diagnosis?",
-    "Shortness of Breath": "A 34-year-old woman presents with 3 days of worsening shortness of breath, low-grade fever, and a dry cough. She denies chest pain or recent travel. Pulse oximetry is 91% on room air.",
-    "Abdominal Pain": "A 22-year-old female presents with lower right quadrant abdominal pain, nausea, and fever. The pain started around the umbilicus and migrated to the right lower quadrant over the past 12 hours.",
-    "Pediatric Fever": "A 2-year-old child has a fever of 39.5°C, irritability, and a rash on the trunk and arms. The child received all standard vaccinations and has no sick contacts. What should be considered in the differential diagnosis?"
+    "Headache case": "A 35-year-old female presents with a throbbing headache, nausea, and sensitivity to light. It started on one side of her head and worsens with activity. No prior trauma.",
+    "Chest pain": "A 58-year-old male presents with chest tightness radiating to his left arm, shortness of breath, and sweating. Symptoms began while climbing stairs.",
+    "Abdominal pain": "A 24-year-old complains of right lower quadrant abdominal pain, nausea, and mild fever. The pain started around the belly button and migrated.",
+    "BMI calculation": "A patient weighs 85 kg and is 1.75 meters tall. Calculate the BMI and interpret whether it's underweight, normal, overweight, or obese."
 }
 
-# Custom CSS
 css = """
 .markdown-body .katex { 
     font-size: 1.2em; 
@@ -102,12 +92,7 @@ css = """
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    gr.Markdown(
-        """
-        # Medical Diagnosis Assistant
-        This AI assistant uses structured reasoning to evaluate clinical cases and assist with diagnostic decision-making. Includes LaTeX support for medical calculations.
-        """
-    )
+    gr.Markdown("# Medical Diagnostic Assistant\nThis AI assistant helps analyze symptoms and provide preliminary diagnostic reasoning using LaTeX-rendered medical formulas where needed.")
 
     gr.HTML("""
     <script>
@@ -126,7 +111,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             messageStyle: 'none'
         };
     }
-
     function rerender() {
         if (window.MathJax && window.MathJax.Hub) {
             window.MathJax.Hub.Queue(['Typeset', window.MathJax.Hub]);
@@ -147,21 +131,21 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 top_k_slider = gr.Slider(1, 100, step=1, value=50, label="Top-k")
                 top_p_slider = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
                 repetition_penalty_slider = gr.Slider(1.0, 2.0, value=1.0, label="Repetition Penalty")
-        
+
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", render_markdown=True, type="messages", elem_id="chatbot", show_copy_button=True)
+            chatbot = gr.Chatbot(label="Chat", render_markdown=True, type="messages", show_copy_button=True)
             with gr.Row():
-                user_input = gr.Textbox(label="Describe patient symptoms...", placeholder="Type a clinical case here...", scale=3)
+                user_input = gr.Textbox(label="Describe symptoms or ask a medical question", placeholder="Type your message here...", scale=3)
                 submit_button = gr.Button("Send", variant="primary", scale=1)
                 clear_button = gr.Button("Clear", scale=1)
-            gr.Markdown("**Try these example cases:**")
+            gr.Markdown("**Try these examples:**")
             with gr.Row():
-                example1_button = gr.Button("Chest Pain")
-                example2_button = gr.Button("Shortness of Breath")
-                example3_button = gr.Button("Abdominal Pain")
-                example4_button = gr.Button("Pediatric Fever")
+                example1 = gr.Button("Headache case")
+                example2 = gr.Button("Chest pain")
+                example3 = gr.Button("Abdominal pain")
+                example4 = gr.Button("BMI calculation")
 
-    submit_button.click(
+    submit_button.stream(
         fn=generate_response,
         inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],
         outputs=[chatbot, history_state]
@@ -171,15 +155,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=user_input
     )
 
-    clear_button.click(
-        fn=lambda: ([], []),
-        inputs=None,
-        outputs=[chatbot, history_state]
-    )
+    clear_button.click(fn=lambda: ([], []), inputs=None, outputs=[chatbot, history_state])
 
-    example1_button.click(fn=lambda: gr.update(value=example_messages["Chest Pain"]), inputs=None, outputs=user_input)
-    example2_button.click(fn=lambda: gr.update(value=example_messages["Shortness of Breath"]), inputs=None, outputs=user_input)
-    example3_button.click(fn=lambda: gr.update(value=example_messages["Abdominal Pain"]), inputs=None, outputs=user_input)
-    example4_button.click(fn=lambda: gr.update(value=example_messages["Pediatric Fever"]), inputs=None, outputs=user_input)
+    example1.click(lambda: gr.update(value=example_messages["Headache case"]), None, user_input)
+    example2.click(lambda: gr.update(value=example_messages["Chest pain"]), None, user_input)
+    example3.click(lambda: gr.update(value=example_messages["Abdominal pain"]), None, user_input)
+    example4.click(lambda: gr.update(value=example_messages["BMI calculation"]), None, user_input)
 
 demo.launch(ssr_mode=False)