Adding checkpointing, wandb, and new mlm script

README.md

@@ -1,19 +1,47 @@
+---
+language: no
+license: CC-BY 4.0
+tags:
+- spanish
+- roberta
+pipeline_tag: fill-mask
+widget:
+- text: "Lo hizo en un abrir y cerar de <mask>."
+---
+
 # BERTIN
+
 BERTIN is a series of BERT-based models for Spanish. This one is a RoBERTa-large model trained from scratch on the Spanish portion of mC4 using [Flax](https://github.com/google/flax), including training scripts.
 
 This is part of the
 [Flax/Jax Community Week](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104), organised by [HuggingFace](https://huggingface.co/) and TPU usage sponsored by Google.
 
+## Spanish mC4
+
+The Spanish portion of mC4 containes about 416 million records and 235 billion words.
+
+```bash
+$ zcat c4/multilingual/c4-es*.tfrecord*.json.gz | wc -l
+416057992
+```
+
+```bash
+$ zcat c4/multilingual/c4-es*.tfrecord-*.json.gz | jq -r '.text | split(" ") | length' | paste -s -d+ - | bc
+235303687795
+```
+
 ## Team members
+
 - Javier de la Rosa (versae)
 - Manu Romero (mrm8488)
 - María Grandury (mariagrandury)
 - Ari Polakov (aripo99)
 - Pablogps
 - daveni
-- Sri Lakshmi 
+- Sri Lakshmi
 
 ## Useful links
+
 - [Community Week timeline](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104#summary-timeline-calendar-6)
 - [Community Week README](https://github.com/huggingface/transformers/blob/master/examples/research_projects/jax-projects/README.md)
 - [Community Week thread](https://discuss.huggingface.co/t/bertin-pretrain-roberta-large-from-scratch-in-spanish/7125)

perplexity.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+import kenlm
+from datasets import load_dataset
+from tqdm import tqdm
+
+
+def pp(log_score, length):
+    return 10.0 ** (-log_score / length)
+
+# http://dl.fbaipublicfiles.com/cc_net/lm/es.arpa.bin
+model = kenlm.Model("es.arpa.bin")
+mc4 = load_dataset("mc4", "es", streaming=True)
+with open("mc4-es-perplexity.txt", "w") as f:
+    for sample in tqdm(mc4["train"].shuffle(buffer_size=100_000), total=416057992):
+        lines = sample["text"].split("\n")
+        doc_log_score, doc_length = 0, 0
+        for line in lines:
+            log_score = model.score(line)
+            length = len(line.split()) + 1
+            doc_log_score += log_score
+            doc_length += length
+        f.write(f"{pp(doc_log_score, doc_length)}\n")

run_mlm_flax.py

@@ -56,22 +56,6 @@ from transformers import (
 )
 
 
-# Cache the result
-has_tensorboard = is_tensorboard_available()
-if has_tensorboard:
-    try:
-        from flax.metrics.tensorboard import SummaryWriter
-    except ImportError as ie:
-        has_tensorboard = False
-        print(f"Unable to display metrics through TensorBoard because some package are not installed: {ie}")
-
-else:
-    print(
-        "Unable to display metrics through TensorBoard because the package is not installed: "
-        "Please run pip install tensorboard to enable."
-    )
-
-
 MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_MASKED_LM_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
@@ -126,6 +110,9 @@ class DataTrainingArguments:
     dataset_config_name: Optional[str] = field(
         default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
     )
+    dataset_streaming: bool = field(
+        default=False, metadata={"help": "Whether dataset_name should be retrieved using streaming if available."}
+    )
     train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
     validation_file: Optional[str] = field(
         default=None,
@@ -269,7 +256,7 @@ def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndar
     return batch_idx
 
 
-def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
+def write_train_metric(summary_writer, train_metrics, train_time, step):
     summary_writer.scalar("train_time", train_time, step)
 
     train_metrics = get_metrics(train_metrics)
@@ -278,6 +265,8 @@ def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
         for i, val in enumerate(vals):
             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
 
+
+def write_eval_metric(summary_writer, eval_metrics, step):
     for metric_name, value in eval_metrics.items():
         summary_writer.scalar(f"eval_{metric_name}", value, step)
 
@@ -315,10 +304,6 @@ if __name__ == "__main__":
 
     # Log on each process the small summary:
     logger = logging.getLogger(__name__)
-    #logger.warning(
-    #    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
-    #    + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
-    #)
 
     # Set the verbosity to info of the Transformers logger (on main process only):
     logger.info(f"Training/evaluation parameters {training_args}")
@@ -337,7 +322,7 @@ if __name__ == "__main__":
     # download the dataset.
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
-        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)
+        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, streaming=data_args.dataset_streaming)
 
         if "validation" not in datasets.keys():
             datasets["validation"] = load_dataset(
@@ -345,12 +330,14 @@ if __name__ == "__main__":
                 data_args.dataset_config_name,
                 split=f"train[:{data_args.validation_split_percentage}%]",
                 cache_dir=model_args.cache_dir,
+                streaming=data_args.dataset_streaming,
             )
             datasets["train"] = load_dataset(
                 data_args.dataset_name,
                 data_args.dataset_config_name,
                 split=f"train[{data_args.validation_split_percentage}%:]",
                 cache_dir=model_args.cache_dir,
+                streaming=data_args.dataset_streaming,
             )
     else:
         data_files = {}
@@ -469,10 +456,32 @@ if __name__ == "__main__":
             num_proc=data_args.preprocessing_num_workers,
             load_from_cache_file=not data_args.overwrite_cache,
         )
-
     # Enable tensorboard only on the master node
+    has_tensorboard = is_tensorboard_available()
     if has_tensorboard and jax.process_index() == 0:
-        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
+        try:
+            from flax.metrics.tensorboard import SummaryWriter
+            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
+            # Enable Weight&Biases
+            import wandb
+            wandb.init(
+                entity='wandb',
+                project='hf-flax-bertin-roberta-es',
+                sync_tensorboard=True,
+            )
+            wandb.config.update(training_args)
+            wandb.config.update(model_args)
+            wandb.config.update(data_args)
+        except ImportError as ie:
+            has_tensorboard = False
+            logger.warning(
+                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
+            )
+    else:
+        logger.warning(
+            "Unable to display metrics through TensorBoard because the package is not installed: "
+            "Please run pip install tensorboard to enable."
+        )
 
     # Data collator
     # This one will take care of randomly masking the tokens.
@@ -521,7 +530,7 @@ if __name__ == "__main__":
         learning_rate=linear_decay_lr_schedule_fn,
         b1=training_args.adam_beta1,
         b2=training_args.adam_beta2,
-        eps=1e-8,
+        eps=training_args.adam_epsilon,
         weight_decay=training_args.weight_decay,
         mask=decay_mask_fn,
     )
@@ -601,7 +610,7 @@ if __name__ == "__main__":
         train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
 
         # Gather the indexes for creating the batch and do a training step
-        for i, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
+        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
             samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
             model_inputs = data_collator(samples, pad_to_multiple_of=16)
 
@@ -610,11 +619,31 @@ if __name__ == "__main__":
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
             train_metrics.append(train_metric)
 
-        train_time += time.time() - train_start
+            cur_step = epoch * (num_train_samples // train_batch_size) + step
 
-        epochs.write(
-            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
-        )
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
+                # Save metrics
+                train_metric = jax_utils.unreplicate(train_metric)
+                train_time += time.time() - train_start
+                if has_tensorboard and jax.process_index() == 0:
+                    write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                epochs.write(
+                    f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
+                )
+
+                train_metrics = []
+
+            if training_args.save_strategy == "steps" and cur_step and cur_step % training_args.save_steps == 0:
+                if jax.process_index() == 0:
+                    params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+                    model.save_pretrained(
+                        Path(str(training_args.output_dir)) / "checkpoints" / f"checkpoint-{cur_step}",
+                        params=params,
+                        push_to_hub=training_args.push_to_hub,
+                        temp_dir=True,
+                        commit_message=f"Saving weights and logs of step {cur_step}",
+                    )
 
         # ======================== Evaluating ==============================
         num_eval_samples = len(tokenized_datasets["validation"])
@@ -645,7 +674,7 @@ if __name__ == "__main__":
         # Save metrics
         if has_tensorboard and jax.process_index() == 0:
             cur_step = epoch * (len(tokenized_datasets["train"]) // train_batch_size)
-            write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
+            write_eval_metric(summary_writer, eval_metrics, cur_step)
 
         # save checkpoint after each epoch and push checkpoint to the hub
         if jax.process_index() == 0:

tokens.py

@@ -3,12 +3,14 @@ from datasets import load_dataset
 from tokenizers import ByteLevelBPETokenizer
 
 # Load dataset
-dataset = load_dataset("large_spanish_corpus", split="train")
+dataset = load_dataset("oscar", "unshuffled_deduplicated_es")
+
 # Instantiate tokenizer
 tokenizer = ByteLevelBPETokenizer()
 def batch_iterator(batch_size=100_000_000):
     for i in range(0, len(dataset), batch_size):
         yield dataset["text"][i: i + batch_size]
+
 # Customized training
 tokenizer.train_from_iterator(batch_iterator(), vocab_size=50265, min_frequency=2, special_tokens=[
     "<s>",