Update app.py

app.py

@@ -29,11 +29,12 @@ model = model.to(device)  # Move model to appropriate device
 model = model.eval()
 
 # Override the chat function to remove hardcoded .cuda()
-def modified_chat(inputs, ocr_type='ocr', ocr_box='', render=False, save_render_file=''):
+def modified_chat(inputs, *args, ocr_type='ocr', **kwargs):
     input_ids = torch.as_tensor(inputs.input_ids).to(device)  # Use .to(device)
-    # Replace all .cuda() calls with .to(device)
-    # Handle the remaining logic as needed in the function
-    # ...
+    # Additional processing logic here
+    # Example: replace with actual model inference code if necessary
+    # res = model(input_ids)
+    return f"Processed input: {input_ids}, OCR Type: {ocr_type}"
 
 # Replace the model's chat method with the modified version
 model.chat = modified_chat
@@ -68,7 +69,7 @@ if image_file is not None:
     if st.button("Run OCR"):
         # Use GOT-OCR2 model for plain text OCR (structured documents)
         with torch.no_grad():
-            res_plain = model.chat(tokenizer, temp_file_path, ocr_type='ocr')
+            res_plain = model.chat(tokenizer, temp_file_path, ocr_type='ocr')  # Ensure the correct parameters are passed
 
         # Perform formatted text OCR
         with torch.no_grad():