Update app.py

app.py

@@ -1,60 +1,67 @@
-# app.py
-
+from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import sympy
 
-# Load the pre-trained model and processor outside the function for efficiency
-processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
-model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
+# Use a public model identifier. If you need a private model, remember to authenticate.
+model_name = "google/pix2struct-textcaps-base"
+model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
+processor = Pix2StructProcessor.from_pretrained(model_name)
 
 def solve_math_problem(image):
     try:
-        # Ensure the image is in RGB format
+        # Ensure the image is in RGB format.
         image = image.convert("RGB")
         
-        # Resize and normalize the image as expected by the model
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values
-        
-        # Generate the text (this extracts the handwritten equation)
-        generated_ids = model.generate(pixel_values)
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
-        # Clean and prepare the extracted text
-        problem_text = generated_text.strip().replace(' ', '')
+        # Preprocess the image and text. Note that header_text is omitted as it's not used for non-VQA tasks.
+        inputs = processor(
+            images=[image],
+            text="Solve the following math problem:",
+            return_tensors="pt",
+            max_patches=2048
+        )
         
-        # Use sympy to parse and solve the equation
-        # Handle simple arithmetic and algebraic equations
-        expr = sympy.sympify(problem_text)
-        solution = sympy.solve(expr)
+        # Generate the solution with generation parameters.
+        predictions = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            early_stopping=True,
+            num_beams=4,
+            temperature=0.2
+        )
         
-        # Format the solution for display
-        if isinstance(solution, list):
-            solution = ', '.join([str(s) for s in solution])
-        else:
-            solution = str(solution)
+        # Decode the problem text and generated solution.
+        problem_text = processor.decode(
+            inputs["input_ids"][0],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True
+        )
+        solution = processor.decode(
+            predictions[0],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True
+        )
         
-        return f"**Problem:** {problem_text}\n\n**Solution:** {solution}"
+        return f"Problem: {problem_text}\nSolution: {solution}"
         
     except Exception as e:
-        return f"**Error processing image:** {str(e)}"
+        return f"Error processing image: {str(e)}"
 
-# Create the Gradio interface
+# 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"
+        image_mode="RGB"  # This forces the input to be RGB.
     ),
-    outputs=gr.Markdown(),
+    outputs=gr.Textbox(label="Solution", show_copy_button=True),
     title="Handwritten Math Problem Solver",
-    description="Upload an image of a handwritten math problem, and the app will attempt to solve it.",
+    description="Upload an image of a handwritten math problem (algebra, arithmetic, etc.) and get the solution",
     examples=[
         ["example_addition.png"],
         ["example_algebra.jpg"]
     ],
+    theme="soft",
     allow_flagging="never"
 )