Update app.py

GPU Request: Added import spaces and the @spaces.GPU decorator to symptom_checker_chat.
State Management:
The click and submit events now use chat_history as an input and output.
The symptom_checker_chat function accepts history from the state and returns the updated list to both the chatbot and chat_history.
Robust Parsing: Replaced the fragile rfind() logic with a much more reliable method of decoding only the newly generated tokens.
UI Cleanup:
Added text_box to the outputs of the event handlers.
The function now returns "" as its last value to clear the textbox after each submission.

app.py

@@ -3,8 +3,10 @@ import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
 from PIL import Image
 import os
+import spaces # <-- FIX 1: IMPORT SPACES
 
 # Get the Hugging Face token from the environment variables
+# Make sure to set this as a "Secret" in your Hugging Face Space settings
 hf_token = os.environ.get("HF_TOKEN")
 
 # Initialize the processor and model
@@ -16,7 +18,8 @@ model_id = "google/medgemma-4b-it"
 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
     dtype = torch.bfloat16
 else:
-    dtype = torch.float32
+    # Fallback to float16 if bfloat16 is not available
+    dtype = torch.float16
 
 model_loaded = False
 # Load the processor and model from Hugging Face
@@ -34,20 +37,20 @@ except Exception as e:
     print(f"Error loading model: {e}")
     # We will display an error in the UI if the model fails to load.
 
-
 # This is the core function for the chatbot
+@spaces.GPU # <-- FIX 1: ADD THE GPU DECORATOR
 def symptom_checker_chat(user_input, history, image_input):
     """
     Manages the conversational flow for the symptom checker.
     """
     if not model_loaded:
         history.append((user_input, "Error: The model could not be loaded. Please check the Hugging Face Space logs."))
-        return history, None
+        # <-- FIX 3 & 4: Return values match new outputs
+        return history, history, None, ""
 
     # System prompt to guide the model's behavior
     system_prompt = """
     You are an expert, empathetic AI medical assistant. Your role is to analyze a user's symptoms and provide a helpful, safe, and informative response.
-
     Here is your workflow:
     1.  Analyze the user's initial input, which may include text and an image.
     2.  If the information is insufficient, ask specific, relevant clarifying questions to better understand the symptoms (e.g., "How long have you had this symptom?", "Can you describe the pain? Is it sharp or dull?").
@@ -55,7 +58,6 @@ def symptom_checker_chat(user_input, history, image_input):
     4.  For each possible condition, briefly explain why it might be relevant.
     5.  Provide a clear, actionable plan, such as "It would be best to monitor your symptoms," or "You should consider consulting a healthcare professional."
     6.  **Crucially, you must ALWAYS end every single response with the following disclaimer, formatted exactly like this, on a new line:**
-
     ***Disclaimer: I am an AI assistant and not a medical professional. This is not a diagnosis. Please consult a doctor for any health concerns.***
     """
     
@@ -63,15 +65,20 @@ def symptom_checker_chat(user_input, history, image_input):
     conversation = [{"role": "system", "content": system_prompt}]
     for user, assistant in history:
         conversation.append({"role": "user", "content": user})
-        conversation.append({"role": "assistant", "content": assistant})
-    
-    # Add the current user input
-    conversation.append({"role": "user", "content": user_input})
-    
+        if assistant: # Ensure assistant message is not None
+            conversation.append({"role": "assistant", "content": assistant})
+            
+    # Add the current user input with a special image token if an image is present
+    if image_input:
+        # MedGemma expects the text to start with <image> token if an image is provided
+        conversation.append({"role": "user", "content": f"<image>\n{user_input}"})
+    else:
+        conversation.append({"role": "user", "content": user_input})
+        
     # Apply the chat template
     prompt = processor.tokenizer.apply_chat_template(
-        conversation, 
-        tokenize=False, 
+        conversation,
+        tokenize=False,
         add_generation_prompt=True
     )
 
@@ -85,16 +92,12 @@ def symptom_checker_chat(user_input, history, image_input):
     # Generate the output from the model
     try:
         outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
-        response_text = processor.decode(outputs[0], skip_special_tokens=True)
         
-        # Clean the response to only get the assistant's part
-        # This logic finds the last assistant message in the generated text
-        last_assistant_marker = "assistant\n"
-        last_occurrence = response_text.rfind(last_assistant_marker)
-        if last_occurrence != -1:
-            clean_response = response_text[last_occurrence + len(last_assistant_marker):].strip()
-        else:
-            clean_response = "I'm sorry, I encountered an issue processing your request. Please try again."
+        # <-- FIX 2: ROBUST RESPONSE PARSING
+        # Decode only the newly generated tokens, not the whole conversation
+        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}")
@@ -103,8 +106,8 @@ def symptom_checker_chat(user_input, history, image_input):
     # Update the history
     history.append((user_input, clean_response))
     
-    return history, None # Return updated history and clear the image input
-
+    # <-- FIX 3 & 4: Return values to update state, clear image box, and clear text box
+    return history, history, None, ""
 
 # Create the Gradio Interface using Blocks for more control
 with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
@@ -116,10 +119,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}")
     )
     
     # Chatbot component to display the conversation
-    chatbot = gr.Chatbot(label="Conversation", height=500)
+    chatbot = gr.Chatbot(label="Conversation", height=500, avatar_images=("user.png", "bot.png")) # Added avatars for fun
     
     # State to store the conversation history
-    chat_history = gr.State([])
+    chat_history = gr.State([]) # <-- FIX 3: This state will now be used correctly
 
     with gr.Row():
         # Image input
@@ -137,27 +140,30 @@ with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}")
 
     # Function to clear all inputs
     def clear_all():
-        return [], None, "", None
+        return [], [], None, "" # <-- FIX 3: Correctly clear the state and chatbot
 
     # Clear button
     clear_btn = gr.Button("Start New Conversation")
-    clear_btn.click(clear_all, outputs=[chatbot, image_box, text_box, chat_history], queue=False)
-    
+    # <-- FIX 3: The outputs list now correctly targets the state
+    clear_btn.click(clear_all, outputs=[chatbot, chat_history, image_box, text_box], queue=False)
+        
     # Define what happens when the user submits
     submit_btn.click(
         fn=symptom_checker_chat,
-        inputs=[text_box, chatbot, image_box],
-        outputs=[chatbot, image_box] # Also clear the image box after submission
+        # <-- FIX 3 & 4: Corrected inputs and outputs
+        inputs=[text_box, chat_history, image_box],
+        outputs=[chatbot, chat_history, image_box, text_box]
     )
-    
+        
     # Define what happens when the user just presses Enter in the textbox
     text_box.submit(
         fn=symptom_checker_chat,
-        inputs=[text_box, chatbot, image_box],
-        outputs=[chatbot, image_box]
+        # <-- FIX 3 & 4: Corrected inputs and outputs
+        inputs=[text_box, chat_history, image_box],
+        outputs=[chatbot, chat_history, image_box, text_box]
     )
 
-
 # Launch the Gradio app
 if __name__ == "__main__":
     demo.launch(debug=True) # Debug mode for more detailed logs
+