fix

EasyLM/data.py

@@ -175,13 +175,14 @@ class HuggingfaceDataset(object):
         self._index = 0
 
     def __iter__(self):
+        if not self._eval_dataset and self._train_epochs > 0:
+            self._dataset = self._dataset.shuffle(seed=42, buffer_size=10000)
         chunk_size = self.config.batch_size * self.config.seq_length
-        total_tokens = 0
         while True:
             token_buffer = []
             loss_mask_buffer = []
-            if not self._eval_dataset:
-                self._shuffle()
+            if not self._eval_dataset and self._train_epochs > 0:
+                self._dataset.set_epoch(self._train_epochs)
             for index, example in enumerate(self._dataset):
                 self._index = index
                 if not self._eval_dataset and self._dataset_loc > index:
@@ -217,12 +218,7 @@ class HuggingfaceDataset(object):
                 break
             else:
                 self._dataset_loc = 0
-                self._shuffle()
                 self._train_epochs += 1
-                print(f"TRAIN {self._train_epochs} EPOCH DONE")
-
-    def _shuffle(self):
-        self._dataset = self._dataset.shuffle(buffer_size=100)
 
     def get_state_dict(self):
         return dict(