Update app.py

app.py

@@ -13,7 +13,6 @@ processor = None
 model_loaded = False
 
 try:
-    # We load the model and its dedicated processor separately.
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         torch_dtype=torch.bfloat16,
@@ -31,8 +30,8 @@ except Exception as e:
 @spaces.GPU()
 def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
     """
-    Analyzes symptoms by directly using the model and processor with the correct,
-    two-step templating and processing logic.
+    Analyzes symptoms by MANUALLY constructing the prompt string to ensure all special
+    tokens are correctly placed, bypassing the faulty chat template abstraction.
     """
     if not model_loaded:
         return "Error: The AI model could not be loaded. Please check the Space logs."
@@ -42,53 +41,51 @@ def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
         return "Please describe your symptoms or upload an image for analysis."
 
     try:
-        # --- DEFINITIVE PROMPT & INPUT PREPARATION ---
-        
-        # 1. Combine all text inputs into a single string for the user's turn.
-        # Add the <image> placeholder only if an image is provided.
+        # --- DEFINITIVE MANUAL PROMPT CONSTRUCTION ---
+
         system_instruction = (
             "You are an expert, empathetic AI medical assistant. "
             "Analyze the potential medical condition based on the following information. "
             "Provide a list of possible conditions, your reasoning, and a clear, actionable next-steps plan."
         )
-        
-        user_content = ""
-        if symptom_image:
-            # The model expects the <image> token to know where to place the image.
-            user_content += "<image>\n"
-        
-        # Combine user text and system instructions for the user's message.
-        user_content += f"{symptoms_text}\n\n{system_instruction}"
 
-        messages = [
-            {"role": "user", "content": user_content}
-        ]
+        # 1. Manually build the prompt string as a list of parts.
+        prompt_parts = ["<start_of_turn>user"]
 
-        # 2. Use the tokenizer to apply the model's specific chat template.
-        # This correctly formats the text with all required special tokens.
-        prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        # 2. CRUCIAL: Add the <image> placeholder *only* if an image exists.
+        if symptom_image:
+            prompt_parts.append("<image>")
+        
+        # 3. Add all text content.
+        prompt_parts.append(f"{symptoms_text}\n\n{system_instruction}")
+        
+        # 4. Signal the start of the model's turn.
+        prompt_parts.append("<start_of_turn>model")
+        
+        # 5. Join all parts into a single string. This is our final prompt.
+        prompt = "\n".join(prompt_parts)
 
-        # 3. Use the main processor to prepare the final model inputs.
-        # It takes the formatted text and the PIL image and creates the tensors.
+        # 6. Use the processor with our manually built prompt. It will now find the <image>
+        # token and correctly process the associated image object.
         inputs = processor(
             text=prompt,
-            images=symptom_image,  # This can be None if no image is provided
+            images=symptom_image,  # This will be None for text-only, which is now handled correctly.
             return_tensors="pt"
         ).to(model.device)
 
-        # 4. Generation parameters
+        # 7. Generation parameters
         generate_kwargs = {
             "max_new_tokens": 512,
             "do_sample": True,
             "temperature": 0.7,
         }
 
-        print("Generating model output directly...")
+        print("Generating model output with manually constructed prompt...")
         
-        # 5. Generate the response
+        # 8. Generate the response
         generate_ids = model.generate(**inputs, **generate_kwargs)
 
-        # 6. Decode only the newly generated tokens back into a string.
+        # 9. Decode only the newly generated tokens. This logic is correct.
         input_token_len = inputs["input_ids"].shape[-1]
         result = processor.batch_decode(generate_ids[:, input_token_len:], skip_special_tokens=True)[0]
         
@@ -103,7 +100,7 @@ def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
 # --- Gradio Interface (No changes needed) ---
 with gr.Blocks(theme=gr.themes.Soft(), title="AI Symptom Analyzer") as demo:
     gr.HTML("""
-        <div style="text-align: center; background: linear-gradient(135deg, #66eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 10px; margin-bottom: 2rem;">
+        <div style="text-align: center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 10px; margin-bottom: 2rem;">
             <h1>🩺 AI Symptom Analyzer</h1>
             <p>Advanced symptom analysis powered by Google's MedGemma AI</p>
         </div>