Update app.py

app.py

@@ -2,24 +2,26 @@ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import gradio as gr
 from PIL import Image
 
-# Load the pre-trained Pix2Struct model and processor
+# Load the pre-trained Pix2Struct model and processor.
 model_name = "google/pix2struct-mathqa-base"
 model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
 processor = Pix2StructProcessor.from_pretrained(model_name)
 
 def solve_math_problem(image):
     try:
-        # Preprocess the image
-        image = image.convert("RGB")  # Ensure RGB format
+        # Ensure the image is in RGB format.
+        image = image.convert("RGB")
+        
+        # Preprocess the image and text.
+        # Note: We omit the header_text parameter because this is not a VQA task.
         inputs = processor(
-            images=[image],  # Wrap in list
-            text="Solve the following math problem:",  # More specific prompt
+            images=[image],                           # Provide a list of images.
+            text="Solve the following math problem:", # Prompt text.
             return_tensors="pt",
-            max_patches=2048,  # Increased from default 1024 for better math handling
-            header_text="Math Problem"  # Add header text
+            max_patches=2048                          # Increase the maximum patches for better math handling.
         )
         
-        # Generate the solution
+        # Generate the solution with specified generation parameters.
         predictions = model.generate(
             **inputs,
             max_new_tokens=200,
@@ -28,33 +30,38 @@ def solve_math_problem(image):
             temperature=0.2
         )
         
-        # Decode the output
+        # Decode the input text and the model prediction.
+        # Here, we access "input_ids" via the dictionary key.
+        problem_text = processor.decode(
+            inputs["input_ids"][0],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True
+        )
         solution = processor.decode(
-            predictions[0], 
+            predictions[0],
             skip_special_tokens=True,
             clean_up_tokenization_spaces=True
         )
         
-        # Format the solution
-        return f"Problem: {processor.decode(inputs.input_ids[0])}\nSolution: {solution}"
+        return f"Problem: {problem_text}\nSolution: {solution}"
         
     except Exception as e:
         return f"Error processing image: {str(e)}"
 
-# Gradio interface with explicit image handling
+# Set up the Gradio interface.
 demo = gr.Interface(
     fn=solve_math_problem,
     inputs=gr.Image(
         type="pil", 
         label="Upload Handwritten Math Problem",
-        image_mode="RGB",  # Force RGB format
+        image_mode="RGB",  # Force RGB conversion.
         source="upload"
     ),
     outputs=gr.Textbox(label="Solution", show_copy_button=True),
     title="Handwritten Math Problem Solver",
     description="Upload an image of a handwritten math problem (algebra, arithmetic, etc.) and get the solution",
     examples=[
-        ["example_addition.png"],  # Make sure to upload these files
+        ["example_addition.png"],  # Ensure these example files exist in your working directory.
         ["example_algebra.jpg"]
     ],
     theme="soft",
@@ -62,4 +69,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()