Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
 import time
+from typing import List, Dict, Generator
 
 # =======================================================
 # Load Model
@@ -28,45 +29,40 @@ print(f"Model device: {next(model.parameters()).device}")
 
 
 # =======================================================
-# Generate Doctor Response (Stateless + Clean Replies)
+# Response Generation
 # =======================================================
-def generate_doctor_response(history):
-    user_message = history[-1]["content"]
-
+def generate_doctor_response(user_message: str) -> Generator[str, None, None]:
+    """Generate medical advice response with streaming output."""
     if not user_message.strip():
-        history.append({"role": "assistant", "content": "⚠️ Please describe your symptoms or ask a question."})
-        yield history
+        yield "⚠️ Please describe your symptoms or ask a question."
         return
 
-    # 🩺 New Prompt (no 'Patient:' or 'Doctor:' lines)
-    prompt = f"""
-You are a compassionate and professional medical expert.
-Your role is to help users by providing clear, empathetic, and accurate medical information.
+    # Enhanced prompt - asks ONE relevant follow-up question
+    prompt = f"""You are a compassionate medical AI assistant. Provide helpful, accurate medical information.
 
 Guidelines:
-1. Do NOT include words like 'Doctor:' or 'Patient:' in your replies.
-2. Respond naturally and directly to the user's concern.
-3. Keep answers short, clear, and medically sound.
-4. Add a disclaimer when appropriate: 
-   ⚕️ *This is AI-generated information and not a substitute for professional medical advice.*
+- Respond directly without role labels like "Doctor:" or "Assistant:"
+- Be concise (2-3 sentences)
+- Provide helpful information about the symptoms
+- Ask ONE relevant follow-up question to better understand the condition
+- Include disclaimer for serious symptoms: "⚕️ Please consult a healthcare professional for proper diagnosis."
 
-Now, please respond to the user's message below:
+User's question: {user_message}
 
-User: {user_message}
-Assistant:
-"""
+Response:"""
 
-    # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
 
     gen_config = GenerationConfig(
         temperature=0.7,
         top_p=0.9,
+        top_k=50,
         do_sample=True,
-        max_new_tokens=500,
+        max_new_tokens=350,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.2
+        repetition_penalty=1.15,
+        no_repeat_ngram_size=3
     )
 
     input_len = inputs["input_ids"].shape[1]
@@ -77,96 +73,121 @@ Assistant:
     generated_ids = output_ids[0][input_len:]
     response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
-    # Keep concise output
-    response = ". ".join(response.split(". ")[:3]).strip()
-    if response.lower().startswith("assistant:"):
-        response = response[10:].strip()
-    if len(response) < 10:
+    response = clean_response(response)
+    
+    # Stream response
+    for i in range(0, len(response), 3):
+        chunk = response[:i + 3]
+        yield chunk + "▌"
+        time.sleep(0.012)
+
+    yield response
+
+
+def clean_response(response: str) -> str:
+    """Clean and format the model's response."""
+    # Remove common prefixes
+    prefixes = ["assistant:", "doctor:", "response:", "answer:"]
+    response_lower = response.lower()
+    for prefix in prefixes:
+        if response_lower.startswith(prefix):
+            response = response[len(prefix):].strip()
+            break
+    
+    # Limit length
+    sentences = response.split('. ')
+    if len(sentences) > 4:
+        response = '. '.join(sentences[:4]) + '.'
+    
+    if response and response[-1] not in '.!?':
+        response += '.'
+    
+    if len(response.strip()) < 15:
         response = "I understand your concern. Could you please provide more details about your symptoms?"
+    
+    return response
+
+
+# =======================================================
+# Chat Handler
+# =======================================================
+def respond(message: str, history: List[Dict]) -> tuple:
+    """Handle user message and generate response."""
+    user_message = message.strip()
+    if not user_message:
+        return "", history
 
-    # Stream response token by token
+    history.append({"role": "user", "content": user_message})
     history.append({"role": "assistant", "content": ""})
-    for i in range(0, len(response), 4):
-        chunk = response[:i + 4]
-        history[-1]["content"] = chunk + "▌"
-        yield history.copy()
-        time.sleep(0.015)
 
-    history[-1]["content"] = response
-    yield history
+    for partial_response in generate_doctor_response(user_message):
+        history[-1]["content"] = partial_response
+        yield "", history
+
+    return "", history
+
+
+def clear_chat():
+    """Clear the chat history."""
+    return []
 
 
 # =======================================================
 # Gradio Interface
 # =======================================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🩺 AI Doctor Chat Assistant")
+    
+    gr.Markdown("""
+    # 🩺 AI Medical Assistant
+    
+    Get medical information and guidance. This AI will ask relevant follow-up questions to better understand your condition.
+    
+    ⚠️ **Disclaimer:** For informational purposes only. Always consult healthcare professionals for medical advice.
+    """)
 
     chatbot = gr.Chatbot(
-        label="💬 Doctor Consultation",
+        label="💬 Medical Consultation",
         type='messages',
         avatar_images=(
-            "https://cdn-icons-png.flaticon.com/512/706/706830.png",  # Patient
-            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"   # Doctor
+            "https://cdn-icons-png.flaticon.com/512/706/706830.png",
+            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
         ),
-        height=500
+        height=500,
+        show_copy_button=True
     )
 
     with gr.Row():
         user_input = gr.Textbox(
-            placeholder="Type your symptoms or question here...",
-            label="🧍 Your Message",
+            placeholder="Describe your symptoms...",
+            label="💭 Your Message",
             lines=2,
             scale=4
         )
 
     with gr.Row():
         send_btn = gr.Button("💬 Send", variant="primary", scale=1)
-        clear_btn = gr.Button("🧹 Clear Chat", scale=1)
+        clear_btn = gr.Button("🧹 Clear", variant="secondary", scale=1)
 
     gr.Examples(
         examples=[
-            "I have a fever of 102°F since yesterday",
-            "I've been having headaches for the past week",
-            "I feel very tired all the time",
-            "I have a sore throat and body aches",
+            "I have a fever of 102°F and body aches",
+            "I've been having headaches for a week",
+            "I feel extremely tired all the time",
+            "I have a sore throat and cough"
         ],
         inputs=user_input,
-        label="💡 Example Questions"
     )
 
-    # =======================================================
-    # Respond Function — Model forgets, Chat UI remembers
-    # =======================================================
-    def respond(message, history):
-        user_message = message.strip()
-        if not user_message:
-            return "", history
-
-        # Show user message in chat
-        history.append({"role": "user", "content": user_message})
-
-        # Model sees only current message (no memory)
-        temp_history = [{"role": "user", "content": user_message}]
-
-        for updated_history in generate_doctor_response(temp_history):
-            if len(history) == 0 or history[-1]["role"] != "assistant":
-                history.append({"role": "assistant", "content": updated_history[-1]["content"]})
-            else:
-                history[-1]["content"] = updated_history[-1]["content"]
-            yield "", history
-
-    # =======================================================
-    # Button & Input Bindings
-    # =======================================================
-    send_btn.click(respond, [user_input, chatbot], [user_input, chatbot])
-    user_input.submit(respond, [user_input, chatbot], [user_input, chatbot])
-    clear_btn.click(lambda: [], None, chatbot, queue=False)
+    # Event handlers
+    send_btn.click(respond, [user_input, chatbot], [user_input, chatbot], queue=True)
+    user_input.submit(respond, [user_input, chatbot], [user_input, chatbot], queue=True)
+    clear_btn.click(clear_chat, outputs=chatbot, queue=False)
 
 
 # =======================================================
-# Launch App
+# Launch
 # =======================================================
 if __name__ == "__main__":
-    demo.queue()
-    demo.launch(share=True)
+    print("🚀 Starting AI Medical Assistant...")
+    demo.queue(max_size=20)
+    demo.launch(share=True, show_error=True)