Update app.py

app.py

@@ -124,3 +124,27 @@ if image_file is not None:
 
         # Clean up the temporary file after use
         os.remove(temp_file_path)
+
+# Override the chat function to remove hardcoded .cuda()
+def modified_chat(tokenizer, temp_file_path, ocr_type='ocr', *args, **kwargs):
+    # Load the image data, perform OCR and get text
+    with open(temp_file_path, 'rb') as f:
+        image_data = f.read()
+    
+    # Assuming OCR process to extract text from image
+    # Use the OCR process here to obtain text. Replace 'extracted_text' with actual extracted text.
+    extracted_text = "some OCR processed text"  # Placeholder, replace with actual OCR result
+
+    # Tokenize the extracted text
+    inputs = tokenizer(extracted_text, return_tensors="pt", truncation=True, padding=True)
+    
+    # Move input_ids to the appropriate device
+    input_ids = inputs['input_ids'].to(device)  # Move to the appropriate device
+    # Perform any necessary processing using the model
+    # Example: res = model(input_ids)
+    
+    # Return the processed result
+    return f"Processed input: {input_ids}, OCR Type: {ocr_type}"
+
+# Replace the model's chat method with the modified version
+model.chat = modified_chat