ci calib dataset should not be fixed at 128 rows for all tests

ModelCloud · Jan 29, 2025 · c5c6287 · c5c6287
1 parent 4755ada
commit c5c6287
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/tests/models/model_test.py b/tests/models/model_test.py
@@ -126,15 +126,15 @@ def load_tokenizer(self, model_id_or_path, trust_remote_code=False):
         return tokenizer
 
     @classmethod
-    def load_dataset(self, tokenizer):
+    def load_dataset(self, tokenizer, rows: int = 128):
         traindata = load_dataset("json", data_files="/monster/data/model/dataset/c4-train.00000-of-01024.json.gz", split="train")
 
         datas = []
         for index, sample in enumerate(traindata):
             tokenized = tokenizer(sample['text'])
             if len(tokenized.data['input_ids']) < self.INPUTS_MAX_LENGTH:
                 datas.append(tokenized)
-                if len(datas) >= 128:
+                if len(datas) >= rows:
                     break
 
         return datas

diff --git a/tests/test_dynamic.py b/tests/test_dynamic.py
@@ -36,7 +36,7 @@
 
 class TestDynamic(ModelTest):
     NATIVE_MODEL_ID = "/monster/data/model/TinyLlama-1.1B-Chat-v1.0"  # "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    tmp_dir = None
+    tmp_quant_path = None
 
     def calculate_avg_ppl(self, model, tokenizer):
         ppl = Perplexity(
@@ -57,13 +57,13 @@ def calculate_avg_ppl(self, model, tokenizer):
 
     @classmethod
     def setUpClass(cls):
-        cls.tmp_dir = tempfile.TemporaryDirectory()
+        cls.tmp_quant_path = tempfile.TemporaryDirectory()
         cls.tokenizer = AutoTokenizer.from_pretrained(cls.NATIVE_MODEL_ID, use_fast=True)
 
         if not cls.tokenizer.pad_token_id:
             cls.tokenizer.pad_token_id = cls.tokenizer.eos_token_id
 
-        cls.calibration_dataset = cls.load_dataset(cls.tokenizer)
+        cls.calibration = cls.load_dataset(cls.tokenizer, rows=32)
 
         # support dynamic override of bits, group_size, desc_act, sym for each layer/module match
         dynamic = {
@@ -83,14 +83,14 @@ def setUpClass(cls):
             cls.NATIVE_MODEL_ID,
             quantize_config=quantize_config,
         )
-        model.quantize(cls.calibration_dataset, batch_size=4)
+        model.quantize(cls.calibration, batch_size=4)
 
-        model.save(cls.tmp_dir.name)
+        model.save(cls.tmp_quant_path.name)
 
     @classmethod
     def tearDownClass(cls):
-        cls.tmp_dir.cleanup()
-        assert not os.path.exists(cls.tmp_dir.name)
+        cls.tmp_quant_path.cleanup()
+        assert not os.path.exists(cls.tmp_quant_path.name)
 
     @parameterized.expand(
         [
@@ -100,7 +100,7 @@ def tearDownClass(cls):
     )
     def test_dynamic_bits(self, backend):
         model = GPTQModel.load(
-            self.tmp_dir.name,
+            self.tmp_quant_path.name,
             backend=backend,
         )
 
@@ -123,7 +123,7 @@ def test_skip_module(self):
             self.NATIVE_MODEL_ID,
             quantize_config=QuantizeConfig(dynamic=dynamic),
         )
-        model.quantize(self.calibration_dataset, batch_size=4)
+        model.quantize(self.calibration, batch_size=4)
 
         for name, submodule in model.named_modules():
             if name == 'model.model.layers.0.self_attn.q_proj' and isinstance(submodule, BaseQuantLinear):  # module 0 was skipped