Update src/app/model.py

src/app/model.py

@@ -1,7 +1,7 @@
 # Necessary imports
-# import os
+import os
 import sys
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
 from typing import Any
 import torch
 from transformers import AutoModel, AutoTokenizer
@@ -11,11 +11,11 @@ from src.logger import logging
 from src.exception import CustomExceptionHandling
 
 
-# # Load the Environment Variables from .env file
-# load_dotenv()
+# Load the Environment Variables from .env file
+load_dotenv()
 
-# # Access token for using the model
-# access_token = os.environ.get("ACCESS_TOKEN")
+# Access token for using the model
+access_token = os.environ.get("ACCESS_TOKEN")
 
 
 def load_model_and_tokenizer(model_name: str, device: str) -> Any:
@@ -37,10 +37,12 @@ def load_model_and_tokenizer(model_name: str, device: str) -> Any:
             trust_remote_code=True,
             attn_implementation="sdpa",
             torch_dtype=torch.bfloat16,
-            # token=access_token
+            token=access_token
         )
         model = model.to(device=device)
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=True, token=access_token
+        )
         model.eval()
 
         # Log the successful loading of the model and tokenizer