Update app.py

app.py

@@ -22,222 +22,180 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-print("Model loaded successfully!")
+print("✅ Model loaded successfully!")
 print(f"Device map: {model.hf_device_map}")
 print(f"Model device: {next(model.parameters()).device}")
 
 
 # =======================================================
-# Generate Doctor Response - Interactive Medical Consultation
+# Generate Doctor Response (Refined for natural tone)
 # =======================================================
 def generate_doctor_response(history):
     user_message = history[-1]["content"]
 
     if not user_message.strip():
-        history.append({"role": "assistant", "content": "How can I help you today?"})
+        history.append({"role": "assistant", "content": "⚠️ Please describe your symptoms or ask a question."})
         yield history
         return
 
-    # Build conversation context from history
-    conversation_context = ""
-    if len(history) > 1:
-        # Include previous exchanges for context
-        for msg in history[:-1]:
-            if msg["role"] == "user":
-                conversation_context += f"PATIENT: {msg['content']}\n"
-            else:
-                conversation_context += f"DOCTOR: {msg['content']}\n"
-    
-    # Medical conversation prompt - like real doctor-patient interaction
-    prompt = f"""You are an experienced medical doctor conducting a patient consultation. Have a natural, interactive conversation where you:
+    # 🩺 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.
 
-- Ask relevant follow-up questions to understand symptoms better
-- Gather medical history (medications, lifestyle, family history)
-- Provide medical assessment and recommendations
-- Suggest medications with dosages when appropriate
-- Give diet and lifestyle advice
-- Explain what tests or next steps are needed
+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.
 
-Respond naturally as a caring doctor would. Keep responses concise (2-4 sentences). Ask ONE specific follow-up question when you need more information.
+Patient says: "{user_message}"
 
-Previous conversation:
-{conversation_context}
-PATIENT: {user_message}
-DOCTOR:"""
+Dr. Aiden:
+"""
 
     # Tokenize input
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
 
     gen_config = GenerationConfig(
-        temperature=0.75,
-        top_p=0.92,
-        top_k=45,
+        temperature=0.7,
+        top_p=0.9,
         do_sample=True,
-        max_new_tokens=250,
+        max_new_tokens=600,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.2,
-        no_repeat_ngram_size=3
+        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 response
-    response = clean_doctor_response(response)
+    # Clean up the response
+    response = clean_medical_response(response)
 
-    # Stream response token by token
+    # Stream response (simulated)
     history.append({"role": "assistant", "content": ""})
-    for i in range(0, len(response), 4):
-        chunk = response[:i + 4]
-        history[-1]["content"] = chunk + "▌"
+    for i in range(0, len(response), 5):
+        history[-1]["content"] = response[:i + 5] + "▌"
         yield history.copy()
-        time.sleep(0.012)
+        time.sleep(0.01)
 
     history[-1]["content"] = response
     yield history
 
 
-def clean_doctor_response(response: str) -> str:
-    """Clean the doctor's response to be natural and conversational."""
-    
-    # Remove role labels if present
-    prefixes_to_remove = ["doctor:", "assistant:", "response:", "patient:"]
-    response_lower = response.lower()
-    for prefix in prefixes_to_remove:
-        if response_lower.startswith(prefix):
-            response = response[len(prefix):].strip()
-            break
-    
-    # Stop at repetitive patterns or gibberish
-    stop_phrases = ["accordingly", "respectively", "speaking correctly", "faithfully yours"]
-    for phrase in stop_phrases:
-        if phrase in response.lower():
-            # Find first occurrence and cut there
-            idx = response.lower().find(phrase)
-            response = response[:idx].strip()
-            break
-    
-    # Limit to reasonable number of sentences (4-6 max)
-    sentences = [s.strip() + '.' for s in response.split('.') if s.strip()]
-    if len(sentences) > 6:
-        response = ' '.join(sentences[:6])
-    else:
-        response = ' '.join(sentences)
-    
-    # Remove incomplete sentences at the end
-    if response and response[-1] not in '.!?':
-        last_period = response.rfind('.')
-        if last_period > 0:
-            response = response[:last_period + 1]
-    
-    # Clean up extra spaces
-    response = ' '.join(response.split())
-    
-    # Fallback for very short or empty responses
-    if len(response.strip()) < 20:
-        response = "Could you tell me more about your symptoms? When did they start?"
-    
+# =======================================================
+# Clean the 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()
+
+    response = response.replace("Dr. Aiden:", "").strip()
+
+    # Ensure punctuation
+    if response and response[-1] not in ".!?":
+        response += "."
+
+    # 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.*"
+
     return response.strip()
 
 
 # =======================================================
-# Gradio Interface
+# Gradio UI
 # =======================================================
 with gr.Blocks(theme=gr.themes.Soft(), css="""
     .medical-header {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        background: linear-gradient(135deg, #2c3e50 0%, #3498db 100%);
         padding: 20px;
-        border-radius: 10px;
+        border-radius: 12px;
         color: white;
         text-align: center;
         margin-bottom: 20px;
+        box-shadow: 0 4px 12px rgba(0,0,0,0.15);
     }
 """) as demo:
     
     gr.HTML("""
         <div class="medical-header">
-            <h1>🏥 AI Doctor Consultation</h1>
-            <p>Interactive Medical Conversation • Just Like Visiting Your Doctor</p>
+            <h1>🏥 Dr. Aiden – AI Medical Consultation</h1>
+            <p>Friendly • Professional • Science-Backed Guidance</p>
         </div>
     """)
 
-    gr.Markdown("""
-    ### 💬 How This Works:
-    - Describe your symptoms or health concerns
-    - The AI doctor will ask questions to understand your condition
-    - You'll get medical advice, medication suggestions, and lifestyle recommendations
-    - Have a natural back-and-forth conversation just like a real doctor visit
-    """)
-
     chatbot = gr.Chatbot(
-        label="💬 Doctor-Patient Consultation",
+        label="💬 Your Consultation with Dr. Aiden",
         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=500,
+        height=550,
         show_copy_button=True
     )
 
     with gr.Row():
         user_input = gr.Textbox(
-            placeholder="Describe your symptoms or answer the doctor's questions...",
-            label="🧍 Patient (You)",
-            lines=2,
+            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
         )
 
     with gr.Row():
-        send_btn = gr.Button("💬 Send", variant="primary", scale=1)
-        clear_btn = gr.Button("🔄 New Consultation", scale=1)
+        send_btn = gr.Button("💬 Ask Dr. Aiden", variant="primary", size="lg")
+        clear_btn = gr.Button("🧹 New Consultation", size="lg")
 
-    gr.Markdown("### 💡 Example Conversations")
+    gr.Markdown("### 💡 Example Questions")
     gr.Examples(
         examples=[
-            "Hi Doctor, I've been having fever and body aches for 2 days",
-            "I have numbness in my toes and difficulty walking",
-            "I've been feeling very tired all the time lately",
-            "I have chest pain and shortness of breath",
-            "I get headaches almost every day",
-            "I have stomach pain after eating"
+            "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,
     )
 
-    gr.Markdown("""
-    ---
-    ⚠️ **Medical Disclaimer:** This AI provides general medical information for educational purposes. 
-    It is NOT a substitute for professional medical advice. Always consult a qualified healthcare 
-    provider for diagnosis and treatment. In case of emergency, call emergency services immediately.
-    """)
-
     # =======================================================
-    # Respond Function with Context Memory
+    # Respond Function (stateless model, persistent chat)
     # =======================================================
     def respond(message, history):
         user_message = message.strip()
         if not user_message:
             return "", history
 
-        # Add user message to history
+        # Show user input
         history.append({"role": "user", "content": user_message})
 
-        # Generate response with full conversation context
-        for updated_history in generate_doctor_response(history):
-            # Update the last assistant message
-            if history[-1]["role"] == "assistant":
+        # 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 & Input Bindings
+    # Button Bindings
     # =======================================================
     send_btn.click(respond, [user_input, chatbot], [user_input, chatbot])
     user_input.submit(respond, [user_input, chatbot], [user_input, chatbot])
@@ -249,12 +207,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
 # =======================================================
 if __name__ == "__main__":
     print("="*60)
-    print("🏥 AI Doctor Consultation System Starting...")
-    print("   Interactive medical conversation with context memory")
+    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"
-    )
+    )