Update app.py

app.py

@@ -3,11 +3,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
 import time
 
-# =======================================================
-# Global session state for multi-step questioning
-# =======================================================
-session_answers = {}
-
 # =======================================================
 # Load Model
 # =======================================================
@@ -32,10 +27,9 @@ print(f"Device map: {model.hf_device_map}")
 print(f"Model device: {next(model.parameters()).device}")
 
 # =======================================================
-# Generate Doctor Response
+# Generate Doctor Response (Stateless)
 # =======================================================
 def generate_doctor_response(history):
-    global session_answers
     user_message = history[-1]["content"]
 
     if not user_message.strip():
@@ -43,41 +37,34 @@ def generate_doctor_response(history):
         yield history
         return
 
-    # Build prompt with context
+    # Single-turn prompt (no past context)
     prompt = """# 🩺 You are a highly knowledgeable Medical Expert
 
-        **Roles:**
-        - **Doctor:** Diagnose symptoms, ask relevant follow-up questions, and provide medical advice.
-        - **Nutritionist:** Give diet and lifestyle recommendations.
-        - **Medical Teacher:** Explain complex medical terms and conditions in simple, understandable language.
-        
-        **Guidelines:**
-        1. Ask **only one follow-up question at a time** to gather necessary information.
-        2. Provide advice only after collecting enough information.
-        3. Be **empathetic, professional, and conversational**.
-        4. Provide disclaimers when needed:  
-           ⚕️ *This is AI-generated information and not a substitute for professional medical advice. Please consult a healthcare provider for proper diagnosis and treatment.*
-        5. Keep answers concise, clear, and actionable.
-        
-        Use this expertise to respond naturally to any patient message, balancing teaching, advice, and medical guidance.
-        """
-
-    recent_history = history[-10:-1] if len(history) > 10 else history[:-1]
-    for msg in recent_history:
-        role = "Patient" if msg["role"] == "user" else "Doctor"
-        content = msg['content'].replace("⚕️ *Note: This is AI-generated information*", "").strip()
-        prompt += f"{role}: {content}\n"
-    prompt += f"Patient: {user_message}\nDoctor:"
+**Roles:**
+- **Doctor:** Diagnose symptoms, ask relevant follow-up questions, and provide medical advice.
+- **Nutritionist:** Give diet and lifestyle recommendations.
+- **Medical Teacher:** Explain complex medical terms and conditions in simple, understandable language.
+
+**Guidelines:**
+1. Ask **only one follow-up question at a time** if necessary.
+2. Provide advice after collecting enough information.
+3. Be **empathetic, professional, and conversational**.
+4. Include disclaimers when needed:  
+   ⚕️ *This is AI-generated information and not a substitute for professional medical advice. Please consult a healthcare provider for proper diagnosis and treatment.*
+5. Keep answers concise, clear, and actionable.
+
+Respond naturally to the following patient message:
+
+Patient: """ + user_message + "\nDoctor:"
 
     # Tokenize input
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    # Generation configuration for concise, interactive answers
     gen_config = GenerationConfig(
         temperature=0.7,
         top_p=0.9,
         do_sample=True,
-        max_new_tokens=500,  # short answers
+        max_new_tokens=500,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
         repetition_penalty=1.2
@@ -91,27 +78,25 @@ def generate_doctor_response(history):
     generated_ids = output_ids[0][input_len:]
     response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
-    # Take only first 2-3 sentences to make it concise
+    # Keep concise output
     response = ". ".join(response.split(". ")[:3]).strip()
     if response.lower().startswith("doctor:"):
         response = response[7:].strip()
     if len(response) < 10:
         response = "I understand your concern. Could you please provide more details about your symptoms?"
 
-    # Add assistant placeholder for streaming
-    history.append({"role": "assistant", "content": ""})
-
     # Stream response token by token
+    history.append({"role": "assistant", "content": ""})
     for i in range(0, len(response), 4):
-        chunk = response[:i+4]
+        chunk = response[:i + 4]
         history[-1]["content"] = chunk + "▌"
         yield history.copy()
         time.sleep(0.015)
 
-    # Final response
     history[-1]["content"] = response
     yield history
 
+
 # =======================================================
 # Gradio Interface
 # =======================================================
@@ -151,20 +136,38 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         label="💡 Example Questions"
     )
 
+    # =======================================================
+    # Respond Function — Model forgets, Chat UI remembers
+    # =======================================================
     def respond(message, history):
-        if history is None:
-            history = []
-        if not message.strip():
+        user_message = message.strip()
+        if not user_message:
             return "", history
-        history.append({"role": "user", "content": message})
-        for updated_history in generate_doctor_response(history):
-            yield "", updated_history
 
+        # Show user message in chat
+        history.append({"role": "user", "content": user_message})
+
+        # Model sees only current message (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
+    # =======================================================
     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)
 
-# Launch
+
+# =======================================================
+# Launch App
+# =======================================================
 if __name__ == "__main__":
     demo.queue()
     demo.launch(share=True)