tweak data loader to log metrics

codeparrot_training.py

@@ -22,6 +22,7 @@ class ConstantLengthDataset(IterableDataset):
         self.dataset = dataset
         self.seq_length = seq_length
         self.input_characters = seq_length * chars_per_token * num_of_sequences
+        self.epoch = 0
     
     def __iter__(self):
         iterator = iter(self.dataset)
@@ -36,6 +37,8 @@ class ConstantLengthDataset(IterableDataset):
                     buffer_len += len(buffer[-1])
                 except StopIteration:
                     iterator = iter(self.dataset)
+                    self.epoch += 1
+                    logger.info(f"Dataset epoch: {self.epoch}")
             tokenized_inputs = tokenizer(buffer, truncation=False)['input_ids']
             all_token_ids = []
             for tokenized_input in tokenized_inputs: