docker

huggingface · Feb 7, 2024 · 34fdf75 · 34fdf75
1 parent db21e83
commit 34fdf75
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -11,13 +11,14 @@ class LLaMaInferenceTest(unittest.TestCase):
 
     def test_foo(self):
         ckpt = "meta-llama/Llama-2-7b-hf"
+        hf_token = token=os.getenv("HF_HUB_READ_TOKEN", None)
 
-        tokenizer = AutoTokenizer.from_pretrained(ckpt)
+        tokenizer = AutoTokenizer.from_pretrained(ckpt, token=hf_token)
 
         prompt = "Hey, are you conscious? Can you talk to me?"
         inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-        model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, token=os.getenv("HF_HUB_READ_TOKEN", None))
+        model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, token=hf_token)
         model.to(device)
 
         # Generate