Update app.py

app.py

@@ -1,159 +1,223 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
+import re
 
+# =======================================================
+# Model Loading with Error Handling
+# =======================================================
 model_name = "augtoma/qCammel-13"
 print("Loading tokenizer and model...")
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    trust_remote_code=True,
-    low_cpu_mem_usage=True
-)
-model.eval()
-print("✅ Model loaded successfully!")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True
+    )
+    model.eval()
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise
 
 
 # =======================================================
-# Global Session Memory
+# Constants
 # =======================================================
-session = {
-    "name": None,
-    "age": None,
-    "gender": None,
-    "symptoms": None,
-    "duration": None,
-    "stage": "intro"
-}
+MEDICAL_DISCLAIMER = """
+⚠️ **IMPORTANT MEDICAL DISCLAIMER:**
+- This is an AI simulation for educational purposes only
+- For medical emergencies, call emergency services immediately (911/999/local emergency number)
+- Always consult a licensed healthcare provider for diagnosis and treatment
+- Do not use this for serious, worsening, or life-threatening symptoms
+- This AI cannot diagnose conditions or prescribe medications
+"""
+
+EMERGENCY_KEYWORDS = [
+    "chest pain", "can't breathe", "breathing difficulty", "suicide", 
+    "heart attack", "stroke", "severe bleeding", "unconscious", 
+    "overdose", "severe pain", "emergency"
+]
 
 
 # =======================================================
-# Helper: Extract Name
+# Helper Functions
 # =======================================================
 def extract_name(text):
-    text = text.lower().replace("yes", "").replace("i am", "").replace("i'm", "")
-    text = text.replace("my name is", "").replace("name", "").replace("is", "").strip()
-    return text.title() if text else "Patient"
+    """Extract name from user input"""
+    text = text.lower()
+    # Remove common phrases
+    for phrase in ["yes", "i am", "i'm", "my name is", "name is", "it's", "its"]:
+        text = text.replace(phrase, "")
+    
+    # Clean and capitalize
+    text = re.sub(r'[^\w\s]', '', text).strip()
+    return text.title() if text and len(text) > 1 else "Patient"
+
+
+def extract_age(text):
+    """Extract age from user input"""
+    # Look for numbers in the text
+    numbers = re.findall(r'\b(\d{1,3})\b', text)
+    for num in numbers:
+        age = int(num)
+        if 1 <= age <= 120:  # Reasonable age range
+            return age
+    return None
+
+
+def check_emergency(text):
+    """Check if message contains emergency keywords"""
+    text_lower = text.lower()
+    return any(keyword in text_lower for keyword in EMERGENCY_KEYWORDS)
+
+
+def create_initial_session():
+    """Create a fresh session dictionary"""
+    return {
+        "name": None,
+        "age": None,
+        "gender": None,
+        "symptoms": None,
+        "duration": None,
+        "medication": None,
+        "stage": "intro"
+    }
 
 
 # =======================================================
-# Doctor Response Logic
+# Diet Suggestions
 # =======================================================
-def doctor_response(user_message):
-    global session
-    user_message = user_message.strip().lower()
-
-    # Step 1: Greeting
-    if session["stage"] == "intro":
-        session["stage"] = "ask_name"
-        return "👨‍⚕️ Hello! I’m Dr. Aiden. May I know your name, please?"
-
-    # Step 2: Get Name
-    elif session["stage"] == "ask_name":
-        session["name"] = extract_name(user_message)
-        session["stage"] = "ask_age"
-        return f"Nice to meet you, {session['name']}! How old are you?"
-
-    # Step 3: Get Age
-    elif session["stage"] == "ask_age":
-        words = user_message.split()
-        for w in words:
-            if w.isdigit():
-                session["age"] = int(w)
-                session["stage"] = "ask_gender"
-                return f"Got it, {session['name']}. Are you male or female?"
-        return "Please tell me your age in numbers, like 20 or 25."
-
-    # Step 4: Get Gender
-    elif session["stage"] == "ask_gender":
-        if "male" in user_message:
-            session["gender"] = "male"
-        elif "female" in user_message:
-            session["gender"] = "female"
-        else:
-            return "Could you please specify whether you are male or female?"
-        session["stage"] = "ask_symptoms"
-        return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}. What symptoms are you experiencing?"
-
-    # Step 5: Get Symptoms
-    elif session["stage"] == "ask_symptoms":
-        session["symptoms"] = user_message
-        session["stage"] = "ask_duration"
-        return "Since when have you been feeling this way?"
-
-    # Step 6: Duration
-    elif session["stage"] == "ask_duration":
-        session["duration"] = user_message
-        session["stage"] = "ask_medicine"
-        return "Got it. Are you taking any medications or treatments currently?"
-
-    # Step 7: Medicine intake
-    elif session["stage"] == "ask_medicine":
-        session["medication"] = user_message
-        session["stage"] = "consult"
-        return "Thank you. Let’s discuss what could be going on and how you can manage it safely."
-
-    # Step 8: Personalized Consultation
-    elif session["stage"] == "consult":
-        name = session["name"]
-        age = session["age"]
-        gender = session["gender"]
-        symptoms = session.get("symptoms", "")
-        duration = session.get("duration", "")
-        medication = session.get("medication", "")
-
-        # ------------- Handle diet-related questions -----------------
-        if any(word in user_message for word in ["diet", "food", "meal", "eat", "nutrition"]):
-            # Diet suggestions tailored to the condition
-            if "fever" in symptoms:
-                diet = (
-                    f"🍎 {name}, since you have a fever, try keeping your meals light and hydrating.\n\n"
-                    "- 🥣 **Breakfast:** Oatmeal or boiled egg with fruit (like banana or apple)\n"
-                    "- 🍲 **Lunch:** Rice with lentil soup or boiled vegetables\n"
-                    "- 🍛 **Dinner:** Light soup or grilled chicken with plain rice\n"
-                    "- 💧 **Hydration:** Drink 8–10 glasses of water, coconut water, or clear soups\n"
-                    "- 🚫 **Avoid:** Fried, oily, or spicy foods\n\n"
-                    "Eat small portions often — this helps your body recover faster."
-                )
-            elif "stomach" in symptoms or "vomit" in symptoms:
-                diet = (
-                    f"🥗 {name}, for stomach discomfort, stick to a **bland diet**:\n\n"
-                    "- 🍞 Toast, plain rice, boiled potatoes, or bananas\n"
-                    "- 💧 Sip water, oral rehydration solution, or herbal tea\n"
-                    "- 🚫 Avoid milk, fried, or spicy foods\n"
-                    "- 🍌 Eat small meals to avoid nausea"
-                )
-            elif "cold" in symptoms or "flu" in symptoms or "cough" in symptoms:
-                diet = (
-                    f"🍊 {name}, for flu or cold, focus on **immunity-boosting foods**:\n\n"
-                    "- 🍋 Citrus fruits, honey with warm water, and soups\n"
-                    "- 🍲 Chicken soup helps clear congestion\n"
-                    "- ☕ Ginger or green tea for throat relief\n"
-                    "- 🚫 Avoid sugary and chilled drinks"
-                )
-            else:
-                diet = (
-                    f"🥗 {name}, eat a balanced diet: fruits, vegetables, lean protein, and whole grains. "
-                    "Stay hydrated and avoid processed or fried foods."
-                )
+def get_diet_suggestion(name, symptoms):
+    """Get personalized diet suggestions based on symptoms"""
+    symptoms_lower = symptoms.lower()
+    
+    if "fever" in symptoms_lower:
+        return f"""🍎 **Diet Plan for {name} (Fever)**
+
+**Meals:**
+- 🥣 **Breakfast:** Oatmeal or boiled egg with banana/apple
+- 🍲 **Lunch:** Rice with lentil soup or boiled vegetables
+- 🍛 **Dinner:** Light soup or grilled chicken with plain rice
+
+**Hydration:**
+- 💧 Drink 8–10 glasses of water daily
+- Coconut water or clear soups
+- Herbal teas (chamomile, ginger)
+
+**Avoid:**
+- 🚫 Fried, oily, or spicy foods
+- Heavy meals
+- Caffeinated drinks
+
+**Tips:** Eat small portions frequently to help your body recover."""
+
+    elif any(word in symptoms_lower for word in ["stomach", "vomit", "nausea", "diarrhea"]):
+        return f"""🥗 **Diet Plan for {name} (Stomach Issues)**
+
+**BRAT Diet:**
+- 🍌 Bananas
+- 🍚 Rice (plain white rice)
+- 🍎 Applesauce
+- 🍞 Toast (dry or with minimal butter)
+
+**Hydration:**
+- 💧 Sip water frequently
+- Oral rehydration solution (ORS)
+- Clear broths
+
+**Avoid:**
+- 🚫 Dairy products
+- Fried or spicy foods
+- Caffeine and alcohol
+- Raw vegetables
+
+**Tips:** Start with small amounts and gradually increase as you feel better."""
+
+    elif any(word in symptoms_lower for word in ["cold", "flu", "cough", "sore throat"]):
+        return f"""🍊 **Diet Plan for {name} (Cold/Flu)**
+
+**Immunity Boosters:**
+- 🍋 Citrus fruits (oranges, lemons)
+- 🍯 Honey with warm water
+- 🧄 Garlic (natural antibiotic)
+- 🧅 Onions
+
+**Meals:**
+- 🍲 Chicken soup (helps clear congestion)
+- ☕ Ginger or green tea
+- 🥣 Warm porridge or oatmeal
+
+**Hydration:**
+- 💧 Plenty of warm fluids
+- Herbal teas
+
+**Avoid:**
+- 🚫 Sugary drinks
+- Chilled beverages
+- Dairy (may increase mucus)
+
+**Tips:** Stay warm and get plenty of rest."""
+
+    elif any(word in symptoms_lower for word in ["headache", "migraine"]):
+        return f"""🧘 **Diet Plan for {name} (Headache)**
+
+**Helpful Foods:**
+- 💧 Stay well-hydrated (dehydration causes headaches)
+- 🥔 Magnesium-rich: bananas, nuts, spinach
+- 🐟 Omega-3: fish, walnuts
+- 🧂 Small amounts of salt if low blood pressure
+
+**Avoid:**
+- 🚫 Processed foods
+- Artificial sweeteners
+- Aged cheeses
+- Alcohol
+- Excessive caffeine
+
+**Tips:** Eat regular meals; skipping meals can trigger headaches."""
+
+    else:
+        return f"""🥗 **General Healthy Diet Plan for {name}**
+
+**Balanced Meals:**
+- 🍎 Fresh fruits and vegetables
+- 🍗 Lean proteins (chicken, fish, legumes)
+- 🍚 Whole grains (brown rice, oats)
+- 🥜 Nuts and seeds
+
+**Hydration:**
+- 💧 8–10 glasses of water daily
+
+**Avoid:**
+- 🚫 Processed and fried foods
+- Excessive sugar and salt
+- Junk food
+
+**Tips:** Eat at regular intervals and maintain portion control."""
 
-            return diet + "\n\n⚕️ *Note: This is general advice — not a substitute for medical care.*"
 
-        # ----------------- Normal medical consultation ----------------
-        prompt = f"""
-You are Dr. Aiden — a warm, caring, and professional doctor.
-You are consulting a {age}-year-old {gender} named {name}.
+# =======================================================
+# AI Response Generation
+# =======================================================
+def generate_ai_response(session, user_message):
+    """Generate AI consultation response using the model"""
+    try:
+        # Build the prompt
+        prompt = f"""You are Dr. Aiden — a warm, caring, and professional doctor.
+You are consulting a {session['age']}-year-old {session['gender']} named {session['name']}.
 
 Patient details:
-- Symptoms: {symptoms}
-- Duration: {duration}
-- Current medications: {medication}
+- Symptoms: {session['symptoms']}
+- Duration: {session['duration']}
+- Current medications: {session['medication']}
 
 Now the patient says: "{user_message}"
 
@@ -164,10 +228,11 @@ Include:
 3. Simple home remedies or OTC medicines (if safe)
 4. Diet, rest, and hydration tips
 5. When to visit a real doctor
-6. End with a reassuring tone
-"""
+6. End with a reassuring tone"""
 
+        # Tokenize and generate
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
+        
         gen_cfg = GenerationConfig(
             temperature=0.7,
             top_p=0.9,
@@ -180,66 +245,240 @@ Include:
         with torch.no_grad():
             output = model.generate(**inputs, generation_config=gen_cfg)
 
+        # Decode the output
         output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-        output_text = output_text.split("Doctor:")[-1].strip()
-
+        
+        # Try to extract just the response part
+        # The model might repeat the prompt, so we try to find where the actual response starts
+        if "reassuring tone" in output_text:
+            parts = output_text.split("reassuring tone")
+            if len(parts) > 1:
+                output_text = parts[-1].strip()
+        
+        # If the output contains the prompt, try to split it
+        if user_message in output_text:
+            parts = output_text.split(user_message)
+            if len(parts) > 1:
+                output_text = parts[-1].strip()
+        
+        # Clean up
+        output_text = output_text.strip()
+        
+        # If output is too short or empty, provide a fallback
+        if len(output_text) < 20:
+            output_text = f"Hello {session['name']}, I understand you have {session['symptoms']} for {session['duration']}. This could be due to several reasons. I recommend staying hydrated, getting plenty of rest, and monitoring your temperature. If your fever persists beyond 5 days or gets worse, please visit a doctor in person for proper examination."
+        
+        # Ensure proper ending punctuation
         if not output_text.endswith((".", "!", "?")):
             output_text += "."
 
+        # Add disclaimer
         output_text += "\n\n⚕️ *Note: This advice is AI-generated and not a substitute for professional medical care.*"
+        
         return output_text
+        
+    except Exception as e:
+        print(f"Error in generate_ai_response: {e}")
+        return f"I understand you're experiencing {session['symptoms']}, {session['name']}. Since this has been going on for {session['duration']}, I recommend consulting with a healthcare provider for a proper examination and diagnosis.\n\n⚕️ *Note: This advice is AI-generated and not a substitute for professional medical care.*"
 
 
+# =======================================================
+# Doctor Response Logic
+# =======================================================
+def doctor_response(user_message, session):
+    """Generate doctor's response based on conversation stage"""
+    if not user_message or not user_message.strip():
+        return "Please type your message.", session
+    
+    user_message_clean = user_message.strip()
+    user_message_lower = user_message_clean.lower()
+    
+    # Check for emergency situations
+    if check_emergency(user_message_clean):
+        emergency_response = """🚨 **EMERGENCY ALERT** 🚨
+
+Your message suggests a potentially serious medical emergency.
+
+**PLEASE TAKE IMMEDIATE ACTION:**
+- Call emergency services NOW (911 in US, 999 in UK, or your local emergency number)
+- Go to the nearest emergency room
+- Do not delay seeking immediate medical attention
+
+This AI chatbot cannot handle emergencies. Your safety is the priority."""
+        return emergency_response, session
+    
+    # Stage-based conversation flow
+    try:
+        if session["stage"] == "intro":
+            session["stage"] = "ask_name"
+            return "👨‍⚕️ Hello! I'm Dr. Aiden, your AI medical assistant.\n\n" + MEDICAL_DISCLAIMER + "\n\nMay I know your name, please?", session
+        
+        elif session["stage"] == "ask_name":
+            session["name"] = extract_name(user_message_clean)
+            session["stage"] = "ask_age"
+            return f"Nice to meet you, {session['name']}! How old are you?", session
+        
+        elif session["stage"] == "ask_age":
+            age = extract_age(user_message_clean)
+            if age:
+                session["age"] = age
+                session["stage"] = "ask_gender"
+                return f"Got it, {session['name']}. Are you male or female?", session
+            else:
+                return "Please tell me your age in numbers (e.g., 25, 30, 45).", session
+        
+        elif session["stage"] == "ask_gender":
+            if "male" in user_message_lower and "female" not in user_message_lower:
+                session["gender"] = "male"
+            elif "female" in user_message_lower:
+                session["gender"] = "female"
+            else:
+                return "Could you please specify whether you are male or female?", session
+            
+            session["stage"] = "ask_symptoms"
+            return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}.\n\nWhat symptoms are you experiencing? Please describe them in detail.", session
+        
+        elif session["stage"] == "ask_symptoms":
+            if len(user_message_clean) < 5:
+                return "Please describe your symptoms in more detail so I can help you better.", session
+            session["symptoms"] = user_message_clean
+            session["stage"] = "ask_duration"
+            return "Since when have you been feeling this way? (e.g., 2 days, 1 week, etc.)", session
+        
+        elif session["stage"] == "ask_duration":
+            session["duration"] = user_message_clean
+            session["stage"] = "ask_medicine"
+            return "Are you currently taking any medications or treatments? If not, just say 'no'.", session
+        
+        elif session["stage"] == "ask_medicine":
+            session["medication"] = user_message_clean
+            session["stage"] = "consult"
+            # Provide summary and automatically generate initial consultation
+            summary = f"""Thank you for providing all the information, {session['name']}.
+
+📋 **Summary:**
+- **Name:** {session['name']}
+- **Age:** {session['age']} years old
+- **Gender:** {session['gender']}
+- **Symptoms:** {session['symptoms']}
+- **Duration:** {session['duration']}
+- **Current medication:** {session['medication']}
+
+Let me assess your condition...\n\n"""
+            
+            # Generate initial consultation automatically
+            initial_consultation = generate_ai_response(session, "Please assess my condition and give me advice")
+            return summary + initial_consultation, session
+        
+        elif session["stage"] == "consult":
+            # Handle diet-specific queries
+            if any(word in user_message_lower for word in ["diet", "food", "meal", "eat", "nutrition", "drink"]):
+                diet_plan = get_diet_suggestion(session["name"], session["symptoms"])
+                return diet_plan + "\n\n⚕️ *Note: This is general advice. Consult a healthcare provider for personalized dietary recommendations, especially if you have allergies or other conditions.*", session
+            
+            # Generate AI consultation response for any other question
+            return generate_ai_response(session, user_message_clean), session
+        
+        else:
+            # Reset if stage is unknown
+            session["stage"] = "intro"
+            return "Let's start over. Say 'Hi Doctor' to begin your consultation.", session
+    
+    except Exception as e:
+        print(f"Error in doctor_response: {e}")
+        return "I apologize, but I encountered an error. Please try again or start a new consultation.", session
+
 
 # =======================================================
 # Gradio Interface
 # =======================================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.HTML("""
-    <div style="text-align:center; background-color:#4C7DFF; color:white; padding:20px; border-radius:10px;">
-        <h1>🏥 Doctor Consultation with Dr. Aiden</h1>
-        <p>AI-powered doctor interview — step-by-step and caring conversation</p>
+    <div style="text-align:center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color:white; padding:25px; border-radius:15px; box-shadow: 0 4px 6px rgba(0,0,0,0.1);">
+        <h1 style="margin:0; font-size: 2.5em;">🏥 Dr. Aiden - AI Medical Assistant</h1>
+        <p style="margin:10px 0 0 0; font-size: 1.1em;">Compassionate AI-powered health consultation</p>
+        <p style="margin:5px 0 0 0; font-size: 0.9em; opacity: 0.9;">⚠️ For emergencies, call 911 immediately</p>
     </div>
     """)
-
+    
+    # Session state
+    session_state = gr.State(create_initial_session)
+    
     chatbot = gr.Chatbot(
         label="👨‍⚕️ Chat with Dr. Aiden",
-        height=550,
+        height=500,
         type='messages',
         avatar_images=(
             "https://cdn-icons-png.flaticon.com/512/706/706830.png",
             "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
         )
     )
-
-    user_input = gr.Textbox(placeholder="Say 'Hi Doctor' to start your consultation...", label="Your Message", lines=2)
-    send_btn = gr.Button("💬 Send", variant="primary")
-    clear_btn = gr.Button("🧹 New Consultation")
-
-    def respond(message, history):
+    
+    with gr.Row():
+        user_input = gr.Textbox(
+            placeholder="Type your message here (say 'Hi Doctor' to start)...", 
+            label="Your Message", 
+            lines=2,
+            scale=4
+        )
+    
+    with gr.Row():
+        send_btn = gr.Button("💬 Send", variant="primary", scale=1)
+        clear_btn = gr.Button("🔄 New Consultation", variant="secondary", scale=1)
+    
+    gr.HTML("""
+    <div style="margin-top: 20px; padding: 15px; background-color: #fff3cd; border-left: 4px solid #ffc107; border-radius: 5px;">
+        <strong>⚠️ Important Reminders:</strong>
+        <ul style="margin: 10px 0; padding-left: 20px;">
+            <li>This is an AI simulation for educational purposes only</li>
+            <li>Always consult a real doctor for medical diagnosis and treatment</li>
+            <li>For emergencies, call 911 or your local emergency services immediately</li>
+            <li>Do not use this for serious or life-threatening conditions</li>
+        </ul>
+    </div>
+    """)
+    
+    def respond(message, history, session):
+        """Handle user message and update chat"""
         if history is None:
             history = []
-        response = doctor_response(message)
+        
+        if not message or not message.strip():
+            return "", history, session
+        
+        response, updated_session = doctor_response(message, session)
+        
         history.append({"role": "user", "content": message})
         history.append({"role": "assistant", "content": response})
-        return "", history
-
+        
+        return "", history, updated_session
+    
     def reset():
-        global session
-        session = {
-            "name": None,
-            "age": None,
-            "gender": None,
-            "symptoms": None,
-            "duration": None,
-            "stage": "intro"
-        }
-        return []
-    send_btn.click(respond, [user_input, chatbot], [user_input, chatbot])
-    user_input.submit(respond, [user_input, chatbot], [user_input, chatbot])
-    clear_btn.click(reset, None, chatbot, queue=False)
+        """Reset conversation and session"""
+        return [], create_initial_session()
+    
+    # Event handlers
+    send_btn.click(
+        respond, 
+        [user_input, chatbot, session_state], 
+        [user_input, chatbot, session_state]
+    )
+    
+    user_input.submit(
+        respond, 
+        [user_input, chatbot, session_state], 
+        [user_input, chatbot, session_state]
+    )
+    
+    clear_btn.click(
+        reset, 
+        None, 
+        [chatbot, session_state], 
+        queue=False
+    )
 
 if __name__ == "__main__":
-    print("🏥 Launching Dr. Aiden...")
+    print("🏥 Launching Dr. Aiden Medical Assistant...")
+    print("⚠️  Remember: This is for educational purposes only!")
     demo.queue()
-    demo.launch(share=True)
+    demo.launch(share=True)