Adding reading streaming files from local disk

mc4/mc4.py

@@ -283,6 +283,7 @@ class Mc4(datasets.GeneratorBasedBuilder):
     BUILDER_CONFIG_CLASS = Mc4Config
 
     def __init__(self, *args, writer_batch_size=None, **kwargs):
+        self.filepaths = kwargs.pop(filepaths, {})
         self.sampling_method = kwargs.pop("sampling_method", None)
         if self.sampling_method:
             seed = kwargs.pop("seed", None)
@@ -290,19 +291,20 @@ class Mc4(datasets.GeneratorBasedBuilder):
                 self.rng = default_rng(seed)
             else:
                 self.rng = default_rng()
-            self.perplexity_model = kwargs.pop("perplexity_model", None)
-            self.sampling_factor = kwargs.pop("sampling_factor", None)
-            self.boundaries = kwargs.pop("boundaries", None)
-            # Loading 5-gram model
-            # http://dl.fbaipublicfiles.com/cc_net/lm/es.arpa.bin
-            logger.info("loading model = %s", self.perplexity_model)
-            self.pp_model = kenlm.Model(self.perplexity_model)
-            if self.sampling_method == "gaussian":
-                self.should_keep_doc = self._should_keep_doc_gaussian
-            elif self.sampling_method == "random":
+            if self.sampling_method == "random":
                 self.should_keep_doc = self._should_keep_doc_random
             else:
-                self.should_keep_doc = self._should_keep_doc_step
+                self.perplexity_model = kwargs.pop("perplexity_model", None)
+                self.sampling_factor = kwargs.pop("sampling_factor", None)
+                self.boundaries = kwargs.pop("boundaries", None)
+                # Loading 5-gram model
+                # http://dl.fbaipublicfiles.com/cc_net/lm/es.arpa.bin
+                logger.info("loading model = %s", self.perplexity_model)
+                self.pp_model = kenlm.Model(self.perplexity_model)
+                if self.sampling_method == "gaussian":
+                    self.should_keep_doc = self._should_keep_doc_gaussian
+                else:
+                    self.should_keep_doc = self._should_keep_doc_step
 
         super().__init__(*args, writer_batch_size=writer_batch_size, **kwargs)
 
@@ -341,7 +343,9 @@ class Mc4(datasets.GeneratorBasedBuilder):
         return self.rng.uniform() < weighted_perplexity
 
     def _should_keep_doc_random(self, doc, factor=None, boundaries=None):
-        return self.rng.uniform() <= 0.5
+        if factor is None:
+            factor = 0.5
+        return self.rng.uniform() <= factor
 
     def _info(self):
         return datasets.DatasetInfo(
@@ -371,8 +375,18 @@ class Mc4(datasets.GeneratorBasedBuilder):
                 for lang in self.config.languages
                 for index in range(_N_SHARDS_PER_SPLIT[lang][split])
             ]
-        train_downloaded_files = dl_manager.download(data_urls["train"])
-        validation_downloaded_files = dl_manager.download(data_urls["validation"])
+        if "train" in self.filepaths:
+            train_downloaded_files = self.filepaths["train"]
+            if not isinstance(train_downloaded_files, (tuple, list)):
+                train_downloaded_files = [train_downloaded_files]
+        else:
+            train_downloaded_files = dl_manager.download(data_urls["train"])
+        if "validation" in self.filepaths:
+            validation_downloaded_files = self.filepaths["validation"]
+            if not isinstance(validation_downloaded_files, (tuple, list)):
+                validation_downloaded_files = [validation_downloaded_files]
+        else:
+            validation_downloaded_files = dl_manager.download(data_urls["validation"])
         return [
             datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepaths": train_downloaded_files}),
             datasets.SplitGenerator(
@@ -385,6 +399,14 @@ class Mc4(datasets.GeneratorBasedBuilder):
         id_ = 0
         for filepath in filepaths:
             logger.info("generating examples from = %s", filepath)
+            if filepath.endswith("json") or filepath.endswith("jsonl"):
+                with open(filepath, "r", encoding="utf-8") as f:
+                    for line in f:
+                        if line:
+                            example = json.loads(line)
+                            yield id_, example
+                            id_ += 1
+        else:
             with gzip.open(open(filepath, "rb"), "rt", encoding="utf-8") as f:
                 if self.sampling_method:
                     logger.info("sampling method = %s", self.sampling_method)

run_mlm_flax_stream.py

@@ -178,10 +178,10 @@ class DataTrainingArguments:
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]
-                assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
+                assert extension in ["csv", "json", "txt", "gz"], "`train_file` should be a csv, a json or a txt file."
             if self.validation_file is not None:
                 extension = self.validation_file.split(".")[-1]
-                assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
+                assert extension in ["csv", "json", "txt", "gz"], "`validation_file` should be a csv, a json or a txt file."
 
 
 @flax.struct.dataclass
@@ -386,6 +386,11 @@ if __name__ == "__main__":
     # 'text' is found. You can easily tweak this behavior (see below).
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
+        filepaths = {}
+        if data_args.train_file:
+            filepaths["train"] = data_args.train_file
+        if data_args.validation_file:
+            filepaths["validation"] = data_args.validation_file
         dataset = load_dataset(
             data_args.dataset_name,
             data_args.dataset_config_name,
@@ -397,6 +402,7 @@ if __name__ == "__main__":
             boundaries=sampling_args.boundaries,
             perplexity_model=sampling_args.perplexity_model,
             seed=training_args.seed,
+            filepaths=filepaths,
         )
 
     if model_args.config_name: