fix tokenization error on empty data

huggingface · Dec 6, 2023 · 3e3f0c8 · 3e3f0c8
1 parent 7e007ff
commit 3e3f0c8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/datatrove/pipeline/tokens/tokenizer.py b/src/datatrove/pipeline/tokens/tokenizer.py
@@ -182,6 +182,10 @@ def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> Do
         unshuf_filename = self.get_output_filename(rank, "unshuffled")
         logger.info(f'Tokenizing in "{unshuf_filename}"...')
         outputfile: TokenizedFile = self.write_unshuffled(data, unshuf_filename)
+        if len(outputfile) == 0:
+            self.output_folder.close()
+            logger.warning("No data saved.")
+            return
         if self.shuffle:
             shuffled_filename = self.get_output_filename(rank, "shuffled")
             # get new TokenizedFile, shuffling docs from original one