Update app.py

app.py

@@ -29,11 +29,23 @@ model = model.to(device)  # Move model to appropriate device
 model = model.eval()
 
 # Override the chat function to remove hardcoded .cuda()
-def modified_chat(inputs, *args, ocr_type='ocr', **kwargs):
-    input_ids = torch.as_tensor(inputs.input_ids).to(device)  # Use .to(device)
-    # Additional processing logic here
-    # Example: replace with actual model inference code if necessary
-    # res = model(input_ids)
+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
+    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)  # Use .to(device)
+    
+    # Perform any necessary processing using the model
+    # Example: res = model(input_ids)  # Uncomment and implement model processing
+    
     return f"Processed input: {input_ids}, OCR Type: {ocr_type}"
 
 # Replace the model's chat method with the modified version
@@ -123,28 +135,4 @@ if image_file is not None:
             st.markdown(highlighted_text)
 
         # 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
+        os.remove(temp_file_path)