Update weights to checkpoint 140000

README.md

@@ -25,15 +25,40 @@ For a demo of the model, head over to the Hugging Face Spaces for the **[Netherf
 
 ## Dataset
 
+
 `t5-base-dutch-demo` is fine-tuned on three mixed news sources:
 
  1. **CNN DailyMail** translated to Dutch with MarianMT.
  2. **XSUM** translated to Dutch with MarianMt.
  3. News article summaries distilled from the nu.nl website.
+
+The total number of training examples in this dataset is 1366592.
  
 ## Training
 
-The pre-trained model [t5-base-dutch](https://huggingface.co/flax-community/t5-base-dutch) was fine-tuned with a constant learning rate of 0.0005 and a batch size of 64 for 10.000 steps. 
-The performance of this model can be improved with longer training. Unfortunately due to a bug, an earlier training script would not save intermediate checkpoints and had been started for 6 epochs, which would have it finish past the TPU-VM availability schedule. Since there was limited time left, the fine-tuning was restarted without evaluation and for only half an epoch (10.000 steps).
+Training consisted of fine-tuning [t5-base-dutch](https://huggingface.co/flax-community/t5-base-dutch) with
+the following parameters:
+
+ * Constant learning rate 0.0005
+ * Batch size 8
+ * 1 epoch (170842 steps)
+
+## Evaluation
+
+The performance of the summarization model is measured with the Rouge metric from the
+Huggingface Datasets library.
+
+```
+    "rouge{n}" (e.g. `"rouge1"`, `"rouge2"`) where: {n} is the n-gram based scoring,
+    "rougeL": Longest common subsequence based scoring.
+    "rougeLSum": rougeLsum splits text using "
+"
+```
 
+ * Rouge1: 28.7066 
+ * Rouge2: 9.5498 
+ * RougeL: 22.8103
+ * rougeLsum: 24.2696
 
+These scores are expected to improve when the model is trained and evaluation configured
+for the CNN DM and XSUM datasets (translated to Dutch) individually.

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "flax-community/t5-base-dutch-demo",
+  "_name_or_path": "./",
   "architectures": [
     "T5ForConditionalGeneration"
   ],

output/events.out.tfevents.1626477704.t1v-n-0e7426e8-w-0.83817.3.v2 → events.out.tfevents.1626708806.yeb-z390-k80.10632.3.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:091b61fc500aae0368d977c5c0fd73632a32aabebdc0e7fba4129f26b6c8abdf
-size 6630102
+oid sha256:daeab64aaf6dd18fc097ee6bed7cd5e4e765e75716ca80c47777ad3b849b3679
+size 19440898

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8299c056e5ca07f93db2db052d61cb941710e0925c62486ee0c9775116e0a6bf
+oid sha256:8ba1de1818d02f938ad913208487e569e15df1ce66ea9a2fa9580bb9f2a32f19
 size 891548548

output/ckpt-9999/config.json

@@ -1,58 +0,0 @@
-{
-  "_name_or_path": ".",
-  "architectures": [
-    "T5ForConditionalGeneration"
-  ],
-  "d_ff": 3072,
-  "d_kv": 64,
-  "d_model": 768,
-  "decoder_start_token_id": 0,
-  "dropout_rate": 0.1,
-  "eos_token_id": 1,
-  "feed_forward_proj": "relu",
-  "gradient_checkpointing": false,
-  "initializer_factor": 1.0,
-  "is_encoder_decoder": true,
-  "layer_norm_epsilon": 1e-06,
-  "model_type": "t5",
-  "n_positions": 512,
-  "num_decoder_layers": 12,
-  "num_heads": 12,
-  "num_layers": 12,
-  "output_past": true,
-  "pad_token_id": 0,
-  "relative_attention_num_buckets": 32,
-  "task_specific_params": {
-    "summarization": {
-      "early_stopping": true,
-      "length_penalty": 2.0,
-      "max_length": 200,
-      "min_length": 30,
-      "no_repeat_ngram_size": 3,
-      "num_beams": 4,
-      "prefix": "summarize: "
-    },
-    "translation_en_to_de": {
-      "early_stopping": true,
-      "max_length": 300,
-      "num_beams": 4,
-      "prefix": "translate English to German: "
-    },
-    "translation_en_to_fr": {
-      "early_stopping": true,
-      "max_length": 300,
-      "num_beams": 4,
-      "prefix": "translate English to French: "
-    },
-    "translation_en_to_ro": {
-      "early_stopping": true,
-      "max_length": 300,
-      "num_beams": 4,
-      "prefix": "translate English to Romanian: "
-    }
-  },
-  "torch_dtype": "float32",
-  "transformers_version": "4.9.0.dev0",
-  "use_cache": true,
-  "vocab_size": 32103
-}

output/ckpt-9999/flax_model.msgpack

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8299c056e5ca07f93db2db052d61cb941710e0925c62486ee0c9775116e0a6bf
-size 891548548

output/ckpt-9999/opt_state.msgpack

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7c912197fd24feea06a22802e5bfbd9935100bb392a8d1966e230891aeaec658
-size 1783097336

output/ckpt-9999/training_state.json

@@ -1 +0,0 @@
-{"step": 10000}

output/events.out.tfevents.1626504033.t1v-n-0e7426e8-w-0.89661.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5b3b9f725bfa1e9befedd29c8c0319001a6ddc3597c6dfa30c754913531f26bc
-size 40

output/events.out.tfevents.1626504547.t1v-n-0e7426e8-w-0.93479.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1d3bd0981ae5d2bb0ac2ffef88a1eac66f198c1a58e207b37e216a9997428160
-size 40

output/events.out.tfevents.1626505238.t1v-n-0e7426e8-w-0.95128.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6b26b65cf4c438d4270d906ca1ed332fbe65f924ae82a22289dda08f95d5919f
-size 40

output/events.out.tfevents.1626506421.t1v-n-0e7426e8-w-0.96635.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d1a031810dfc4c6e7c52913e5261afc3fa3d5cf5a68695b76bbffd177b065e27
-size 40

output/events.out.tfevents.1626507299.t1v-n-0e7426e8-w-0.98584.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4874aada340bc85728ebb4b7f8329a0eb6618a19f0b646abb1f1b5f2e9fc84fe
-size 40

output/events.out.tfevents.1626508342.t1v-n-0e7426e8-w-0.101251.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ed227170f48707b46db61d657803869c0be10d350b75f29b0844a6ef8a9e0cd3
-size 40

output/flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8299c056e5ca07f93db2db052d61cb941710e0925c62486ee0c9775116e0a6bf
+oid sha256:f8e339d352658c4fae724883dc700cc559e7ab3eb7116139f6f0d187fe7720e1
 size 891548548

output/opt_state.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7c912197fd24feea06a22802e5bfbd9935100bb392a8d1966e230891aeaec658
+oid sha256:9f2eb4ce40eafe6435258b3761c281883b93221092ca701e0cd1f21b78264297
 size 1783097336

output/training_state.json

@@ -1 +1 @@
-{"step": 10000}
+{"step": 140001}

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:837e804cfcfee38ffdbb87dc80de834a7c5aec62634910e6b2514794f848bba2
+oid sha256:7a9f97e22703c1a5cf20353b9e859f377c5fa709e5a96ee15ad40d66674b67fa
 size 891650495

run.sh

@@ -1,12 +1,12 @@
 #!/bin/bash
-export CUDA_VISIBLE_DEVICES=1
+export CUDA_VISIBLE_DEVICES="1"
 
 MODEL="flax-community/t5-base-dutch"
 OUTPUT="./output"
 
-TRAIN="/home/yeb/cnnuxsum/cnnuxsum_train.json"
-VAL="/home/yeb/cnnuxsum/cnnuxsum_val.json"
-TEST="/home/yeb/cnnuxsum/cnnuxsum_test.json"
+TRAIN="/home/yeb/Developer/data/cnnuxsum/cnnuxsum_train.json"
+VAL="/home/yeb/Developer/data/cnnuxsum/cnnuxsum_val.json"
+TEST="/home/yeb/Developer/data/cnnuxsum/cnnuxsum_test.json"
 
 mkdir -p "${OUTPUT}"
 
@@ -15,48 +15,28 @@ python ./run_summarization_flax.py \
     --learning_rate "5e-4" \
     --warmup_steps 500 \
     --do_train \
+    --do_predict \
+    --do_eval \
     --train_file "${TRAIN}" \
     --validation_file "${VAL}" \
     --test_file "${TEST}" \
-    --max_train_samples 640000 \
-    --max_eval_samples 512 \
-    --max_predict_samples 64 \
+    --max_train_samples 1366592 \
+    --max_eval_samples 32 \
+    --max_predict_samples 8 \
     --text_column "complete_text" \
     --summary_column "summary_text" \
-    --source_prefix "summarize: " \
     --max_source_length 1024 \
     --max_target_length 142 \
     --output_dir "${OUTPUT}" \
     --per_device_train_batch_size=8 \
-    --per_device_eval_batch_size=2 \
+    --per_device_eval_batch_size=8 \
     --overwrite_output_dir \
     --num_train_epochs="1" \
-    --logging_steps="50" \
-    --save_steps="2000" \
-    --eval_steps="25000000" \
-    --num_beams 4
-
-#     \
-#    --do_predict
-#    --do_eval \
-
-
-#     \
-#    --prediction_debug \
-#    --predict_with_generate
-
-
-
+    --logging_steps="100" \
+    --save_steps="20000" \
+    --eval_steps="5000" \
+    --num_beams 4 \
+    --prediction_debug \
+    --predict_with_generate
 
 #    --source_prefix "summarize: " \
-
-#      --lr_scheduler_type="constant" \
-
-#    --task "summarization" \
-#      --early_stopping "true" \
-#      --length_penalty "2.0" \
-#      --max_length 300 \
-#      --min_length 75 \
-#      --no_repeat_ngram_size 3 \
-#      --num_beams 4 \
-#      --prefix "summarize: " \

run_summarization_flax.py

@@ -90,20 +90,34 @@ class ModelArguments:
     )
     model_type: Optional[str] = field(
         default=None,
-        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
+        metadata={
+            "help": "If training from scratch, pass a model type from the list: "
+            + ", ".join(MODEL_TYPES)
+        },
     )
     config_name: Optional[str] = field(
-        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+        default=None,
+        metadata={
+            "help": "Pretrained config name or path if not the same as model_name"
+        },
     )
     tokenizer_name: Optional[str] = field(
-        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+        default=None,
+        metadata={
+            "help": "Pretrained tokenizer name or path if not the same as model_name"
+        },
     )
     cache_dir: Optional[str] = field(
-        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+        default=None,
+        metadata={
+            "help": "Where do you want to store the pretrained models downloaded from s3"
+        },
     )
     use_fast_tokenizer: bool = field(
         default=True,
-        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+        metadata={
+            "help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."
+        },
     )
     dtype: Optional[str] = field(
         default="float32",
@@ -120,27 +134,41 @@ class DataTrainingArguments:
     """
 
     dataset_name: Optional[str] = field(
-        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+        default=None,
+        metadata={"help": "The name of the dataset to use (via the datasets library)."},
     )
     dataset_config_name: Optional[str] = field(
-        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+        default=None,
+        metadata={
+            "help": "The configuration name of the dataset to use (via the datasets library)."
+        },
     )
     text_column: Optional[str] = field(
         default=None,
-        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
+        metadata={
+            "help": "The name of the column in the datasets containing the full texts (for summarization)."
+        },
     )
     summary_column: Optional[str] = field(
         default=None,
-        metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."},
+        metadata={
+            "help": "The name of the column in the datasets containing the summaries (for summarization)."
+        },
+    )
+    train_file: Optional[str] = field(
+        default=None, metadata={"help": "The input training data file (a text file)."}
     )
-    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
     validation_file: Optional[str] = field(
         default=None,
-        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+        metadata={
+            "help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."
+        },
     )
     test_file: Optional[str] = field(
         default=None,
-        metadata={"help": "An optional input evaluation data file to predict the perplexity on (a text file)."},
+        metadata={
+            "help": "An optional input evaluation data file to predict the perplexity on (a text file)."
+        },
     )
     max_source_length: Optional[int] = field(
         default=1024,
@@ -191,10 +219,16 @@ class DataTrainingArguments:
         metadata={"help": "The number of processes to use for the preprocessing."},
     )
     source_prefix: Optional[str] = field(
-        default=None, metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
+        default=None,
+        metadata={
+            "help": "A prefix to add before every source text (useful for T5 models)."
+        },
     )
     predict_with_generate: bool = field(
-        default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
+        default=False,
+        metadata={
+            "help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."
+        },
     )
     num_beams: Optional[int] = field(
         default=None,
@@ -204,52 +238,52 @@ class DataTrainingArguments:
         },
     )
     overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
     )
     prediction_debug: bool = field(
         default=False,
-        metadata={
-            "help": "Whether to show some examples of the model prediction"
-        },
+        metadata={"help": "Whether to show some examples of the model prediction"},
     )
 
     def __post_init__(self):
-        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
+        if (
+            self.dataset_name is None
+            and self.train_file is None
+            and self.validation_file is None
+        ):
+            raise ValueError(
+                "Need either a dataset name or a training/validation file."
+            )
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]
-                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
+                assert extension in [
+                    "csv",
+                    "json",
+                ], "`train_file` should be a csv or a json file."
             if self.validation_file is not None:
                 extension = self.validation_file.split(".")[-1]
-                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
+                assert extension in [
+                    "csv",
+                    "json",
+                ], "`validation_file` should be a csv or a json file."
         if self.val_max_target_length is None:
             self.val_max_target_length = self.max_target_length
 
 
-summarization_name_mapping = {
-    "amazon_reviews_multi": ("review_body", "review_title"),
-    "big_patent": ("description", "abstract"),
-    "cnn_dailymail": ("article", "highlights"),
-    "orange_sum": ("text", "summary"),
-    "pn_summary": ("article", "summary"),
-    "psc": ("extract_text", "summary_text"),
-    "samsum": ("dialogue", "summary"),
-    "thaisum": ("body", "summary"),
-    "xglue": ("news_body", "news_title"),
-    "xsum": ("document", "summary"),
-    "wiki_summary": ("article", "highlights"),
-}
-
-
 class TrainState(train_state.TrainState):
     dropout_rng: jnp.ndarray
 
     def replicate(self):
-        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+        return jax_utils.replicate(self).replace(
+            dropout_rng=shard_prng_key(self.dropout_rng)
+        )
 
 
-def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
+def data_loader(
+    rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False
+):
     """
     Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
     Shuffle batches if `shuffle` is `True`.
@@ -273,7 +307,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
         yield batch
 
 
-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)
@@ -282,21 +316,35 @@ 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)
 
 
 def create_learning_rate_fn(
-    train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
+    train_ds_size: int,
+    train_batch_size: int,
+    num_train_epochs: int,
+    num_warmup_steps: int,
+    learning_rate: float,
 ) -> Callable[[int], jnp.array]:
     """Returns a linear warmup, linear_decay learning rate function."""
     steps_per_epoch = train_ds_size // train_batch_size
     num_train_steps = steps_per_epoch * num_train_epochs
-    warmup_fn = optax.linear_schedule(init_value=learning_rate, end_value=learning_rate, transition_steps=num_warmup_steps)
+    warmup_fn = optax.linear_schedule(
+        init_value=learning_rate,
+        end_value=learning_rate,
+        transition_steps=num_warmup_steps,
+    )
     decay_fn = optax.linear_schedule(
-        init_value=learning_rate, end_value=learning_rate, transition_steps=num_train_steps - num_warmup_steps
+        init_value=learning_rate,
+        end_value=learning_rate,
+        transition_steps=num_train_steps - num_warmup_steps,
+    )
+    schedule_fn = optax.join_schedules(
+        schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]
     )
-    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
 
     return schedule_fn
 
@@ -306,11 +354,15 @@ def main():
     # or by passing the --help flag to this script.
     # We now keep distinct sets of args, for a cleaner separation of concerns.
 
-    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    parser = HfArgumentParser(
+        (ModelArguments, DataTrainingArguments, TrainingArguments)
+    )
     if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
         # If we pass only one argument to the script and it's the path to a json file,
         # let's parse it to get our arguments.
-        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+        model_args, data_args, training_args = parser.parse_json_file(
+            json_file=os.path.abspath(sys.argv[1])
+        )
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
@@ -334,11 +386,7 @@ def main():
         state = jax_utils.unreplicate(state)
         logger.info(f"SAVING CHECKPOINT IN {save_dir}")
         save_dir = f"{save_dir}/ckpt-{mb_item(state.step) - 1}"
-        model.save_pretrained(
-            save_dir,
-            params=state.params,
-            push_to_hub=False
-        )
+        model.save_pretrained(save_dir, params=state.params, push_to_hub=False)
         if with_opt:
             with open(os.path.join(save_dir, "opt_state.msgpack"), "wb") as f:
                 f.write(to_bytes(state.opt_state))
@@ -352,9 +400,13 @@ def main():
         #     commit_message=f"Saving weights and logs of step {cur_step}",
         # )
         if with_opt:
-            with open(os.path.join(training_args.output_dir, "opt_state.msgpack"), "wb") as f:
+            with open(
+                os.path.join(training_args.output_dir, "opt_state.msgpack"), "wb"
+            ) as f:
                 f.write(to_bytes(state.opt_state))
-            with open(os.path.join(training_args.output_dir, "training_state.json"), "w") as f:
+            with open(
+                os.path.join(training_args.output_dir, "training_state.json"), "w"
+            ) as f:
                 json.dump({"step": state.step.item()}, f)
         logger.info("checkpoint saved")
 
@@ -386,7 +438,10 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         dataset = load_dataset(
-            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            cache_dir=model_args.cache_dir,
+            keep_in_memory=False,
         )
     else:
         data_files = {}
@@ -399,27 +454,37 @@ def main():
         if data_args.test_file is not None:
             data_files["test"] = data_args.test_file
             extension = data_args.test_file.split(".")[-1]
-        dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
+        dataset = load_dataset(
+            extension, data_files=data_files, cache_dir=model_args.cache_dir
+        )
     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
     # https://huggingface.co/docs/datasets/loading_datasets.html.
 
     # Load pretrained model and tokenizer
 
     if model_args.config_name:
-        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
+        config = AutoConfig.from_pretrained(
+            model_args.config_name, cache_dir=model_args.cache_dir
+        )
     elif model_args.model_name_or_path:
-        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path, cache_dir=model_args.cache_dir
+        )
     else:
         config = CONFIG_MAPPING[model_args.model_type]()
         logger.warning("You are instantiating a new config instance from scratch.")
 
     if model_args.tokenizer_name:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.tokenizer_name,
+            cache_dir=model_args.cache_dir,
+            use_fast=model_args.use_fast_tokenizer,
         )
     elif model_args.model_name_or_path:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.model_name_or_path,
+            cache_dir=model_args.cache_dir,
+            use_fast=model_args.use_fast_tokenizer,
         )
     else:
         raise ValueError(
@@ -429,7 +494,10 @@ def main():
 
     if model_args.model_name_or_path:
         model = FlaxAutoModelForSeq2SeqLM.from_pretrained(
-            model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
+            model_args.model_name_or_path,
+            config=config,
+            seed=training_args.seed,
+            dtype=getattr(jnp, model_args.dtype),
         )
     else:
         model = FlaxAutoModelForSeq2SeqLM.from_config(
@@ -437,7 +505,9 @@ def main():
         )
 
     if model.config.decoder_start_token_id is None:
-        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+        raise ValueError(
+            "Make sure that `config.decoder_start_token_id` is correctly defined"
+        )
 
     prefix = data_args.source_prefix if data_args.source_prefix is not None else ""
 
@@ -450,13 +520,14 @@ def main():
     elif training_args.do_predict:
         column_names = dataset["test"].column_names
     else:
-        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
+        logger.info(
+            "There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`."
+        )
         return
 
     # Get the column names for input/target.
-    dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
     if data_args.text_column is None:
-        text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
+        text_column = column_names[0]
     else:
         text_column = data_args.text_column
         if text_column not in column_names:
@@ -464,7 +535,7 @@ def main():
                 f"--text_column' value '{data_args.text_column}' needs to be one of: {', '.join(column_names)}"
             )
     if data_args.summary_column is None:
-        summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
+        summary_column = column_names[1]
     else:
         summary_column = data_args.summary_column
         if summary_column not in column_names:
@@ -487,18 +558,28 @@ def main():
         targets = examples[summary_column]
         inputs = [prefix + inp for inp in inputs]
         model_inputs = tokenizer(
-            inputs, max_length=data_args.max_source_length, padding="max_length", truncation=True, return_tensors="np"
+            inputs,
+            max_length=data_args.max_source_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="np",
         )
 
         # Setup the tokenizer for targets
         with tokenizer.as_target_tokenizer():
             labels = tokenizer(
-                targets, max_length=max_target_length, padding="max_length", truncation=True, return_tensors="np"
+                targets,
+                max_length=max_target_length,
+                padding="max_length",
+                truncation=True,
+                return_tensors="np",
             )
 
         model_inputs["labels"] = labels["input_ids"]
         decoder_input_ids = shift_tokens_right_fn(
-            jnp.array(labels["input_ids"]), config.pad_token_id, config.decoder_start_token_id
+            jnp.array(labels["input_ids"]),
+            config.pad_token_id,
+            config.decoder_start_token_id,
         )
         model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)
 
@@ -544,7 +625,9 @@ def main():
             raise ValueError("--do_predict requires a test dataset")
         predict_dataset = dataset["test"]
         if data_args.max_predict_samples is not None:
-            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
+            predict_dataset = predict_dataset.select(
+                range(data_args.max_predict_samples)
+            )
         predict_dataset = predict_dataset.map(
             preprocess_function,
             batched=True,
@@ -553,6 +636,14 @@ def main():
             load_from_cache_file=not data_args.overwrite_cache,
             desc="Running tokenizer on prediction dataset",
         )
+        eval_batch_size = (
+            int(training_args.per_device_eval_batch_size) * jax.device_count()
+        )
+        pred_steps = len(predict_dataset) // eval_batch_size
+        if pred_steps == 0:
+            raise Exception(
+                "The length of the prediction dataset // eval batch size is 0. Increase prediction dataset size"
+            )
 
     # Metric
     metric = load_metric("rouge")
@@ -578,13 +669,28 @@ def main():
             for index in random.sample(range(len(decoded_labels)), 3):
                 logger.info(f'reference: "{decoded_labels[index]}"')
                 logger.info(f'predicted: "{decoded_preds[index]}"')
-                logger.info('---')
+                logger.info("---")
 
-        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+        result = metric.compute(
+            predictions=decoded_preds, references=decoded_labels, use_stemmer=True
+        )
         # Extract a few results from ROUGE
         result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
 
-        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+        try:
+            result_blue = bleu.compute(
+                predictions=decoded_preds, references=decoded_labels_bleu
+            )
+            result_blue = result_blue["score"]
+        except Exception as e:
+            logger.info(f"Error occurred during bleu {e}")
+            result_blue = 0.0 * 100
+        result["blue"] = result_blue
+
+
+        prediction_lens = [
+            np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds
+        ]
         result["gen_len"] = np.mean(prediction_lens)
         result = {k: round(v, 4) for k, v in result.items()}
         return result
@@ -595,7 +701,7 @@ def main():
         try:
             from flax.metrics.tensorboard import SummaryWriter
 
-            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
+            summary_writer = SummaryWriter(log_dir=Path(training_args.logging_dir))
         except ImportError as ie:
             has_tensorboard = False
             logger.warning(
@@ -613,7 +719,9 @@ def main():
 
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    train_batch_size = (
+        int(training_args.per_device_train_batch_size) * jax.device_count()
+    )
     eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
     steps_per_epoch = len(train_dataset) // train_batch_size
     total_train_steps = steps_per_epoch * num_epochs
@@ -634,13 +742,36 @@ def main():
     # Note that this mask is specifically adapted for FlaxBart.
     # For FlaxT5, one should correct the layer norm parameter naming
     # accordingly - see `run_t5_mlm_flax.py` e.g.
-    def decay_mask_fn(params):
-        flat_params = traverse_util.flatten_dict(params)
-        layer_norm_params = [
-            (name, "scale") for name in ["self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"]
-        ]
-        flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params}
-        return traverse_util.unflatten_dict(flat_mask)
+    if config.model_type in ["t5", "mt5", "byt5"]:
+
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale") for name in ["layer_norm", "final_layer_norm"]
+            ]
+            flat_mask = {
+                path: (path[-1] != "bias" and path[-2:] not in layer_norm_params)
+                for path in flat_params
+            }
+            return traverse_util.unflatten_dict(flat_mask)
+
+    else:
+
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale")
+                for name in [
+                    "self_attn_layer_norm",
+                    "layernorm_embedding",
+                    "final_layer_norm",
+                ]
+            ]
+            flat_mask = {
+                path: (path[-1] != "bias" and path[-2:] not in layer_norm_params)
+                for path in flat_params
+            }
+            return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
     adamw = optax.adamw(
@@ -653,7 +784,9 @@ def main():
     )
 
     # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
+    state = TrainState.create(
+        apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng
+    )
 
     # label smoothed cross entropy
     def loss_fn(logits, labels, padding_mask, label_smoothing_factor=0.0):
@@ -665,9 +798,12 @@ def main():
         confidence = 1.0 - label_smoothing_factor
         low_confidence = (1.0 - confidence) / (vocab_size - 1)
         normalizing_constant = -(
-            confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+            confidence * jnp.log(confidence)
+            + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+        )
+        soft_labels = onehot(
+            labels, vocab_size, on_value=confidence, off_value=low_confidence
         )
-        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
 
         loss = optax.softmax_cross_entropy(logits, soft_labels)
         loss = loss - normalizing_constant
@@ -683,8 +819,12 @@ def main():
 
         def compute_loss(params):
             labels = batch.pop("labels")
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-            loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor)
+            logits = state.apply_fn(
+                **batch, params=params, dropout_rng=dropout_rng, train=True
+            )[0]
+            loss = loss_fn(
+                logits, labels, batch["decoder_attention_mask"], label_smoothing_factor
+            )
             return loss
 
         grad_fn = jax.value_and_grad(compute_loss)
@@ -693,7 +833,10 @@ def main():
 
         new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
 
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics = {
+            "loss": loss,
+            "learning_rate": linear_decay_lr_schedule_fn(state.step),
+        }
         metrics = jax.lax.pmean(metrics, axis_name="batch")
 
         return new_state, metrics
@@ -702,7 +845,9 @@ def main():
     def eval_step(params, batch, label_smoothing_factor=0.0):
         labels = batch.pop("labels")
         logits = model(**batch, params=params, train=False)[0]
-        loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor)
+        loss = loss_fn(
+            logits, labels, batch["decoder_attention_mask"], label_smoothing_factor
+        )
 
         # summarize metrics
         metrics = {"loss": loss}
@@ -711,21 +856,36 @@ def main():
 
     # Define generation function
     max_length = (
-        data_args.val_max_target_length if data_args.val_max_target_length is not None else model.config.max_length
+        data_args.val_max_target_length
+        if data_args.val_max_target_length is not None
+        else model.config.max_length
+    )
+    num_beams = (
+        data_args.num_beams
+        if data_args.num_beams is not None
+        else model.config.num_beams
     )
-    num_beams = data_args.num_beams if data_args.num_beams is not None else model.config.num_beams
     gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
 
     def generate_step(params, batch):
         model.params = params
-        output_ids = model.generate(batch["input_ids"], attention_mask=batch["attention_mask"], **gen_kwargs)
+        output_ids = model.generate(
+            batch["input_ids"], attention_mask=batch["attention_mask"], **gen_kwargs
+        )
         return output_ids.sequences
 
     # Create parallel version of the train and eval step
     p_train_step = jax.pmap(
-        partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,)
+        partial(
+            train_step, label_smoothing_factor=training_args.label_smoothing_factor
+        ),
+        "batch",
+        donate_argnums=(0,),
+    )
+    p_eval_step = jax.pmap(
+        partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor),
+        "batch",
     )
-    p_eval_step = jax.pmap(partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
     p_generate_step = jax.pmap(generate_step, "batch")
 
     # Replicate the train state on each device
@@ -734,11 +894,16 @@ def main():
     logger.info("***** Running training *****")
     logger.info(f"  Num examples = {len(train_dataset)}")
     logger.info(f"  Num Epochs = {num_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
+    logger.info(
+        f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}"
+    )
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size}"
+    )
     logger.info(f"  Total optimization steps = {total_train_steps}")
 
     train_time = 0
+    train_metrics = []
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
     for epoch in epochs:
         # ======================== Training ================================
@@ -746,117 +911,160 @@ def main():
 
         # Create sampling rng
         rng, input_rng = jax.random.split(rng)
-        train_metrics = []
 
         # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
+        train_loader = data_loader(
+            input_rng, train_dataset, train_batch_size, shuffle=True
+        )
         steps_per_epoch = len(train_dataset) // train_batch_size
         # train
-        for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
+        for step in tqdm(
+            range(steps_per_epoch), desc="Training...", position=1, leave=False
+        ):
             batch = next(train_loader)
             state, train_metric = p_train_step(state, batch)
             train_metrics.append(train_metric)
 
-        train_time += time.time() - train_start
-
-        train_metric = unreplicate(train_metric)
-
-        epochs.write(
-            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
-        )
-
-        # save checkpoint after each epoch and push checkpoint to the hub
-        if jax.process_index() == 0:
-            # params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
-            # model.save_pretrained(
-            #     training_args.output_dir,
-            #     params=params,
-            #     push_to_hub=training_args.push_to_hub,
-            #     commit_message=f"Saving weights and logs of epoch {epoch+1}",
-            # )
-            save_checkpoint(model, training_args.output_dir, state)
-
-        # ======================== Evaluating ==============================
-        if training_args.do_eval:
-            eval_metrics = []
-            eval_preds = []
-            eval_labels = []
-
-            eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
-            eval_steps = len(eval_dataset) // eval_batch_size
-            for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
-                # Model forward
-                batch = next(eval_loader)
-                labels = batch["labels"]
-
-                metrics = p_eval_step(state.params, batch)
-                eval_metrics.append(metrics)
-
-                # generation
+            cur_step = epoch * (len(train_dataset) // train_batch_size) + step
+
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
+                # Save metrics
+                train_metric = unreplicate(train_metric)
+                train_time += time.time() - train_start
+
+                if has_tensorboard and jax.process_index() == 0:
+                    logger.info(
+                        f"*** Writing training summary after {cur_step} steps ***"
+                    )
+                    write_train_metric(
+                        summary_writer, train_metrics, train_time, cur_step
+                    )
+
+                epochs.write(
+                    f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
+                )
+
+                train_metrics = []
+
+            if (
+                training_args.do_eval
+                and cur_step % training_args.eval_steps == 0
+                and cur_step > 0
+            ):
+                logger.info(f"*** Evaluation after {cur_step} steps ***")
+                eval_metrics = []
+                eval_preds = []
+                eval_labels = []
+
+                eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
+                eval_steps = len(eval_dataset) // eval_batch_size
+                for _ in tqdm(
+                    range(eval_steps), desc="Evaluating...", position=2, leave=False
+                ):
+                    # Model forward
+                    batch = next(eval_loader)
+                    labels = batch["labels"]
+
+                    metrics = p_eval_step(state.params, batch)
+                    eval_metrics.append(metrics)
+
+                    # generation
+                    if data_args.predict_with_generate:
+                        generated_ids = p_generate_step(state.params, batch)
+                        eval_preds.extend(
+                            jax.device_get(
+                                generated_ids.reshape(-1, gen_kwargs["max_length"])
+                            )
+                        )
+                        eval_labels.extend(
+                            jax.device_get(labels.reshape(-1, labels.shape[-1]))
+                        )
+
+                # normalize eval metrics
+                eval_metrics = get_metrics(eval_metrics)
+                eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
+
+                # compute several metrics
+                mix_desc = ""
                 if data_args.predict_with_generate:
-                    generated_ids = p_generate_step(state.params, batch)
-                    eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                    eval_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
-
-            # normalize eval metrics
-            eval_metrics = get_metrics(eval_metrics)
-            eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
-
-            # compute ROUGE metrics
-            rouge_desc = ""
-            if data_args.predict_with_generate:
-                rouge_metrics = compute_metrics(eval_preds, eval_labels)
-                eval_metrics.update(rouge_metrics)
-                rouge_desc = " ".join([f"Eval {key}: {value} |" for key, value in rouge_metrics.items()])
-
-            # Print metrics and update progress bar
-            desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {rouge_desc})"
-            epochs.write(desc)
-            epochs.desc = desc
-
-            # Save metrics
-            if has_tensorboard and jax.process_index() == 0:
-                cur_step = epoch * (len(train_dataset) // train_batch_size)
-                write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
-
-    # ======================== Prediction loop ==============================
-    if training_args.do_predict:
-        logger.info("*** Predict ***")
-
-        pred_metrics = []
-        pred_generations = []
-        pred_labels = []
-
-        pred_loader = data_loader(input_rng, predict_dataset, eval_batch_size)
-        pred_steps = len(predict_dataset) // eval_batch_size
-        for _ in tqdm(range(pred_steps), desc="Predicting...", position=2, leave=False):
-            # Model forward
-            batch = next(pred_loader)
-            labels = batch["labels"]
-
-            metrics = p_eval_step(state.params, batch)
-            pred_metrics.append(metrics)
-
-            # generation
-            if data_args.predict_with_generate:
-                generated_ids = p_generate_step(state.params, batch)
-                pred_generations.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                pred_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
-
-        # normalize prediction metrics
-        pred_metrics = get_metrics(pred_metrics)
-        pred_metrics = jax.tree_map(jnp.mean, pred_metrics)
-
-        # compute ROUGE metrics
-        rouge_desc = ""
-        if data_args.predict_with_generate:
-            rouge_metrics = compute_metrics(pred_generations, pred_labels)
-            pred_metrics.update(rouge_metrics)
-            rouge_desc = " ".join([f"Predict {key}: {value} |" for key, value in rouge_metrics.items()])
-
-        # Print metrics
-        desc = f"Predict Loss: {pred_metrics['loss']} | {rouge_desc})"
-        logger.info(desc)
+                    mix_metrics = compute_metrics(eval_preds, eval_labels)
+                    eval_metrics.update(mix_metrics)
+                    mix_desc = " ".join(
+                        [f"Eval {key}: {value} |" for key, value in mix_metrics.items()]
+                    )
+
+                # Print metrics and update progress bar
+                desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {mix_desc} )"
+                epochs.write(desc)
+                epochs.desc = desc
+
+                # Save metrics
+                if has_tensorboard and jax.process_index() == 0:
+                    logger.info(
+                        f"*** Writing evaluation summary after {cur_step} steps ***"
+                    )
+                    # cur_step = epoch * (len(train_dataset) // train_batch_size)
+                    write_eval_metric(summary_writer, eval_metrics, cur_step)
+
+                # ======================== Prediction loop ==============================
+                if training_args.do_predict:
+                    logger.info("*** Predict ***")
+
+                    pred_metrics = []
+                    pred_generations = []
+                    pred_labels = []
+
+                    pred_loader = data_loader(
+                        input_rng, predict_dataset, eval_batch_size
+                    )
+                    pred_steps = len(predict_dataset) // eval_batch_size
+                    for _ in tqdm(
+                        range(pred_steps), desc="Predicting...", position=2, leave=False
+                    ):
+                        # Model forward
+                        batch = next(pred_loader)
+                        labels = batch["labels"]
+
+                        metrics = p_eval_step(state.params, batch)
+                        pred_metrics.append(metrics)
+
+                        # generation
+                        if data_args.predict_with_generate:
+                            generated_ids = p_generate_step(state.params, batch)
+                            pred_generations.extend(
+                                jax.device_get(
+                                    generated_ids.reshape(-1, gen_kwargs["max_length"])
+                                )
+                            )
+                            pred_labels.extend(
+                                jax.device_get(labels.reshape(-1, labels.shape[-1]))
+                            )
+
+                    # normalize prediction metrics
+                    pred_metrics = get_metrics(pred_metrics)
+                    pred_metrics = jax.tree_map(jnp.mean, pred_metrics)
+
+                    # compute ROUGE metrics
+                    rouge_desc = ""
+                    if data_args.predict_with_generate:
+                        rouge_metrics = compute_metrics(pred_generations, pred_labels)
+                        pred_metrics.update(rouge_metrics)
+                        rouge_desc = " ".join(
+                            [
+                                f"Predict {key}: {value} |"
+                                for key, value in rouge_metrics.items()
+                            ]
+                        )
+
+                    # Print metrics
+                    desc = f"Predict Loss: {pred_metrics['loss']} | {rouge_desc})"
+                    logger.info(desc)
+
+            if cur_step % training_args.save_steps == 0 and cur_step > 0:
+                logger.info(f"*** Saving checkpoints after {cur_step} steps ***")
+                # save checkpoint after each steps and push checkpoint to the hub
+                if jax.process_index() == 0:
+                    save_checkpoint(model, training_args.output_dir, state)
 
     # save checkpoint after each epoch and push checkpoint to the hub
     if jax.process_index() == 0:
@@ -867,8 +1075,6 @@ def main():
             push_to_hub=training_args.push_to_hub,
             commit_message=f"Saving weights and logs of epoch {epoch+1}",
         )
-        # save_checkpoint(model, training_args.output_dir, state)
-
 
 
 if __name__ == "__main__":