Update app.py

app.py

@@ -23,195 +23,152 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 print("✅ Model loaded successfully!")
-print(f"Device map: {model.hf_device_map}")
-print(f"Model device: {next(model.parameters()).device}")
 
 
 # =======================================================
-# Generate Doctor Response (Refined for natural tone)
+# Global Memory for Doctor Flow
 # =======================================================
-def generate_doctor_response(history):
-    user_message = history[-1]["content"]
-
-    if not user_message.strip():
-        history.append({"role": "assistant", "content": "⚠️ Please describe your symptoms or ask a question."})
-        yield history
-        return
-
-    # 🩺 Refined, Doctor-Like Prompt
-    prompt = f"""
-You are Dr. Aiden, a compassionate, calm, and experienced medical doctor. 
-You speak naturally, like in a real consultation, providing medical reasoning and empathy. 
-You should:
-- Greet the patient kindly and acknowledge their concern.
-- Offer a likely cause in simple medical terms.
-- Suggest possible medicines (with safe dosage and common over-the-counter names).
-- Recommend home remedies, foods, and hydration advice.
-- Share short lifestyle or rest tips to aid recovery.
-- End with reassurance and a disclaimer.
-
-Keep your tone friendly yet professional — like an experienced doctor talking directly to the patient.
-Avoid using headings, bullet points, or medical jargon unless necessary.
-Keep your response under 180 words.
-
-Patient says: "{user_message}"
-
-Dr. Aiden:
-"""
-
-    # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
-
-    gen_config = GenerationConfig(
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        max_new_tokens=600,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.15,
-    )
-
-    input_len = inputs["input_ids"].shape[1]
-    with torch.no_grad():
-        output_ids = model.generate(**inputs, generation_config=gen_config)
-
-    generated_ids = output_ids[0][input_len:]
-    response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-
-    # Clean up the response
-    response = clean_medical_response(response)
-
-    # Stream response (simulated)
-    history.append({"role": "assistant", "content": ""})
-    for i in range(0, len(response), 5):
-        history[-1]["content"] = response[:i + 5] + "▌"
-        yield history.copy()
-        time.sleep(0.01)
-
-    history[-1]["content"] = response
-    yield history
+session = {"name": None, "age": None, "gender": None, "stage": "intro"}
 
 
 # =======================================================
-# Clean the response
+# Generate Doctor Response
 # =======================================================
-def clean_medical_response(response: str) -> str:
-    remove_prefixes = ["assistant:", "doctor:", "dr. aiden:", "response:", "patient:"]
-    for p in remove_prefixes:
-        if response.lower().startswith(p):
-            response = response[len(p):].strip()
+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"] = user_message.split()[0].capitalize()
+        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 25 or 30."
+
+    # 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 warm, professional, and conversational doctor talking naturally with a patient.
+
+Patient Info:
+- Name: {name}
+- Age: {age}
+- Gender: {gender}
+
+Speak in a caring and natural tone (like a friendly doctor in a private clinic).
+
+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
+
+Patient: {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=450,
+            repetition_penalty=1.15,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
 
-    response = response.replace("Dr. Aiden:", "").strip()
+        with torch.no_grad():
+            output = model.generate(**inputs, generation_config=gen_cfg)
 
-    # Ensure punctuation
-    if response and response[-1] not in ".!?":
-        response += "."
+        output_text = tokenizer.decode(output[0], skip_special_tokens=True).strip()
+        output_text = output_text.replace("Doctor:", "").replace("Patient:", "").strip()
 
-    # Add disclaimer if missing
-    if "⚕️" not in response and "consult" not in response.lower():
-        response += "\n\n⚕️ *Please note: This is AI-generated medical guidance, not a substitute for a licensed healthcare provider. Always consult a doctor for personal medical care.*"
+        # 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.*"
 
-    return response.strip()
+        return output_text
 
 
 # =======================================================
-# Gradio UI
+# Gradio Interface
 # =======================================================
-with gr.Blocks(theme=gr.themes.Soft(), css="""
-    .medical-header {
-        background: linear-gradient(135deg, #2c3e50 0%, #3498db 100%);
-        padding: 20px;
-        border-radius: 12px;
-        color: white;
-        text-align: center;
-        margin-bottom: 20px;
-        box-shadow: 0 4px 12px rgba(0,0,0,0.15);
-    }
-""") as demo:
-    
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.HTML("""
-        <div class="medical-header">
-            <h1>🏥 Dr. Aiden – AI Medical Consultation</h1>
-            <p>Friendly • Professional • Science-Backed Guidance</p>
-        </div>
+    <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>
     """)
 
     chatbot = gr.Chatbot(
-        label="💬 Your Consultation with Dr. Aiden",
+        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
-        ),
-        height=550,
-        show_copy_button=True
-    )
-
-    with gr.Row():
-        user_input = gr.Textbox(
-            placeholder="Describe your symptoms or ask a question (e.g., 'I have a fever and sore throat for two days')...",
-            label="🧍 Describe Your Symptoms",
-            lines=3,
-            scale=4
+            "https://cdn-icons-png.flaticon.com/512/706/706830.png",
+            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
         )
-
-    with gr.Row():
-        send_btn = gr.Button("💬 Ask Dr. Aiden", variant="primary", size="lg")
-        clear_btn = gr.Button("🧹 New Consultation", size="lg")
-
-    gr.Markdown("### 💡 Example Questions")
-    gr.Examples(
-        examples=[
-            "I have a fever and headache for two days. What should I take?",
-            "I feel tired all day and have trouble sleeping. What could be wrong?",
-            "I have mild chest tightness when I exercise. Should I worry?",
-            "I'm feeling anxious and stressed. Any natural remedies?",
-            "I have stomach pain after eating. What can I do?",
-            "I caught a cold and sore throat. What treatment do you recommend?",
-        ],
-        inputs=user_input,
     )
 
-    # =======================================================
-    # Respond Function (stateless model, persistent chat)
-    # =======================================================
+    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):
-        user_message = message.strip()
-        if not user_message:
-            return "", history
-
-        # Show user input
-        history.append({"role": "user", "content": user_message})
-
-        # Model sees only current input (no memory)
-        temp_history = [{"role": "user", "content": user_message}]
-
-        for updated_history in generate_doctor_response(temp_history):
-            if len(history) == 0 or history[-1]["role"] != "assistant":
-                history.append({"role": "assistant", "content": updated_history[-1]["content"]})
-            else:
-                history[-1]["content"] = updated_history[-1]["content"]
-            yield "", history
-
-    # =======================================================
-    # Button Bindings
-    # =======================================================
+        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, "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(lambda: [], None, chatbot, queue=False)
+    clear_btn.click(reset, None, chatbot, queue=False)
 
 
 # =======================================================
-# Launch App
+# Launch
 # =======================================================
 if __name__ == "__main__":
-    print("="*60)
-    print("🏥 Dr. Aiden – AI Medical Doctor is starting...")
-    print("="*60)
-    demo.queue(max_size=20)
-    demo.launch(
-        share=True,
-        show_error=True,
-        server_name="0.0.0.0"
-    )
+    print("🏥 Launching Dr. Aiden...")
+    demo.queue()
+    demo.launch(share=True)