Update main.py

main.py

@@ -7,16 +7,18 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 
 token = os.getenv("HUGGINGFACE_TOKEN")
+assert token is not None, "Hugging Face token is missing. Please set the 'HUGGINGFACE_TOKEN' environment variable."
+
 
 #Load pre-trained tokenizer and model (Works)
 model_name = "mistralai/Mistral-7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
 tokenizer.pad_token = tokenizer.eos_token
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="auto",
     torch_dtype="auto",
-    use_auth_token=token
+    token=token
 )