Fix dataset preprocessor for jsonl format

Author: hyunwoongko

USAGE.md

@@ -488,6 +488,10 @@ import os
 from transformers import AutoTokenizer
 from oslo import DatasetPreprocessor
 
+data_names=[
+  "/path/to/wikitext103", "/path/to/lambada", ...
+]
+
 # 1. Create tokenizer
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 tokenizer = AutoTokenizer.from_pretrained(...)
@@ -500,16 +504,36 @@ preprocessor = DatasetPreprocessor(
     append_eod=True,
 )
 
-# 3. Preform preprocessing
-preprocessor.preprocess(
-    data_names=[
-        "/path/to/wikitext103", 
-        "/path/to/lambada", 
-        ...
-    ],
-    extension=".txt",
-    log_interval=100,
-)
+# 3-1. Preform preprocessing (.txt)
+# save_file_name + '.idx' and '.bin' will be created.
+for data_name in data_names:
+    preprocessor.preprocess(
+        open(data_name + ".txt"),
+        save_file_name=data_name,
+        log_interval=100,
+    )
+
+# 3-2. Perform preprocessing (.jsonl, Megatron-LM format)
+# 1 {"text": "blah blah"}
+# 2 {"text": "blah blah"}
+# 3 ...
+for data_name in data_names:
+  preprocessor.preprocess(
+        preprocessor.open_jsonl(
+            data_name, 
+            json_key="text",
+        ),
+        save_file_name=data_name,
+        log_interval=100,
+  )
+
+# 3-3 Perform preprocessing (any other format)
+for data_name in data_names:
+  preprocessor.preprocess(
+        YOUR_OWN_LIST_OF_STRING,
+        save_file_name=data_name,
+        log_interval=100,
+  )
 ```
 
 ### DatasetForCausalLM

oslo/data/preprocess/preprocessor.py

@@ -92,19 +92,18 @@ def __init__(
         else:
             self.eod_token_id = eod_token_id
 
-    def _preprocess(
+    def preprocess(
         self,
-        data_path: str,
-        index_path: str,
-        log_interval: int,
-        kwargs,
+        iterable,
+        save_file_name: str,
+        log_interval: int = 1000,
     ) -> None:
         """
         Preprocess a dataset
 
         Args:
-            data_path (str): dataset bin path
-            index_path (str): dataset index path
+            iterable: iterable of string
+            save_file_name (str): save file name
             log_interval (int) logging interval
         """
 
@@ -114,12 +113,12 @@ def _preprocess(
             eod_token_id=self.eod_token_id,
         )
         binarizer = DatasetBinarizer(self.binarization_impl)
-        index_path, builder = binarizer.create_builder(index_path)
+        index_path, builder = binarizer.create_builder(save_file_name)
 
         with ProcessPoolExecutor() as pool:
             iterator = pool.map(
                 encoder.encode,
-                open(data_path, **kwargs),
+                iterable,
                 chunksize=self.chunksize,
             )
 
@@ -130,30 +129,33 @@ def _preprocess(
                 log_interval=log_interval,
             )
 
-    def preprocess(
-        self,
-        data_names: List[str],
-        extension: str = ".txt",
-        log_interval: int = 1000,
-        **kwargs,
-    ):
+    @staticmethod
+    def open_jsonl(file, json_key):
         """
-        Preprocess datasets
+        Open jsonl file similar with Megatron-LM data format
 
-        Args:
-            data_names (List[str]): dataset names
-            extension (str): data file extension
-            log_interval (int): logging interval
+        Examples:
+            1 {'text': 'blah blah blah ...'}
+            2 {'text': 'blah blah blah ...'}
+            3 {'text': 'blah blah blah ...'}
+            4 ...
+
+            >>> DatasetPreprocessor.open_jsonl(
+            ...     file=FILE_NAME, json_key='text'
+            ... )
         """
 
-        if "." not in extension:
-            extension = "." + extension
+        import json
 
-        for name in data_names:
-            logger.info(f"Start to preprocess {name}.")
-            self._preprocess(
-                name + extension,
-                name,
-                log_interval,
-                kwargs,
-            )
+        if file[-6:].lower() != ".jsonl":
+            file = file + ".jsonl"
+
+        source = open(file)
+
+        while True:
+            line = source.readline()
+
+            if not line:
+                break
+            else:
+                yield json.loads(line)[json_key]