Fix packing test (#2111)

* Fix pack test * same for eval
huggingface · Sep 24, 2024 · a84fc5d · a84fc5d
1 parent 80038a5
commit a84fc5d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -1106,7 +1106,7 @@ def test_sft_trainer_only_train_packing(self):
                 eval_dataset=self.conversational_lm_dataset["test"],
             )
 
-            assert len(trainer.train_dataset["input_ids"]) == 16  # with the used dataset, we end up with 16 sequences
+            assert len(trainer.train_dataset["input_ids"]) == 21  # with the used dataset, we end up with 21 sequences
             assert len(trainer.eval_dataset["input_ids"]) == len(self.conversational_lm_dataset["test"])
 
     def test_sft_trainer_eval_packing(self):
@@ -1131,8 +1131,8 @@ def test_sft_trainer_eval_packing(self):
                 eval_dataset=self.conversational_lm_dataset["test"],
             )
 
-            assert len(trainer.train_dataset["input_ids"]) == 16  # with the used dataset, we end up with 16 sequences
-            assert len(trainer.eval_dataset["input_ids"]) == 1  # with the used dataset, we end up with 1 sequence
+            assert len(trainer.train_dataset["input_ids"]) == 21  # with the used dataset, we end up with 21 sequences
+            assert len(trainer.eval_dataset["input_ids"]) == 2  # with the used dataset, we end up with 2 sequence
 
     def test_sft_trainer_no_packing(self):
         with tempfile.TemporaryDirectory() as tmp_dir: