Update app.py

app.py

@@ -28,50 +28,55 @@ print(f"Model device: {next(model.parameters()).device}")
 
 
 # =======================================================
-# Generate Comprehensive Doctor Response
+# Generate Doctor Response - Interactive Medical Consultation
 # =======================================================
 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."})
+        history.append({"role": "assistant", "content": "How can I help you today?"})
         yield history
         return
 
-    # Enhanced Medical Prompt - Conversational Doctor Style
-    prompt = f"""You are a friendly, experienced doctor having a natural conversation with a patient.
-
-Respond in a warm, conversational way covering:
-- What might be causing their symptoms (keep it simple)
-- Specific medications they can take (with dosages)
-- Foods and drinks that will help
-- Lifestyle tips and rest recommendations
-- When they should see a doctor in person
+    # 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:
 
-Important rules:
-- Write like you're talking to a friend, not writing a formal report
-- NO labels like "Doctor:", "Patient:", "Assessment:", "Medications:" etc.
-- Keep it natural and flowing, not structured lists
-- Be warm and reassuring
-- Keep response under 200 words
-- End with a simple 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
 
-Patient says: {user_message}
+Respond naturally as a caring doctor would. Keep responses concise (2-4 sentences). Ask ONE specific follow-up question when you need more information.
 
-Your response:"""
+Previous conversation:
+{conversation_context}
+PATIENT: {user_message}
+DOCTOR:"""
 
     # Tokenize input
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
 
     gen_config = GenerationConfig(
-        temperature=0.8,
-        top_p=0.95,
-        top_k=40,
+        temperature=0.75,
+        top_p=0.92,
+        top_k=45,
         do_sample=True,
-        max_new_tokens=300,
+        max_new_tokens=250,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
-        repetition_penalty=1.18,
+        repetition_penalty=1.2,
         no_repeat_ngram_size=3
     )
 
@@ -84,76 +89,59 @@ Your response:"""
     response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
     # Clean response
-    response = clean_medical_response(response)
+    response = clean_doctor_response(response)
 
     # Stream response token by token
     history.append({"role": "assistant", "content": ""})
-    for i in range(0, len(response), 5):
-        chunk = response[:i + 5]
+    for i in range(0, len(response), 4):
+        chunk = response[:i + 4]
         history[-1]["content"] = chunk + "▌"
         yield history.copy()
-        time.sleep(0.01)
+        time.sleep(0.012)
 
     history[-1]["content"] = response
     yield history
 
 
-def clean_medical_response(response: str) -> str:
-    """Clean and format the medical response for natural conversation."""
-    # Remove common prefixes
-    prefixes = ["assistant:", "doctor:", "response:", "comprehensive medical response:", 
-                "medical response:", "your response:", "dear patient"]
+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:
+    for prefix in prefixes_to_remove:
         if response_lower.startswith(prefix):
             response = response[len(prefix):].strip()
-            response = response.lstrip(',').strip()
             break
     
-    # Remove section headers/labels
-    lines = response.split('\n')
-    cleaned_lines = []
-    skip_patterns = ['doctor:', 'assistant:', 'patient:', 'assessment:', 'medications:', 
-                    'nutrition:', 'lifestyle:', 'follow-up:', '1.', '2.', '3.', '4.', '5.']
-    
-    for line in lines:
-        line_lower = line.lower().strip()
-        should_skip = False
-        for pattern in skip_patterns:
-            if line_lower.startswith(pattern):
-                should_skip = True
-                break
-        if not should_skip and line.strip():
-            cleaned_lines.append(line)
-    
-    response = '\n'.join(cleaned_lines)
-    
-    # Stop at repetitive text or gibberish
-    if 'accordingly' in response.lower() or 'respectively' in response.lower():
-        sentences = response.split('.')
-        good_sentences = []
-        for sent in sentences:
-            if 'accordingly' not in sent.lower() and 'respectively' not in sent.lower():
-                good_sentences.append(sent)
-            else:
-                break
-        response = '. '.join(good_sentences)
+    # 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 length
-    sentences = [s.strip() for s in response.split('.') if s.strip()]
-    if len(sentences) > 8:
-        response = '. '.join(sentences[:8]) + '.'
+    # 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) + '.'
+        response = ' '.join(sentences)
     
-    # Add disclaimer if not present and response is substantial
-    if len(response) > 50:
-        if '⚕️' not in response and 'consult' not in response.lower():
-            response += '\n\n⚕️ *If symptoms worsen or persist beyond 3-4 days, please see a doctor.*'
+    # 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]
     
-    # Fallback for very short responses
-    if len(response.strip()) < 40:
-        response = "I understand you're not feeling well. Could you tell me a bit more about your symptoms? How long have you had them, and have you taken anything for them yet?"
+    # 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?"
     
     return response.strip()
 
@@ -174,11 +162,19 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
     
     gr.HTML("""
         <div class="medical-header">
-            <h1>🏥 AI Doctor Chat</h1>
-            <p>Friendly Medical Advice • Like Talking to Your Doctor</p>
+            <h1>🏥 AI Doctor Consultation</h1>
+            <p>Interactive Medical Conversation • Just Like Visiting Your Doctor</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",
         type='messages',
@@ -186,69 +182,57 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
             "https://cdn-icons-png.flaticon.com/512/706/706830.png",  # Patient
             "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"   # Doctor
         ),
-        height=550,
+        height=500,
         show_copy_button=True
     )
 
     with gr.Row():
         user_input = gr.Textbox(
-            placeholder="Describe your symptoms in detail (e.g., 'I have fever, headache, and body pain for 3 days')...",
-            label="🧍 Describe Your Symptoms",
-            lines=3,
+            placeholder="Describe your symptoms or answer the doctor's questions...",
+            label="🧍 Patient (You)",
+            lines=2,
             scale=4
         )
 
     with gr.Row():
-        send_btn = gr.Button("💬 Consult Doctor", variant="primary", scale=1, size="lg")
-        clear_btn = gr.Button("🧹 New Consultation", scale=1, size="lg")
+        send_btn = gr.Button("💬 Send", variant="primary", scale=1)
+        clear_btn = gr.Button("🔄 New Consultation", scale=1)
 
-    gr.Markdown("### 💡 Example Consultations")
+    gr.Markdown("### 💡 Example Conversations")
     gr.Examples(
         examples=[
-            "I have a fever of 102°F, headache, and body aches for 2 days. What should I do?",
-            "I've been having persistent headaches and feeling tired. Need advice on diet and lifestyle.",
-            "I have acidity and stomach pain after eating. What medications and diet should I follow?",
-            "I'm feeling stressed and anxious. Suggest lifestyle changes and natural remedies.",
-            "I have high blood pressure. What diet and lifestyle changes should I make?",
-            "I caught a cold with sore throat and cough. What treatment do you recommend?",
+            "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"
         ],
         inputs=user_input,
     )
 
     gr.Markdown("""
     ---
-    ### 🩺 What This AI Doctor Provides:
-    
-    ✅ **Medical Assessment** - Initial diagnosis and condition evaluation  
-    ✅ **Medication Recommendations** - Appropriate medicines with dosage guidance  
-    ✅ **Nutrition & Diet Plans** - Foods and nutrients to help recovery  
-    ✅ **Lifestyle Modifications** - Exercise, sleep, stress management tips  
-    ✅ **Follow-up Advice** - When to see a doctor and warning signs  
-    
-    ⚠️ **Important Medical Disclaimer:**  
-    This AI provides general medical information for educational purposes only. It is NOT a substitute for professional medical advice, diagnosis, or treatment. Always seek the advice of your physician or qualified healthcare provider with any questions about a medical condition. Never disregard professional medical advice or delay seeking it because of something you have read here.
-    
-    🚨 **Emergency:** If you are experiencing a medical emergency, call emergency services immediately.
+    ⚠️ **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
+    # Respond Function with Context Memory
     # =======================================================
     def respond(message, history):
         user_message = message.strip()
         if not user_message:
             return "", history
 
-        # Show user message in chat
+        # Add user message to history
         history.append({"role": "user", "content": user_message})
 
-        # Model sees only current message (stateless for consistent behavior)
-        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:
+        # Generate response with full conversation context
+        for updated_history in generate_doctor_response(history):
+            # Update the last assistant message
+            if history[-1]["role"] == "assistant":
                 history[-1]["content"] = updated_history[-1]["content"]
             yield "", history
 
@@ -265,7 +249,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
 # =======================================================
 if __name__ == "__main__":
     print("="*60)
-    print("🏥 AI Medical Doctor Consultation System Starting...")
+    print("🏥 AI Doctor Consultation System Starting...")
+    print("   Interactive medical conversation with context memory")
     print("="*60)
     demo.queue(max_size=20)
     demo.launch(