Fixes to mc4 fork

mc4/mc4.py

@@ -283,20 +283,20 @@ class Mc4(datasets.GeneratorBasedBuilder):
     BUILDER_CONFIG_CLASS = Mc4Config
 
     def __init__(self, *args, writer_batch_size=None, **kwargs):
-        self.filepaths = kwargs.pop(filepaths, {})
+        self.data_files = kwargs.pop("data_files", {})
         self.sampling_method = kwargs.pop("sampling_method", None)
+        self.perplexity_model = kwargs.pop("perplexity_model", None)
+        self.sampling_factor = kwargs.pop("sampling_factor", None)
+        self.boundaries = kwargs.pop("boundaries", None)
+        self.seed = kwargs.pop("seed", None)
         if self.sampling_method:
-            seed = kwargs.pop("seed", None)
-            if seed is not None:
-                self.rng = default_rng(seed)
+            if self.seed is not None:
+                self.rng = default_rng(self.seed)
             else:
                 self.rng = default_rng()
             if self.sampling_method == "random":
                 self.should_keep_doc = self._should_keep_doc_random
             else:
-                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)
@@ -305,7 +305,6 @@ class Mc4(datasets.GeneratorBasedBuilder):
                     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)
 
     def get_perplexity(self, doc):
@@ -375,14 +374,14 @@ class Mc4(datasets.GeneratorBasedBuilder):
                 for lang in self.config.languages
                 for index in range(_N_SHARDS_PER_SPLIT[lang][split])
             ]
-        if "train" in self.filepaths:
-            train_downloaded_files = self.filepaths["train"]
+        if "train" in self.data_files:
+            train_downloaded_files = self.data_files["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 "validation" in self.data_files:
+            validation_downloaded_files = self.data_files["validation"]
             if not isinstance(validation_downloaded_files, (tuple, list)):
                 validation_downloaded_files = [validation_downloaded_files]
         else:

run_mlm_flax_stream.py

@@ -402,7 +402,7 @@ if __name__ == "__main__":
             boundaries=sampling_args.boundaries,
             perplexity_model=sampling_args.perplexity_model,
             seed=training_args.seed,
-            filepaths={"train": filepaths},
+            data_files=filepaths,
         )
 
     if model_args.config_name: