Update app.py

app.py

@@ -10,36 +10,32 @@ model_name = "google/matcha-base"
 model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
 processor = Pix2StructProcessor.from_pretrained(model_name)
 
-# Move model to GPU if available for faster inference
+# Move model to GPU if available and set to evaluation mode
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
+model.eval()
 
 def solve_math_problem(image):
-    # Preprocess the image and include a clear prompt.
-    # You can adjust the prompt to better match your task if needed.
+    # Preprocess the image and include a prompt.
     inputs = processor(images=image, text="Solve the math problem:", return_tensors="pt")
-    # Ensure the tensors are on the same device as the model
+    # Move all tensors to the same device as the model
     inputs = {key: value.to(device) for key, value in inputs.items()}
     
-    # 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
-    )
+    # Generate the solution using beam search within a no_grad context
+    with torch.no_grad():
+        predictions = model.generate(
+            **inputs,
+            max_new_tokens=150,  # Increase this if longer answers are needed
+            num_beams=5,         # Beam search for more stable outputs
+            early_stopping=True,
+            temperature=0.5      # Lower temperature for more deterministic output
+        )
     
-    # Decode the output to get a string answer, skipping any special tokens.
+    # Decode the generated tokens to a string, skipping special tokens
     solution = processor.decode(predictions[0], skip_special_tokens=True)
     return solution
 
-# Set up a Gradio interface
+# Set up the Gradio interface
 demo = gr.Interface(
     fn=solve_math_problem,
     inputs=gr.Image(type="pil", label="Upload Handwritten Math Problem"),
@@ -51,3 +47,4 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch()
+