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,7 +20,6 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 model.eval()
-
 print("✅ Model loaded successfully!")
 
 
@@ -31,6 +29,18 @@ print("✅ Model loaded successfully!")
 session = {"name": None, "age": None, "gender": None, "stage": "intro"}
 
 
+# =======================================================
+# Helper: Clean name input
+# =======================================================
+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"
+
+
 # =======================================================
 # Generate Doctor Response
 # =======================================================
@@ -45,7 +55,7 @@ def doctor_response(user_message):
 
     # Step 2: Get Name
     elif session["stage"] == "ask_name":
-        session["name"] = user_message.split()[0].capitalize()
+        session["name"] = extract_name(user_message)
         session["stage"] = "ask_age"
         return f"Nice to meet you, {session['name']}! How old are you?"
 
@@ -57,7 +67,7 @@ def doctor_response(user_message):
                 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 25 or 30."
+        return "Please tell me your age in numbers, like 20 or 25."
 
     # Step 4: Get Gender
     elif session["stage"] == "ask_gender":
@@ -78,22 +88,19 @@ def doctor_response(user_message):
         gender = session["gender"]
 
         prompt = f"""
-You are Dr. Aiden — a warm, professional, and conversational doctor talking naturally with a patient.
-
-Patient Info:
-- Name: {name}
-- Age: {age}
-- Gender: {gender}
+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.
 
-Speak in a caring and natural tone (like a friendly doctor in a private clinic).
+Respond like a real doctor would — empathetic, natural, and informative.
 
 Include in your response:
-1. Acknowledgement of their symptoms
-2. Possible causes (simple explanation)
-3. Simple medicines with dosage (if applicable)
-4. Food, rest, and hydration advice
-5. When to see a real doctor
-6. Short closing reassurance
+- 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
 
 Patient: {user_message}
 Doctor:"""
@@ -102,8 +109,8 @@ Doctor:"""
         gen_cfg = GenerationConfig(
             temperature=0.7,
             top_p=0.9,
-            max_new_tokens=450,
-            repetition_penalty=1.15,
+            max_new_tokens=350,
+            repetition_penalty=1.12,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id
         )
@@ -112,7 +119,7 @@ Doctor:"""
             output = model.generate(**inputs, generation_config=gen_cfg)
 
         output_text = tokenizer.decode(output[0], skip_special_tokens=True).strip()
-        output_text = output_text.replace("Doctor:", "").replace("Patient:", "").strip()
+        output_text = output_text.split("Doctor:")[-1].strip()
 
         # Final cleanup
         if not output_text.endswith((".", "!", "?")):