Check for EOS

ratschlab · Jun 30, 2024 · dcaf79b · dcaf79b
1 parent ac4e377
commit dcaf79b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/nanotron/data/petagraph_dataset.py b/src/nanotron/data/petagraph_dataset.py
@@ -266,6 +266,9 @@ def generate(self):
                 if current_tokens is None:
                     current_tokens = new_tokens
                 else:
+                    # Check the last token of the current sequence
+                    # is an EOS token
+                    assert current_tokens[-1] == self._eos_token_id
                     current_tokens = np.concatenate([current_tokens, new_tokens])
 
                 if len(current_tokens) >= self.maxlen: