update test for updated pretraining multipack code

tests/test_packed_pretraining.py

@@ -1,13 +1,13 @@
 """Module for testing streaming dataset sequence packing"""
-import math
 import unittest
 from functools import partial
 
+import torch
 from datasets import load_dataset
 from torch.utils.data import DataLoader
 from transformers import AutoTokenizer
 
-from axolotl.utils.collators import DataCollatorForSeq2Seq
+from axolotl.utils.collators import PretrainingBatchSamplerDataCollatorForSeq2Seq
 from axolotl.utils.data import encode_packed_pretraining
 
 
@@ -19,21 +19,9 @@ class TestPacking(unittest.TestCase):
     def setUp(self) -> None:
         # pylint: disable=duplicate-code
         self.tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-        self.tokenizer.add_special_tokens(
-            {
-                "bos_token": "<s>",
-                "eos_token": "</s>",
-                "unk_token": "<unk>",
-                "pad_token": "[PAD]",
-            }
-        )
-        self.max_seq_length = 8192
-        self.batch_size = 6
-        self.sample_packing_efficiency = 1
-        self.data_collator_kwargs = {
-            "padding": True,
-            "pad_to_multiple_of": 64 * math.ceil(self.max_seq_length / 64),
-        }
+        self.tokenizer.pad_token = "</s>"
+        self.max_seq_length = 2048
+        self.batch_size = 2
 
     def test_packing_stream_dataset(self):
         # pylint: disable=duplicate-code
@@ -43,11 +31,19 @@ def test_packing_stream_dataset(self):
             streaming=True,
         )["train"]
 
+        collate_fn = PretrainingBatchSamplerDataCollatorForSeq2Seq(
+            self.tokenizer,
+            return_tensors="pt",
+            padding=True,
+            pad_to_multiple_of=self.max_seq_length,
+        )
+
         encode = partial(
             encode_packed_pretraining,
             self.tokenizer,
+            collate_fn,
             max_seq_length=self.max_seq_length,
-            sample_packing_efficiency=self.sample_packing_efficiency,
+            batch_size=self.batch_size,
         )
 
         dataset = dataset.map(
@@ -57,28 +53,27 @@ def test_packing_stream_dataset(self):
             remove_columns=dataset.features.keys(),
         )
 
-        data_collator_fn = DataCollatorForSeq2Seq(
-            self.tokenizer,
-            return_tensors="pt",
-            **self.data_collator_kwargs,
-        )
-
         trainer_loader = DataLoader(
             dataset,
-            batch_size=self.batch_size,
-            collate_fn=data_collator_fn,
+            batch_size=1,
+            collate_fn=None,
             drop_last=True,
         )
         idx = 0
         for data in trainer_loader:
             if idx > 10:
                 break
-            assert data["input_ids"].shape == (self.batch_size, self.max_seq_length)
-            assert data["position_ids"].shape == (self.batch_size, self.max_seq_length)
-            assert data["labels"].shape == (self.batch_size, self.max_seq_length)
-            assert data["attention_mask"].shape == (
-                self.batch_size,
-                self.max_seq_length,
+            assert data["input_ids"].shape == torch.Size(
+                [1, self.batch_size * self.max_seq_length]
+            )
+            assert data["position_ids"].shape == torch.Size(
+                [1, self.batch_size * self.max_seq_length]
+            )
+            assert data["labels"].shape == torch.Size(
+                [1, self.batch_size * self.max_seq_length]
+            )
+            assert data["attention_mask"].shape == torch.Size(
+                [1, self.batch_size * self.max_seq_length]
             )
             idx += 1