Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
 import time
-from typing import List, Dict, Generator
 
 # =======================================================
 # Load Model
@@ -29,39 +28,52 @@ print(f"Model device: {next(model.parameters()).device}")
 
 
 # =======================================================
-# Response Generation
+# Generate Comprehensive Doctor Response
 # =======================================================
-def generate_doctor_response(user_message: str) -> Generator[str, None, None]:
-    """Generate medical advice response with streaming output."""
+def generate_doctor_response(history):
+    user_message = history[-1]["content"]
+
     if not user_message.strip():
-        yield "⚠️ Please describe your symptoms or ask a question."
+        history.append({"role": "assistant", "content": "⚠️ Please describe your symptoms or ask a question."})
+        yield history
         return
 
-    # Enhanced prompt - asks ONE relevant follow-up question
-    prompt = f"""You are a compassionate medical AI assistant. Provide helpful, accurate medical information.
+    # Enhanced Medical Prompt - Comprehensive Doctor Approach
+    prompt = f"""You are an experienced and compassionate medical doctor providing a comprehensive consultation.
+
+Based on the patient's concern, provide a detailed response that includes:
+
+1. **Assessment**: Acknowledge their symptoms and provide initial medical assessment
+2. **Possible Causes**: Explain potential causes or conditions
+3. **Medications**: Recommend appropriate over-the-counter or prescription medications (with dosages when relevant)
+4. **Nutrition & Diet**: Suggest specific foods, nutrients, or dietary changes that can help
+5. **Lifestyle Modifications**: Recommend lifestyle changes, exercises, rest, or habits to adopt
+6. **Follow-up**: Advise when to see a doctor or what warning signs to watch for
 
 Guidelines:
-- 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."
+- Do NOT use labels like "Doctor:" or "Patient:" in your response
+- Be professional, empathetic, and thorough like a real doctor
+- Provide specific, actionable recommendations
+- Use medical terminology but explain it simply
+- Structure your response clearly with the categories above
+- End with: "⚕️ *Please consult a healthcare provider for proper diagnosis and personalized treatment plan.*"
 
-User's question: {user_message}
+Patient's concern: {user_message}
 
-Response:"""
+Comprehensive Medical Response:"""
 
+    # Tokenize input
     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_p=0.92,
         top_k=50,
         do_sample=True,
-        max_new_tokens=350,
+        max_new_tokens=600,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.15,
+        repetition_penalty=1.18,
         no_repeat_ngram_size=3
     )
 
@@ -73,121 +85,166 @@ Response:"""
     generated_ids = output_ids[0][input_len:]
     response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
-    response = clean_response(response)
-    
-    # Stream response
-    for i in range(0, len(response), 3):
-        chunk = response[:i + 3]
-        yield chunk + "▌"
-        time.sleep(0.012)
+    # Clean response
+    response = clean_medical_response(response)
 
-    yield response
+    # Stream response token by token
+    history.append({"role": "assistant", "content": ""})
+    for i in range(0, len(response), 5):
+        chunk = response[:i + 5]
+        history[-1]["content"] = chunk + "▌"
+        yield history.copy()
+        time.sleep(0.01)
+
+    history[-1]["content"] = response
+    yield history
 
 
-def clean_response(response: str) -> str:
-    """Clean and format the model's response."""
+def clean_medical_response(response: str) -> str:
+    """Clean and format the medical response."""
     # Remove common prefixes
-    prefixes = ["assistant:", "doctor:", "response:", "answer:"]
+    prefixes = ["assistant:", "doctor:", "response:", "comprehensive medical response:", "medical response:"]
     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]) + '.'
+    # Remove any remaining role labels
+    lines = response.split('\n')
+    cleaned_lines = []
+    for line in lines:
+        if not line.lower().strip().startswith(('doctor:', 'assistant:', 'patient:')):
+            cleaned_lines.append(line)
+    response = '\n'.join(cleaned_lines)
     
+    # Ensure proper ending
     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?"
+    # Add disclaimer if not present
+    if '⚕️' not in response and 'consult' not in response.lower():
+        response += '\n\n⚕️ *Please consult a healthcare provider for proper diagnosis and personalized treatment plan.*'
     
-    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
-
-    history.append({"role": "user", "content": user_message})
-    history.append({"role": "assistant", "content": ""})
-
-    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 []
+    # Fallback for very short responses
+    if len(response.strip()) < 30:
+        response = "I understand your concern. To provide you with comprehensive medical guidance including medications, diet, and lifestyle recommendations, could you please describe your symptoms in more detail? For example, when did they start, how severe are they, and have you noticed any triggers?"
+    
+    return response.strip()
 
 
 # =======================================================
 # Gradio Interface
 # =======================================================
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    
-    gr.Markdown("""
-    # 🩺 AI Medical Assistant
+with gr.Blocks(theme=gr.themes.Soft(), css="""
+    .medical-header {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        padding: 20px;
+        border-radius: 10px;
+        color: white;
+        text-align: center;
+        margin-bottom: 20px;
+    }
+""") as demo:
     
-    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.
+    gr.HTML("""
+        <div class="medical-header">
+            <h1>🏥 AI Medical Doctor Consultation</h1>
+            <p>Comprehensive Medical Guidance • Medications • Nutrition • Lifestyle</p>
+        </div>
     """)
 
     chatbot = gr.Chatbot(
-        label="💬 Medical Consultation",
+        label="💬 Doctor-Patient Consultation",
         type='messages',
         avatar_images=(
-            "https://cdn-icons-png.flaticon.com/512/706/706830.png",
-            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
+            "https://cdn-icons-png.flaticon.com/512/706/706830.png",  # Patient
+            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"   # Doctor
         ),
-        height=500,
+        height=550,
         show_copy_button=True
     )
 
     with gr.Row():
         user_input = gr.Textbox(
-            placeholder="Describe your symptoms...",
-            label="💭 Your Message",
-            lines=2,
+            placeholder="Describe your symptoms in detail (e.g., 'I have fever, headache, and body pain for 3 days')...",
+            label="🧍 Describe Your Symptoms",
+            lines=3,
             scale=4
         )
 
     with gr.Row():
-        send_btn = gr.Button("💬 Send", variant="primary", scale=1)
-        clear_btn = gr.Button("🧹 Clear", variant="secondary", scale=1)
+        send_btn = gr.Button("💬 Consult Doctor", variant="primary", scale=1, size="lg")
+        clear_btn = gr.Button("🧹 New Consultation", scale=1, size="lg")
 
+    gr.Markdown("### 💡 Example Consultations")
     gr.Examples(
         examples=[
-            "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"
+            "I have a fever of 102°F, headache, and body aches for 2 days. What should I do?",
+            "I've been having persistent headaches and feeling tired. Need advice on diet and lifestyle.",
+            "I have acidity and stomach pain after eating. What medications and diet should I follow?",
+            "I'm feeling stressed and anxious. Suggest lifestyle changes and natural remedies.",
+            "I have high blood pressure. What diet and lifestyle changes should I make?",
+            "I caught a cold with sore throat and cough. What treatment do you recommend?",
         ],
         inputs=user_input,
     )
 
-    # 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)
+    gr.Markdown("""
+    ---
+    ### 🩺 What This AI Doctor Provides:
+    
+    ✅ **Medical Assessment** - Initial diagnosis and condition evaluation  
+    ✅ **Medication Recommendations** - Appropriate medicines with dosage guidance  
+    ✅ **Nutrition & Diet Plans** - Foods and nutrients to help recovery  
+    ✅ **Lifestyle Modifications** - Exercise, sleep, stress management tips  
+    ✅ **Follow-up Advice** - When to see a doctor and warning signs  
+    
+    ⚠️ **Important Medical Disclaimer:**  
+    This AI provides general medical information for educational purposes only. It is NOT a substitute for professional medical advice, diagnosis, or treatment. Always seek the advice of your physician or qualified healthcare provider with any questions about a medical condition. Never disregard professional medical advice or delay seeking it because of something you have read here.
+    
+    🚨 **Emergency:** If you are experiencing a medical emergency, call emergency services immediately.
+    """)
+
+    # =======================================================
+    # Respond Function
+    # =======================================================
+    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 (stateless for consistent behavior)
+        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)
 
 
 # =======================================================
-# Launch
+# Launch App
 # =======================================================
 if __name__ == "__main__":
-    print("🚀 Starting AI Medical Assistant...")
+    print("="*60)
+    print("🏥 AI Medical Doctor Consultation System Starting...")
+    print("="*60)
     demo.queue(max_size=20)
-    demo.launch(share=True, show_error=True)
+    demo.launch(
+        share=True,
+        show_error=True,
+        server_name="0.0.0.0"
+    )