Update src/app/model.py

src/app/model.py

@@ -4,7 +4,7 @@ import sys
 from dotenv import load_dotenv
 from typing import Any
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, AutoProcessor
 
 # Local imports
 from src.logger import logging
@@ -43,13 +43,16 @@ def load_model_and_tokenizer(model_name: str, device: str) -> Any:
         tokenizer = AutoTokenizer.from_pretrained(
             model_name, trust_remote_code=True, token=access_token
         )
+        processor = AutoProcessor.from_pretrained(
+            model_name, trust_remote_code=True, token=access_token
+        )
         model.eval()
 
         # Log the successful loading of the model and tokenizer
         logging.info("Model and tokenizer loaded successfully.")
 
         # Return the model and tokenizer
-        return model, tokenizer
+        return model, tokenizer, processor
 
     # Handle exceptions that may occur during model and tokenizer loading
     except Exception as e: