Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
-import time
 
 # =======================================================
 # Load Model
@@ -21,152 +20,180 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 model.eval()
-
 print("✅ Model loaded successfully!")
-print(f"📦 Device map: {model.hf_device_map}")
-print(f"💻 Model device: {next(model.parameters()).device}")
+
 
 # =======================================================
-# Session memory (for patient info)
+# Global Session Memory
 # =======================================================
-patient_data = {"name": None, "age": None, "gender": None, "collected": False}
+session = {
+    "name": None,
+    "age": None,
+    "gender": None,
+    "symptoms": None,
+    "duration": None,
+    "stage": "intro"
+}
+
 
 # =======================================================
-# Generate Doctor Response
+# Helper: Extract Name
 # =======================================================
-def generate_response(user_message):
-    name, age, gender, collected = (
-        patient_data["name"],
-        patient_data["age"],
-        patient_data["gender"],
-        patient_data["collected"],
-    )
+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"
 
-    # Stage 1 – ask for name
-    if name is None:
-        patient_data["name"] = user_message.strip()
-        return f"👨‍⚕️ Hello! I’m Dr. Aiden. May I know your name, please?" if not name else f"Nice to meet you, {name}! How old are you?"
 
-    # Stage 2 – ask for age
-    if age is None:
-        if user_message.strip().isdigit():
-            patient_data["age"] = user_message.strip()
-            return f"Got it, {name}. Are you male or female?"
-        else:
-            return "Please enter your age in numbers (e.g., 25)."
-
-    # Stage 3 – ask for gender
-    if gender is None:
-        gender_input = user_message.lower()
-        if "male" in gender_input:
-            patient_data["gender"] = "male"
-        elif "female" in gender_input:
-            patient_data["gender"] = "female"
+# =======================================================
+# Doctor Response Logic
+# =======================================================
+def doctor_response(user_message):
+    global session
+    user_message = user_message.strip()
+
+    # Intro
+    if session["stage"] == "intro":
+        session["stage"] = "ask_name"
+        return "👨‍⚕️ Hello! I’m Dr. Aiden. May I know your name, please?"
+
+    # Ask 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?"
+
+    # Ask 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."
+
+    # Ask Gender
+    elif session["stage"] == "ask_gender":
+        if "male" in user_message.lower():
+            session["gender"] = "male"
+        elif "female" in user_message.lower():
+            session["gender"] = "female"
         else:
-            return "Please specify 'male' or 'female'."
-        return f"Thanks, {name}! So you're a {patient_data['age']}-year-old {patient_data['gender']}. What brings you in today?"
-
-    # Stage 4 – consult mode (medical questions)
-    if collected is False:
-        patient_data["collected"] = True
-
-    # Doctor consultation response
-    prompt = f"""
-You are Dr. Aiden — a warm, professional, and empathetic medical doctor having a one-on-one text consultation.
-
-The patient’s details:
-- Name: {name}
-- Age: {age}
-- Gender: {gender}
-
-You are now starting the medical discussion.
-Respond naturally and fully — as if you are talking to a real patient for the first time about their symptoms.
-
-Be sure to include:
-1. A short, caring acknowledgment of the patient's situation
-2. Possible causes (simple explanation)
-3. Basic medicine or home remedy suggestion (only safe over-the-counter advice)
-4. Rest, hydration, and nutrition guidance
-5. When to consult a real doctor
-6. A gentle closing line
-7. End with a disclaimer.
-
-Patient says: "{user_message}"
+            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?"
+
+    # Ask Symptoms
+    elif session["stage"] == "ask_symptoms":
+        session["symptoms"] = user_message
+        session["stage"] = "ask_duration"
+        return "Since when have you been feeling this way?"
+
+    # Ask Duration
+    elif session["stage"] == "ask_duration":
+        session["duration"] = user_message
+        session["stage"] = "consult"
+        return "Got it. Are you taking any medications or treatments currently?"
+
+    # Consultation
+    elif session["stage"] == "consult":
+        name = session["name"]
+        age = session["age"]
+        gender = session["gender"]
+        symptoms = session["symptoms"]
+        duration = session["duration"]
+
+        prompt = f"""
+You are Dr. Aiden — a warm, friendly, and professional doctor having an interview-style consultation.
+The patient is a {age}-year-old {gender} named {name}.
+They have been feeling {symptoms} for {duration}.
+They said: "{user_message}"
+
+Respond like a real doctor — show empathy, analyze the symptoms, suggest likely causes, give simple medication and home care advice.
+
+Include:
+1. Acknowledge their discomfort.
+2. Explain possible causes in simple terms.
+3. Recommend over-the-counter medicines (if safe).
+4. Suggest food, hydration, and rest tips.
+5. Warn when to visit a real doctor.
+6. End with gentle reassurance.
+
 Doctor:"""
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
-
-    gen_config = GenerationConfig(
-        temperature=0.7,
-        top_p=0.92,
-        top_k=50,
-        do_sample=True,
-        max_new_tokens=400,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.1,
-        no_repeat_ngram_size=3,
-    )
+        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_ids = model.generate(**inputs, generation_config=gen_config)
+        with torch.no_grad():
+            output = model.generate(**inputs, generation_config=gen_cfg)
 
-    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+        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.*"
 
-    # Clean and add disclaimer if missing
-    if "⚕️" not in generated_text.lower():
-        generated_text += "\n\n⚕️ *This is AI-generated medical guidance, not a substitute for professional medical care.*"
+        return output_text
 
-    return generated_text.strip()
 
 # =======================================================
-# Gradio UI
+# Gradio Interface
 # =======================================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.HTML("""
-    <div style="background:linear-gradient(135deg,#2b5876,#4e4376);color:white;
-    padding:20px;border-radius:10px;text-align:center;">
-        <h1>🏥 Your Consultation with Dr. Aiden</h1>
-        <p>AI-powered medical chat • Empathetic • Informative • Private</p>
+    <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="👨‍⚕️ Doctor–Patient Chat",
-        height=500,
+        label="👨‍⚕️ Chat with Dr. Aiden",
+        height=550,
+        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
-        ),
-    )
-
-    user_input = gr.Textbox(
-        placeholder="Describe your symptoms (e.g., 'I have a fever for 5 days')",
-        label="💬 Message Dr. Aiden",
-        lines=2,
+            "https://cdn-icons-png.flaticon.com/512/706/706830.png",
+            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
+        )
     )
 
-    send = gr.Button("💬 Send", variant="primary")
-    clear = gr.Button("🧹 New Consultation")
+    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 chat(user_message, history):
-        response = generate_response(user_message)
-        history = history + [[user_message, response]]
+    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 clear_memory():
-        patient_data.update({"name": None, "age": None, "gender": None, "collected": False})
-        return [], "👨‍⚕️ Hello! I’m Dr. Aiden. May I know your name, please?"
+    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)
 
-    send.click(chat, [user_input, chatbot], [user_input, chatbot])
-    user_input.submit(chat, [user_input, chatbot], [user_input, chatbot])
-    clear.click(clear_memory, outputs=[chatbot, user_input])
-
-# =======================================================
-# Launch App
-# =======================================================
 if __name__ == "__main__":
-    print("=" * 60)
-    print("🏥 AI Doctor Chatbot: Dr. Aiden is now online...")
-    print("=" * 60)
+    print("🏥 Launching Dr. Aiden...")
+    demo.queue()
     demo.launch(share=True)