Update app.py

app.py

@@ -34,34 +34,24 @@ except Exception as e:
 def symptom_checker_chat(user_input, history_state, new_image_upload, image_state):
     """
     Manages the conversational flow, persisting the image across turns.
-
-    Args:
-        user_input (str): Text from the user.
-        history_state (list): Stateful conversation history.
-        new_image_upload (PIL.Image): A new image uploaded in this turn.
-        image_state (PIL.Image): The image persisted from a previous turn.
-
-    Returns:
-        tuple: Updated values for all relevant Gradio components.
+    Includes robust error reporting directly in the UI for debugging.
     """
     if not model_loaded:
         history_state.append((user_input, "Error: The model could not be loaded."))
         return history_state, history_state, None, None, ""
 
-    # FIX: Determine which image to use. A new upload takes precedence.
-    # This is the key to solving the "image amnesia" problem.
     current_image = new_image_upload if new_image_upload is not None else image_state
     
-    # If this is the *first* turn with an image, add the <image> token.
-    # Don't add it again if the image is just being carried over in the 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..." # Keep your detailed prompt
+    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...***
+    """
 
-    # Construct the full conversation history
     conversation = [{"role": "system", "content": system_prompt}]
     for turn_input, assistant_output in history_state:
         conversation.append({"role": "user", "content": turn_input})
@@ -73,41 +63,44 @@ def symptom_checker_chat(user_input, history_state, new_image_upload, image_stat
         conversation, tokenize=False, add_generation_prompt=True
     )
 
-    # Process inputs. Crucially, pass `current_image` to the processor.
     try:
+        # Pass the image and text to the processor for encoding
         if current_image:
-            inputs = processor(text=prompt, images=current_image, return_tensors="pt").to(model.device, dtype)
+            # 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)
         else:
             inputs = processor(text=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)
         input_token_len = inputs["input_ids"].shape[1]
         generated_tokens = outputs[:, input_token_len:]
         clean_response = processor.decode(generated_tokens[0], skip_special_tokens=True).strip()
 
     except Exception as e:
-        print(f"Error during model generation: {e}")
-        clean_response = "An error occurred during generation. Please check the logs."
-
-    # Update the text history with the model-aware input
+        # 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)
+        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))
-    
-    # Create the display history without the special tokens
     display_history = [(turn.replace("<image>\n", ""), resp) for turn, resp in history_state]
     
-    # Return everything:
-    # 1. display_history -> to the chatbot UI
-    # 2. history_state -> to the text history state
-    # 3. current_image -> to the image state to be persisted
-    # 4. None -> to the image upload box to clear it
-    # 5. "" -> to the text box to clear it
     return display_history, history_state, current_image, None, ""
 
-# --- Gradio Interface ---
+# --- Gradio Interface (No changes needed here from the last version) ---
 with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
-    gr.Markdown("# AI Symptom Checker powered by MedGemma\n...") # Keep your intro markdown
+    gr.Markdown(
+        """
+        # AI Symptom Checker powered by MedGemma
+        Describe your symptoms in the text box below. You can also upload an image (e.g., a skin rash). The AI assistant will ask clarifying questions before suggesting possible conditions and an action plan.
+        """
+    )
     
-    # FIX: Add a new state component to hold the image across turns
     image_state = gr.State(None)
 
     chatbot = gr.Chatbot(label="Conversation", height=500, bubble_full_width=False)
@@ -117,10 +110,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}")
         image_box = gr.Image(type="pil", label="Upload Image of Symptom (Optional)")
 
     with gr.Row():
-        text_box = gr.Textbox(label="Describe your symptoms...", scale=4)
+        text_box = gr.Textbox(label="Describe your symptoms...", placeholder="e.g., I have a rash on my arm...", scale=4)
         submit_btn = gr.Button("Send", variant="primary", scale=1)
 
-    # FIX: Update the clear function to also clear the new image_state
     def clear_all():
         return [], [], None, None, ""
 
@@ -131,11 +123,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}")
         queue=False
     )
     
-    # Combine submit actions into a single function for DRY principle
     def on_submit(user_input, history, new_image, persisted_image):
         return symptom_checker_chat(user_input, history, new_image, persisted_image)
 
-    # FIX: Update inputs and outputs to include the new image_state
     submit_btn.click(
         fn=on_submit,
         inputs=[text_box, chat_history, image_box, image_state],