app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch

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


# =======================================================
# Global Session Memory
# =======================================================
session = {
    "name": None,
    "age": None,
    "gender": None,
    "symptoms": None,
    "duration": None,
    "stage": "intro"
}


# =======================================================
# Helper: Extract Name
# =======================================================
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"


# =======================================================
# Doctor Response Logic
# =======================================================
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."
                )

            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}.

Patient details:
- Symptoms: {symptoms}
- Duration: {duration}
- Current medications: {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
"""

        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)

        output_text = tokenizer.decode(output[0], skip_special_tokens=True)
        output_text = output_text.split("Doctor:")[-1].strip()

        if not output_text.endswith((".", "!", "?")):
            output_text += "."

        output_text += "\n\n⚕️ *Note: This advice is AI-generated and not a substitute for professional medical care.*"
        return output_text



# =======================================================
# 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>
    """)

    chatbot = gr.Chatbot(
        label="👨‍⚕️ Chat with Dr. Aiden",
        height=550,
        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):
        if history is None:
            history = []
        response = doctor_response(message)
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response})
        return "", history

    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)

if __name__ == "__main__":
    print("🏥 Launching Dr. Aiden...")
    demo.queue()
    demo.launch(share=True)