Update app.py

app.py

@@ -11,12 +11,22 @@ model = Pix2StructForConditionalGeneration.from_pretrained("google/deplot")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+# def extract_chart_data(img: Image.Image):
+#     inputs = processor(images=img, return_tensors="pt").to(device)
+#     outputs = model.generate(**inputs)
+#     result = processor.decode(outputs[0], skip_special_tokens=True)
+#     return result
+
 def extract_chart_data(img: Image.Image):
-    inputs = processor(images=img, return_tensors="pt").to(device)
+    # Provide a header prompt for VQA
+    header_text = "Extract chart data"
+    
+    inputs = processor(images=img, text=header_text, return_tensors="pt").to(device)
     outputs = model.generate(**inputs)
     result = processor.decode(outputs[0], skip_special_tokens=True)
     return result
 
+
 # Gradio interface
 iface = gr.Interface(
     fn=extract_chart_data,