docker

huggingface · Feb 2, 2024 · ae5c3bd · ae5c3bd
1 parent 9891f5d
commit ae5c3bd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -5,12 +5,15 @@
 
 device = "cuda"
 
-assert os.getenv("HF_HUB_READ_TOKEN", None) != None
+
 
 
 class LLaMaInferenceTest(unittest.TestCase):
 
     def test_foo(self):
         ckpt = "meta-llama/Llama-2-7b"
+        assert os.getenv("HF_HUB_READ_TOKEN", None) != None
+        assert type(os.getenv("HF_HUB_READ_TOKEN", None)) == str
+        assert len(os.getenv("HF_HUB_READ_TOKEN", None)) > 3
         model = AutoModelForCausalLM.from_pretrained(ckpt, token=os.getenv("HF_HUB_READ_TOKEN", None))
         model.to(device)
diff --git a/tests/models/llama/test_train_llama.py b/tests/models/llama/test_train_llama.py
@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from transformers import AutoModelForCausalLM
@@ -9,6 +10,9 @@ class LLaMaTrainingTest(unittest.TestCase):
 
     def test_foo(self):
         ckpt = "meta-llama/Llama-2-7b"
+        assert os.getenv("HF_HUB_READ_TOKEN", None) != None
+        assert type(os.getenv("HF_HUB_READ_TOKEN", None)) == str
+        assert len(os.getenv("HF_HUB_READ_TOKEN", None)) > 3
         model = AutoModelForCausalLM.from_pretrained(ckpt, token=os.getenv("HF_HUB_READ_TOKEN", None))
         model.train()
         model.to(device)