Use test data zipfile in CI (#662)

* make fewer requests in tests

* token

Author: awni

.github/workflows/pull_request.yml

@@ -38,4 +38,6 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: |
-          python -m xmlrunner discover -v tests -o test-results/
+          curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
+          unzip test_data.zip
+          HF_HOME="." python -m xmlrunner discover -v tests -o test-results/

tests/test_datsets.py

@@ -21,6 +21,8 @@ def setUpClass(cls):
         cls.test_dir = cls.test_dir_fid.name
         if not os.path.isdir(cls.test_dir):
             os.mkdir(cls.test_dir_fid.name)
+        # Only one HF request
+        AutoTokenizer.from_pretrained(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -37,7 +39,7 @@ def test_text(self):
         data = {"text": "This is an example for the model."}
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -50,7 +52,7 @@ def test_completions(self):
         data = {"prompt": "What is the capital of France?", "completion": "Paris."}
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -69,7 +71,7 @@ def test_chat(self):
         }
         self.save_data(4 * [data])
         args = types.SimpleNamespace(train=True, test=False, data=self.test_dir)
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertEqual(len(train), 4)
         self.assertEqual(len(valid), 4)
@@ -91,7 +93,7 @@ def test_hf(self):
             test=False,
             train=True,
         )
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH)
+        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_PATH, local_files_only=True)
         train, valid, test = datasets.load_dataset(args, tokenizer)
         self.assertTrue(len(train) > 0)
         self.assertTrue(len(train[0]) > 0)

tests/test_generate.py

@@ -81,7 +81,7 @@ def logits_processor(toks, logits):
 
     def test_stream_generate_speculative(self):
         # Use same model as draft model, this is not a speed test
-        draft_model, _ = load(self.HF_MODEL_PATH)
+        draft_model = self.model
 
         results: List[GenerationResponse] = []
         drafted: List[bool] = []

tests/test_prompt_cache.py

@@ -34,6 +34,7 @@ class TestPromptCache(unittest.TestCase):
     def setUpClass(cls):
         cls.test_dir_fid = tempfile.TemporaryDirectory()
         cls.test_dir = cls.test_dir_fid.name
+        cls.model, cls.tokenizer = load(HF_MODEL_PATH)
 
     @classmethod
     def tearDownClass(cls):
@@ -132,7 +133,7 @@ def test_save_load_mixed_cache(self):
                 self.assertTrue(mx.array_equal(v, lv))
 
     def test_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = list(generate_step(prompt, model, max_tokens=4))
         toks, all_logits = zip(*results)
@@ -212,7 +213,7 @@ def test_trim_cache(self):
         self.assertEqual(num_trimmed, 3)
 
     def test_trim_cache_with_generate(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
 
         prompt_cache = make_prompt_cache(model)
@@ -289,7 +290,7 @@ def test_save_load_quantized_cache(self):
         self.assertEqual(metadata, loaded_metadata)
 
     def test_cache_to_quantized(self):
-        model, tokenizer = load(HF_MODEL_PATH)
+        model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]
         results = zip(range(4), generate_step(prompt, model))
         toks, all_logits = zip(*(r[1] for r in results))