app.py

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...")

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


# =======================================================
# Constants
# =======================================================
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 Functions
# =======================================================
def extract_name(text):
    """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"
    }


# =======================================================
# Diet Suggestions
# =======================================================
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."""


# =======================================================
# 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: {session['symptoms']}
- Duration: {session['duration']}
- Current medications: {session['medication']}

Now the patient says: "{user_message}"

Respond as a real doctor would — empathetic, clear, and personalized.
Include:
1. Acknowledge their condition
2. Possible causes
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"""

        # 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,
            max_new_tokens=350,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

        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)
        
        # 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: 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=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"
        )
    )
    
    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 = []
        
        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, updated_session
    
    def reset():
        """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 Medical Assistant...")
    print("⚠️  Remember: This is for educational purposes only!")
    demo.queue()
    demo.launch(share=True)