Update app.py

app.py

@@ -1,38 +1,51 @@
+import torch
 from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 import gradio as gr
 from PIL import Image
 
-# Use a valid model identifier.
-# Replace "google/matcha-base" with your checkpoint if you have one.
+# Use a valid model identifier. Here we use "google/matcha-base".
 model_name = "google/matcha-base"
 
 # Load the pre-trained Pix2Struct model and processor
 model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
 processor = Pix2StructProcessor.from_pretrained(model_name)
 
-# Function to solve handwritten math problems
+# Move model to GPU if available for faster inference
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 def solve_math_problem(image):
-    # Preprocess the image: here we render a prompt asking the model to solve the problem.
+    # Preprocess the image and include a clear prompt.
+    # You can adjust the prompt to better match your task if needed.
     inputs = processor(images=image, text="Solve the math problem:", return_tensors="pt")
+    # Ensure the tensors are on the same device as the model
+    inputs = {key: value.to(device) for key, value in inputs.items()}
     
-    # Generate the solution
-    predictions = model.generate(**inputs, max_new_tokens=100)
+    # Generate the solution using beam search.
+    # Adjust parameters for best performance:
+    # - max_new_tokens: Allows longer responses.
+    # - num_beams: Uses beam search to explore multiple hypotheses.
+    # - early_stopping: Stops decoding once a complete answer is generated.
+    # - temperature: Controls randomness (lower value = more deterministic).
+    predictions = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        num_beams=5,
+        early_stopping=True,
+        temperature=0.5
+    )
     
-    # Decode the output
+    # Decode the output to get a string answer, skipping any special tokens.
     solution = processor.decode(predictions[0], skip_special_tokens=True)
     return solution
 
-# Gradio interface
+# Set up a Gradio interface
 demo = gr.Interface(
     fn=solve_math_problem,
     inputs=gr.Image(type="pil", label="Upload Handwritten Math Problem"),
     outputs=gr.Textbox(label="Solution"),
     title="Handwritten Math Problem Solver",
-    description="Upload an image of a handwritten math problem, and the model will attempt to solve it.",
-    examples=[
-        ["example1.jpg"],  # Add example images if available
-        ["example2.jpg"]
-    ],
+    description="Upload an image of a handwritten math problem and the model will attempt to solve it.",
     theme="soft"
 )