Add token

main.py

@@ -4,23 +4,26 @@ from huggingface_hub import InferenceClient
 import uvicorn
 from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
 
+token = os.getenv("HUGGINGFACE_TOKEN")
 
 #Load pre-trained tokenizer and model (Works)
 model_name = "mistralai/Mistral-7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
 tokenizer.pad_token = tokenizer.eos_token
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="auto",
-    torch_dtype="auto"
+    torch_dtype="auto",
+    use_auth_token=token
 )
 
 
 # Example usage: Generate text
 prompt = "<s>[INST] What's the capital of France? [/INST]"
 
-inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
+inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True ).to(model.device)
 outputs = model.generate(
     **inputs,
     max_new_tokens=100,
@@ -60,7 +63,7 @@ async def generate_text(item: Item):
     # logging.info("Response generated")
 
     inp =f"<s>[INST] {item.prompt} [/INST]"
-    inputs = tokenizer(inp, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
+    inputs = tokenizer(inp, return_tensors="pt", padding=True, return_attention_mask=True ).to(model.device)
 
 
     # input_ids = tokenizer.encode(item.prompt, return_tensors="pt")