Update app.py

app.py

@@ -5,8 +5,8 @@ from retriever import retrieve_documents
 
 # Load Mistral 7B model
 MODEL_NAME = "mistralai/Mistral-7B-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"))
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN")) #, device_map="auto", torch_dtype=torch.float16)
 
 # Create inference pipeline
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)