Update app.py

app.py

@@ -130,7 +130,7 @@ def load_image(image_file):
 def vision_ai_api(image, doc_type):
     """Run the model using a dynamic prompt based on detected doc type."""
     pixel_values = load_image(image).to(torch.float32).to("cpu")
-    generation_config = dict(max_new_tokens=1024, do_sample=True)
+    generation_config = dict(max_new_tokens=512, do_sample=True)
     
     question = front_prompt if doc_type == "front" else back_prompt if doc_type == "back" else "Please provide document details."
     
@@ -143,7 +143,7 @@ def vision_ai_api(image, doc_type):
     gc.collect()  # Force garbage collection
     torch.cuda.empty_cache()  
     
-    return f'User: {question}\nAssistant: {response}'
+    return f'Assistant: {response}'
 
 # ---------------------------
 # PREDICTION PIPELINE