Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
+import time
 
 # =======================================================
 # Load Model
@@ -20,162 +21,152 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 model.eval()
-print("✅ Model loaded successfully!")
-
-
-# =======================================================
-# Global Memory for Doctor Flow
-# =======================================================
-session = {"name": None, "age": None, "gender": None, "stage": "intro"}
 
+print("✅ Model loaded successfully!")
+print(f"📦 Device map: {model.hf_device_map}")
+print(f"💻 Model device: {next(model.parameters()).device}")
 
 # =======================================================
-# Helper: Clean name input
+# Session memory (for patient info)
 # =======================================================
-def extract_name(text):
-    # Remove filler words like "yes", "i am", "my name is"
-    text = text.lower().replace("yes", "").replace("i am", "").replace("i'm", "")
-    text = text.replace("name", "").replace("is", "").replace("it's", "")
-    text = text.strip()
-    # Capitalize nicely
-    return text.title() if text else "Patient"
-
+patient_data = {"name": None, "age": None, "gender": None, "collected": False}
 
 # =======================================================
 # Generate Doctor Response
 # =======================================================
-def doctor_response(user_message):
-    global session
-    user_message = user_message.strip()
-
-    # 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.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["stage"] = "consult"
-        return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}. What brings you in today?"
-
-    # Step 5: Medical Consultation Mode
-    elif session["stage"] == "consult":
-        name = session["name"]
-        age = session["age"]
-        gender = session["gender"]
-
-        prompt = f"""
-You are Dr. Aiden — a friendly, empathetic, and experienced doctor.
-You are chatting casually with a {age}-year-old {gender} patient named {name}.
-The patient is describing their health concern.
-
-Respond like a real doctor would — empathetic, natural, and informative.
+def generate_response(user_message):
+    name, age, gender, collected = (
+        patient_data["name"],
+        patient_data["age"],
+        patient_data["gender"],
+        patient_data["collected"],
+    )
 
-Include in your response:
-- Acknowledge the patient’s concern
-- Possible causes (in simple words)
-- Basic medicine or remedy (if applicable)
-- Rest, hydration, and diet advice
-- When to visit a real doctor
-- End with reassurance
+    # 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?"
 
-Patient: {user_message}
+    # 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"
+        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}"
 Doctor:"""
 
-        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.12,
-            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)
+    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,
+    )
 
-        output_text = tokenizer.decode(output[0], skip_special_tokens=True).strip()
-        output_text = output_text.split("Doctor:")[-1].strip()
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, generation_config=gen_config)
 
-        # Final cleanup
-        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.*"
+    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
 
-        return output_text
+    # 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 generated_text.strip()
 
 # =======================================================
-# Gradio Interface
+# Gradio UI
 # =======================================================
 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>💙 Your Consultation with Dr. Aiden</h1>
-        <p>Empathetic • Knowledgeable • Natural — Your AI Medical Advisor</p>
+    <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>
     """)
 
     chatbot = gr.Chatbot(
-        label="👨‍⚕️ Chat with Dr. Aiden",
-        height=550,
-        type='messages',
+        label="👨‍⚕️ Doctor–Patient Chat",
+        height=500,
         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
+        ),
     )
 
-    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")
+    user_input = gr.Textbox(
+        placeholder="Describe your symptoms (e.g., 'I have a fever for 5 days')",
+        label="💬 Message Dr. Aiden",
+        lines=2,
+    )
 
-    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
+    send = gr.Button("💬 Send", variant="primary")
+    clear = gr.Button("🧹 New Consultation")
 
-    def reset():
-        global session
-        session = {"name": None, "age": None, "gender": None, "stage": "intro"}
-        return []
+    def chat(user_message, history):
+        response = generate_response(user_message)
+        history = history + [[user_message, response]]
+        return "", history
 
-    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)
+    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?"
 
+    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
+# Launch App
 # =======================================================
 if __name__ == "__main__":
-    print("🏥 Launching Dr. Aiden...")
-    demo.queue()
+    print("=" * 60)
+    print("🏥 AI Doctor Chatbot: Dr. Aiden is now online...")
+    print("=" * 60)
     demo.launch(share=True)