Update app.py

app.py

@@ -33,8 +33,8 @@ except Exception as e:
 @spaces.GPU
 def symptom_checker_chat(user_input, history_state, new_image_upload, image_state):
     """
-    Manages the conversational flow, persisting the image across turns.
-    Includes robust error reporting directly in the UI for debugging.
+    Manages the conversational flow by manually building the prompt to ensure
+    correct handling of the <image> token.
     """
     if not model_loaded:
         history_state.append((user_input, "Error: The model could not be loaded."))
@@ -42,34 +42,41 @@ def symptom_checker_chat(user_input, history_state, new_image_upload, image_stat
 
     current_image = new_image_upload if new_image_upload is not None else image_state
     
-    if new_image_upload is not None:
-        model_input_text = f"<image>\n{user_input}"
-    else:
-        model_input_text = user_input
-
-    system_prompt = """
-    You are an expert, empathetic AI medical assistant... (your full prompt here)
-    ***Disclaimer: I am an AI assistant and not a medical professional...***
-    """
-
-    conversation = [{"role": "system", "content": system_prompt}]
+    # --- FIX: Manual Prompt Construction ---
+    # This gives us full control and bypasses the opaque apply_chat_template behavior.
+    
+    # System prompt is not included in the turns, but as a prefix.
+    system_prompt = "You are an expert, empathetic AI medical assistant..." # Keep your full system prompt
+    
+    # Build the prompt from history
+    prompt_parts = []
     for turn_input, assistant_output in history_state:
-        conversation.append({"role": "user", "content": turn_input})
+        # Add a user turn from history
+        prompt_parts.append(f"<start_of_turn>user\n{turn_input}<end_of_turn>\n")
+        # Add a model turn from history
         if assistant_output:
-            conversation.append({"role": "assistant", "content": assistant_output})
-    conversation.append({"role": "user", "content": model_input_text})
-        
-    prompt = processor.apply_chat_template(
-        conversation, tokenize=False, add_generation_prompt=True
-    )
+            prompt_parts.append(f"<start_of_turn>model\n{assistant_output}<end_of_turn>\n")
+
+    # Add the current user turn
+    prompt_parts.append("<start_of_turn>user\n")
+    # The MOST IMPORTANT PART: Add the <image> token if an image is present.
+    # We add it for a new upload OR if we're in a conversation that already had an image.
+    if current_image:
+        prompt_parts.append("<image>\n")
+    prompt_parts.append(f"{user_input}<end_of_turn>\n")
+    
+    # Add the generation prompt for the model to start its response
+    prompt_parts.append("<start_of_turn>model\n")
+    
+    # Join everything into a single string
+    final_prompt = "".join(prompt_parts)
 
     try:
-        # Pass the image and text to the processor for encoding
+        # Process the inputs using our manually built prompt
         if current_image:
-            # FIX 2: Ensure the image is passed as a list. This is more robust.
-            inputs = processor(text=prompt, images=[current_image], return_tensors="pt").to(model.device, dtype)
+            inputs = processor(text=final_prompt, images=[current_image], return_tensors="pt").to(model.device, dtype)
         else:
-            inputs = processor(text=prompt, return_tensors="pt").to(model.device, dtype)
+            inputs = processor(text=final_prompt, return_tensors="pt").to(model.device, dtype)
 
         # Generate the response
         outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
@@ -78,21 +85,29 @@ def symptom_checker_chat(user_input, history_state, new_image_upload, image_stat
         clean_response = processor.decode(generated_tokens[0], skip_special_tokens=True).strip()
 
     except Exception as e:
-        # FIX 1: EXPOSE THE REAL ERROR IN THE UI FOR DEBUGGING
-        # This is the most important change. We will now see the true error message.
         print(f"Caught a critical exception during generation: {e}", flush=True)
+        # Display the real error in the UI for easier debugging
         clean_response = (
             "An error occurred during generation. This is the technical details:\n\n"
             f"```\n{type(e).__name__}: {e}\n```"
         )
 
-    # Update history and return values
-    history_state.append((model_input_text, clean_response))
+    # --- History Management ---
+    # For history, we need to save the user_input along with a marker if an image was present
+    # We use the same <image>\n token we've been using as that marker.
+    history_input = user_input
+    if current_image:
+        history_input = f"<image>\n{user_input}"
+
+    history_state.append((history_input, clean_response))
+    
+    # Create display history without the special tokens
     display_history = [(turn.replace("<image>\n", ""), resp) for turn, resp in history_state]
     
+    # Return all updated values
     return display_history, history_state, current_image, None, ""
 
-# --- Gradio Interface (No changes needed here from the last version) ---
+# --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
     gr.Markdown(
         """
@@ -124,6 +139,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}")
     )
     
     def on_submit(user_input, history, new_image, persisted_image):
+        # We need to handle the case where the user input is empty
+        if not user_input.strip():
+            return history, history, persisted_image, None, ""
         return symptom_checker_chat(user_input, history, new_image, persisted_image)
 
     submit_btn.click(