From eab9779732a4c61e50489c266c0450c238af131d Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sat, 13 Jun 2020 20:12:11 -0400
Subject: [PATCH 01/18] better mask logic

---
 examples/summarization/distillation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/summarization/distillation.py b/examples/summarization/distillation.py
index 8b8cbb4e..529b3475 100644
--- a/examples/summarization/distillation.py
+++ b/examples/summarization/distillation.py
@@ -246,7 +246,7 @@ class SummarizationDistiller(SummarizationTrainer):
                 lm_labels=labels,
                 output_hidden_states=True,
             )
-        dec_mask = invert_mask(self.model.model.last_padding_mask)
+        dec_mask = decoder_input_ids.eq(self.tokenizer.pad_token_id)
         loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, slogits, tlogits)
         if not self.hparams.freeze_decoder and self.alpha_hid > 0:
             hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.layer_to_copy)

From 350eeb7cd29a19f5a59b6771a021185e7edaf2db Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 13:17:09 -0400
Subject: [PATCH 02/18] style

---
 examples/summarization/distillation.py                | 1 -
 examples/summarization/test_summarization_examples.py | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/summarization/distillation.py b/examples/summarization/distillation.py
index 529b3475..30c22326 100644
--- a/examples/summarization/distillation.py
+++ b/examples/summarization/distillation.py
@@ -11,7 +11,6 @@ from torch.nn import functional as F
 
 from lightning_base import generic_train
 from transformers import AdamW, BartConfig, BartForConditionalGeneration, T5Config, T5ForConditionalGeneration
-from transformers.modeling_bart import invert_mask
 
 
 try:
diff --git a/examples/summarization/test_summarization_examples.py b/examples/summarization/test_summarization_examples.py
index 48692a2b..3234eb46 100644
--- a/examples/summarization/test_summarization_examples.py
+++ b/examples/summarization/test_summarization_examples.py
@@ -7,7 +7,6 @@ import unittest
 from pathlib import Path
 from unittest.mock import patch
 
-import pandas as pd
 import torch
 from torch.utils.data import DataLoader
 
@@ -244,6 +243,8 @@ class TestSummarizationDistiller(unittest.TestCase):
         # self.assertEqual(len(contents), 15)
 
         metrics = pickle_load(Path(output_dir) / "metrics.pkl")
+        import pandas as pd
+
         val_df = pd.DataFrame(metrics["val"])
         train_df = pd.DataFrame(metrics["train"])
         test_df = pd.DataFrame(metrics["test"])

From 3e62d96b14ddfc99e3e9acbfb46a27484d8f385f Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 14:28:20 -0400
Subject: [PATCH 03/18] add git-python requirement

---
 examples/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/requirements.txt b/examples/requirements.txt
index 8c1c67c9..cb248bdb 100644
--- a/examples/requirements.txt
+++ b/examples/requirements.txt
@@ -7,3 +7,4 @@ rouge-score
 tensorflow_datasets
 pytorch-lightning==0.7.6  # April 10, 2020 release
 matplotlib
+git-python==1.0.3

From 379e8c7dbaed7fec8eb906a8373028de3d4ea280 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 15:00:35 -0400
Subject: [PATCH 04/18] Cleanup

---
 examples/lightning_base.py             | 16 +++++------
 examples/summarization/callbacks.py    |  7 ++++-
 examples/summarization/distillation.py | 37 +++++++++++++-------------
 examples/summarization/finetune.py     |  5 ++--
 examples/summarization/utils.py        |  6 +++++
 5 files changed, 38 insertions(+), 33 deletions(-)

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index e08ad9e2..a84203c7 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -45,12 +45,6 @@ def set_seed(args: argparse.Namespace):
         torch.cuda.manual_seed_all(args.seed)
 
 
-def count_trainable_parameters(model):
-    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
-    params = sum([np.prod(p.size()) for p in model_parameters])
-    return params
-
-
 class BaseTransformer(pl.LightningModule):
     def __init__(
         self,
@@ -94,7 +88,7 @@ class BaseTransformer(pl.LightningModule):
                 cache_dir=cache_dir,
             )
         else:
-            self.model_type = type(model)
+            self.model_type = None
             self.model = model
 
     def load_hf_checkpoint(self, *args, **kwargs):
@@ -186,7 +180,7 @@ class BaseTransformer(pl.LightningModule):
         )
         parser.add_argument(
             "--tokenizer_name",
-            default="facebook/bart-large",
+            default=None,
             type=str,
             help="Pretrained tokenizer name or path if not the same as model_name",
         )
@@ -234,7 +228,7 @@ class LoggingCallback(pl.Callback):
 
 
 def add_generic_args(parser, root_dir):
-    parser = pl.Trainer.add_argparse_args(parser)
+    #  TODO(SS): allow all pl args? parser = pl.Trainer.add_argparse_args(parser)
     parser.add_argument(
         "--output_dir",
         default=None,
@@ -277,7 +271,7 @@ def generic_train(
     model: BaseTransformer,
     args: argparse.Namespace,
     early_stopping_callback=False,
-    logger=True,
+    logger=True,  # can pass WandbLogger() here
     extra_callbacks=[],
     checkpoint_callback=None,
     logging_callback=None,
@@ -302,6 +296,8 @@ def generic_train(
 
     if args.n_tpu_cores > 0:
         global xm
+        import torch_xla.core.xla_model as xm
+
         train_params["num_tpu_cores"] = args.n_tpu_cores
         train_params["gpus"] = 0
 
diff --git a/examples/summarization/callbacks.py b/examples/summarization/callbacks.py
index 79da7212..d58eacb2 100644
--- a/examples/summarization/callbacks.py
+++ b/examples/summarization/callbacks.py
@@ -2,12 +2,17 @@ import logging
 import os
 from pathlib import Path
 
+import numpy as np
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.utilities import rank_zero_only
 
-from lightning_base import count_trainable_parameters
+
+def count_trainable_parameters(model):
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    return params
 
 
 logger = logging.getLogger(__name__)
diff --git a/examples/summarization/distillation.py b/examples/summarization/distillation.py
index 30c22326..1d7e1f42 100644
--- a/examples/summarization/distillation.py
+++ b/examples/summarization/distillation.py
@@ -23,7 +23,7 @@ try:
     from .initialization_utils import init_student, copy_layers
     from .utils import SummarizationDataset, pickle_load
     from .finetune import main as ft_main
-except ModuleNotFoundError:
+except ImportError:
     from finetune import (
         SummarizationTrainer,
         freeze_params,
@@ -46,7 +46,8 @@ class SummarizationDistiller(SummarizationTrainer):
         super().__init__(hparams, model=student, config=student_cfg)
         self.teacher = teacher
         freeze_params(self.teacher)
-        self.freeze_stuff(d_layers_to_copy)
+        assert len(self.model.model.decoder.layers) == len(d_layers_to_copy)
+        self.sanity_check_gradients()
         self.ce_loss_fct = nn.KLDivLoss(reduction="batchmean")
         self.temperature = 2.0
         self.alpha_mlm = hparams.alpha_mlm
@@ -57,8 +58,7 @@ class SummarizationDistiller(SummarizationTrainer):
         gc.collect()
         torch.cuda.empty_cache()
 
-    def freeze_stuff(self, d_layers_to_copy):
-        assert len(self.model.model.decoder.layers) == len(d_layers_to_copy)
+    def sanity_check_gradients(self):
         assert_all_frozen(self.teacher)
         assert_all_frozen(self.model.model.decoder.embed_tokens)
         assert_all_frozen(self.model.model.encoder.embed_tokens)
@@ -77,7 +77,7 @@ class SummarizationDistiller(SummarizationTrainer):
         }
         d_layers_to_copy = get_layers_to_copy(student_updates["decoder_layers"], teacher.config.decoder_layers)
         e_layers_to_copy: List = get_layers_to_copy(student_updates["encoder_layers"], teacher.config.encoder_layers)
-        hparams.layer_to_copy = d_layers_to_copy
+        hparams.d_layer_to_copy = d_layers_to_copy
         hparams.e_layer_to_copy = e_layers_to_copy
         kw = teacher.config.to_diff_dict()
         kw.update(student_updates)
@@ -248,7 +248,7 @@ class SummarizationDistiller(SummarizationTrainer):
         dec_mask = decoder_input_ids.eq(self.tokenizer.pad_token_id)
         loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, slogits, tlogits)
         if not self.hparams.freeze_decoder and self.alpha_hid > 0:
-            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.layer_to_copy)
+            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_layer_to_copy)
 
         blended_loss = (
             self.alpha_ce * loss_ce
@@ -283,7 +283,7 @@ class T5SummarizationDistiller(SummarizationDistiller):
         d_layers_to_copy = get_layers_to_copy(n_layer, len(teacher.decoder.block))
         e_layers_to_copy: List = get_layers_to_copy(n_layer, len(teacher.encoder.block))
         student_updates = {"num_layers": n_layer}
-        hparams.layer_to_copy = d_layers_to_copy
+        hparams.d_layer_to_copy = d_layers_to_copy
         hparams.e_layer_to_copy = e_layers_to_copy
         kw = teacher.config.to_diff_dict()
 
@@ -304,7 +304,7 @@ class T5SummarizationDistiller(SummarizationDistiller):
         for d in [self.model.encoder, self.model.decoder]:
             freeze_params(d.embed_tokens)
 
-    def freeze_stuff(self, d_layers_to_copy):
+    def sanity_check_gradients(self, d_layers_to_copy):
         """T5"""
         assert len(self.model.decoder.block) == len(d_layers_to_copy)
         assert_all_frozen(self.teacher)
@@ -321,7 +321,6 @@ class T5SummarizationDistiller(SummarizationDistiller):
             freeze_params(self.model.decoder)  # TODO(SS): very suspicious
 
     def _step(self, batch):
-        # assert is_frozen(self.teacher)
         pad_token_id = self.tokenizer.pad_token_id
         source_ids, source_mask, y = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]
         decoder_input_ids = y[:, :-1].contiguous()
@@ -372,7 +371,7 @@ class T5SummarizationDistiller(SummarizationDistiller):
 
         loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, slogits, tlogits)
         if not self.hparams.freeze_decoder and self.alpha_hid > 0:
-            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.layer_to_copy)
+            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_layer_to_copy)
 
         blended_loss = (
             self.alpha_ce * loss_ce
@@ -383,15 +382,6 @@ class T5SummarizationDistiller(SummarizationDistiller):
         return blended_loss, loss_ce, sloss, loss_encoder, hid_loss_enc, hid_loss_dec
 
 
-def distill_main(args):
-    Path(args.output_dir).mkdir(exist_ok=True)
-    if len(os.listdir(args.output_dir)) > 3 and args.do_train:
-        raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
-
-    model = create_module(args)
-    ft_main(args, model=model)
-
-
 def create_module(args):
     t5 = "t5" in args.model_name_or_path
     if args.no_teacher:
@@ -448,6 +438,15 @@ def get_layers_to_copy(n_to_get, tot):
         return all_layers[:n_to_get]
 
 
+def distill_main(args):
+    Path(args.output_dir).mkdir(exist_ok=True)
+    if len(os.listdir(args.output_dir)) > 3 and args.do_train:
+        raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
+
+    model = create_module(args)
+    ft_main(args, model=model)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser = SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index e12b5c40..64ed286a 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -15,16 +15,15 @@ from rouge_score import rouge_scorer, scoring
 from torch import nn
 from torch.utils.data import DataLoader
 
-from durbango import pickle_save
 from lightning_base import BaseTransformer, add_generic_args, generic_train
 from transformers import AutoModelWithLMHead, get_linear_schedule_with_warmup
 
 
 try:
-    from .utils import SummarizationDataset, lmap, flatten_list
+    from .utils import SummarizationDataset, lmap, flatten_list, pickle_save
     from .callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 except ImportError:
-    from utils import SummarizationDataset, lmap, flatten_list
+    from utils import SummarizationDataset, lmap, flatten_list, pickle_save
     from callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 
 logger = logging.getLogger(__name__)
diff --git a/examples/summarization/utils.py b/examples/summarization/utils.py
index c40e5f68..43923d21 100644
--- a/examples/summarization/utils.py
+++ b/examples/summarization/utils.py
@@ -164,5 +164,11 @@ def pickle_load(path):
         return pickle.load(f)
 
 
+def pickle_save(obj, path):
+    """pickle.dump(obj, path)"""
+    with open(path, "wb") as f:
+        return pickle.dump(obj, f)
+
+
 def flatten_list(summary_ids: List[List]):
     return [x for x in itertools.chain.from_iterable(summary_ids)]

From 35a82ee87fb9263a308d717d8c268c603a1c6734 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 15:01:50 -0400
Subject: [PATCH 05/18] boom boom

---
 examples/summarization/finetune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 64ed286a..746cd27b 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -367,7 +367,7 @@ def main(args, model=None):
         model.hparams.test_checkpoint = checkpoints[-1]
         trainer.resume_from_checkpoint = checkpoints[-1]
     trainer.logger.log_hyperparams(model.hparams)
-    trainer.test(model)
+    trainer.test(model)  # NOTE(SS): this will break in DDP, known lightning issue. See evaluate_checkpoint
     return model
 
 

From 6bf996f58d32eb419ab506f36ff50eeaef52cf68 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 15:31:46 -0400
Subject: [PATCH 06/18] boom boom

---
 examples/summarization/finetune.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 746cd27b..d9f920c7 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -293,9 +293,8 @@ class SummarizationTrainer(BaseTransformer):
 
     @staticmethod
     def add_model_specific_args(parser, root_dir):
-        add_generic_args(parser, root_dir)
-
         BaseTransformer.add_model_specific_args(parser, root_dir)
+        add_generic_args(parser, root_dir)
         parser.add_argument(
             "--max_source_length",
             default=1024,
@@ -373,7 +372,6 @@ def main(args, model=None):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    add_generic_args(parser, os.getcwd())
     parser = SummarizationTrainer.add_model_specific_args(parser, os.getcwd())
     args = parser.parse_args()
 

From 5526be3499a6feecaf089f880ef83b8b8964fc85 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 15:35:47 -0400
Subject: [PATCH 07/18] boom boom

---
 examples/lightning_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index a84203c7..212eb777 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -227,7 +227,7 @@ class LoggingCallback(pl.Callback):
                         writer.write("{} = {}\n".format(key, str(metrics[key])))
 
 
-def add_generic_args(parser, root_dir):
+def add_generic_args(parser, root_dir) -> None:
     #  TODO(SS): allow all pl args? parser = pl.Trainer.add_argparse_args(parser)
     parser.add_argument(
         "--output_dir",

From eb84b9d5ee4e55cd7337eae35bc785e04a9aeebe Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 15:36:29 -0400
Subject: [PATCH 08/18] boom boom

---
 examples/lightning_base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index 212eb777..26d8e63e 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -265,6 +265,7 @@ def add_generic_args(parser, root_dir) -> None:
 
     parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
     parser.add_argument("--resume_from_checkpoint", type=str, default=None)
+    parser.add_argument("--val_check_interval", default=1.0, type=float)
 
 
 def generic_train(

From 4606bda735d7850ae720e1e4626b093ef2b2ec62 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 19:47:04 +0000
Subject: [PATCH 09/18] Bash cleanup

---
 examples/lightning_base.py              |  1 -
 examples/summarization/finetune_bart.sh | 19 ++++++++++++-------
 examples/summarization/run_distiller.sh |  2 +-
 3 files changed, 13 insertions(+), 9 deletions(-)
 mode change 100644 => 100755 examples/summarization/finetune_bart.sh

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index 26d8e63e..5c689549 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -318,7 +318,6 @@ def generic_train(
         val_check_interval=args.val_check_interval,
         weights_summary=None,
         resume_from_checkpoint=args.resume_from_checkpoint,
-        auto_scale_batch_size=args.auto_scale_batch_size,
         **train_params,
     )
 
diff --git a/examples/summarization/finetune_bart.sh b/examples/summarization/finetune_bart.sh
old mode 100644
new mode 100755
index b37888f5..d3a86ba1
--- a/examples/summarization/finetune_bart.sh
+++ b/examples/summarization/finetune_bart.sh
@@ -1,6 +1,5 @@
-export OUTPUT_DIR_NAME=bart_sum
-export CURRENT_DIR=${PWD}
-export OUTPUT_DIR=${CURRENT_DIR}/${OUTPUT_DIR_NAME}
+export OUTPUT_DIR_NAME=bart_cnn_finetune
+export OUTPUT_DIR=bart_cnn_finetune
 
 # Make output directory if it doesn't exist
 mkdir -p $OUTPUT_DIR
@@ -8,11 +7,17 @@ mkdir -p $OUTPUT_DIR
 # Add parent directory to python path to access lightning_base.py
 export PYTHONPATH="../":"${PYTHONPATH}"
 
-python finetune.py \
---data_dir=./cnn-dailymail/cnn_dm \
---model_name_or_path=bart-large \
+python -m pdb -c continue finetune.py \
+--data_dir=$CNN_DIR \
+--model_name_or_path=facebook/bart-large \
 --learning_rate=3e-5 \
 --train_batch_size=4 \
 --eval_batch_size=4 \
 --output_dir=$OUTPUT_DIR \
---do_train  $@
+--do_train \
+--do_predict \
+--n_val 1000 \
+--val_check_interval 0.1 \
+--sortish_sampler \
+--num_train_epochs 6 \
+$@
diff --git a/examples/summarization/run_distiller.sh b/examples/summarization/run_distiller.sh
index da4f99be..6fbecad3 100755
--- a/examples/summarization/run_distiller.sh
+++ b/examples/summarization/run_distiller.sh
@@ -3,7 +3,7 @@
 # Add parent directory to python path to access lightning_base.py
 export PYTHONPATH="../":"${PYTHONPATH}"
 
-python finetune.py \
+python distillation.py \
 --learning_rate=3e-4 \
 --do_train \
 --do_predict \

From 9f98a6ac15e2d0d519f2daedd21111e72a0ce6d0 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 16:07:25 -0400
Subject: [PATCH 10/18] docs

---
 examples/lightning_base.py                    |  3 ++
 examples/summarization/README.md              | 44 ++++++-------------
 examples/summarization/distillation.py        |  5 ++-
 examples/summarization/finetune.py            | 10 ++---
 .../{finetune_bart.sh => finetune.sh}         |  8 ++--
 .../{evaluate_cnn.py => run_eval.py}          | 38 ++++++----------
 .../test_summarization_examples.py            |  4 +-
 examples/summarization/utils.py               |  7 +++
 src/transformers/modeling_bart.py             |  2 +-
 9 files changed, 52 insertions(+), 69 deletions(-)
 rename examples/summarization/{finetune_bart.sh => finetune.sh} (71%)
 rename examples/summarization/{evaluate_cnn.py => run_eval.py} (69%)

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index 5c689549..d4f241c0 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -14,6 +14,7 @@ from transformers import (
     AutoModel,
     AutoModelForPreTraining,
     AutoModelForQuestionAnswering,
+    AutoModelForSeq2SeqLM,
     AutoModelForSequenceClassification,
     AutoModelForTokenClassification,
     AutoModelWithLMHead,
@@ -34,6 +35,8 @@ MODEL_MODES = {
     "pretraining": AutoModelForPreTraining,
     "token-classification": AutoModelForTokenClassification,
     "language-modeling": AutoModelWithLMHead,
+    "summarization": AutoModelForSeq2SeqLM,
+    "translation": AutoModelForSeq2SeqLM,
 }
 
 
diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index ad4adc4b..721c42b1 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -1,6 +1,6 @@
-### Get CNN Data
-To be able to reproduce the authors' results on the CNN/Daily Mail dataset you first need to download both CNN and Daily Mail datasets [from Kyunghyun Cho's website](https://cs.nyu.edu/~kcho/DMQA/) (the links next to "Stories") in the same folder. Then uncompress the archives by running:
+### Data
 
+CNN/DailyMail data
 ```bash
 wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/cnn_dm.tgz
 tar -xzvf cnn_dm.tgz
@@ -9,39 +9,21 @@ tar -xzvf cnn_dm.tgz
 this should make a directory called cnn_dm/ with files like `test.source`.
 To use your own data, copy that files format. Each article to be summarized is on its own line.
 
+XSUM Data:
+```bash
+wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/xsum.tar.gz
+tar -xzvf xsum.tar.gz
+```
+
+
 ### Evaluation
 
 To create summaries for each article in dataset, run:
 ```bash
-python evaluate_cnn.py <path_to_test.source> test_generations.txt <model-name>  --score_path rouge_scores.txt
+python run_eval.py <path_to_test.source> test_generations.txt <model-name>  --score_path rouge_scores.txt
 ```
-The default batch size, 8, fits in 16GB GPU memory, but may need to be adjusted to fit your system.
+The default batch size, 4, fits in 16GB GPU memory, but may need to be adjusted to fit your system.
+
 
 ### Training
-Run/modify `finetune_bart.sh` or `finetune_t5.sh`
-
-### Stanford CoreNLP Setup
-```
-ptb_tokenize () {
-    cat $1 | java edu.stanford.nlp.process.PTBTokenizer -ioFileList -preserveLines > $2
-}
-
-sudo apt install openjdk-8-jre-headless
-sudo apt-get install ant
-wget http://nlp.stanford.edu/software/stanford-corenlp-full-2018-10-05.zip
-unzip stanford-corenlp-full-2018-10-05.zip
-cd stanford-corenlp-full-2018-10-05
-export CLASSPATH=stanford-corenlp-3.9.2.jar:stanford-corenlp-3.9.2-models.jar
-```
-Then run `ptb_tokenize` on `test.target` and your generated hypotheses.
-### Rouge Setup
-Install `files2rouge` following the instructions at [here](https://github.com/pltrdy/files2rouge).
-I also needed to run `sudo apt-get install libxml-parser-perl`
-
-```python
-from files2rouge import files2rouge
-from files2rouge import settings
-files2rouge.run(<path_to_tokenized_hypo>,
-                <path_to_tokenized_target>,
-               saveto='rouge_output.txt')
-```
+Run/modify `finetune.sh`
diff --git a/examples/summarization/distillation.py b/examples/summarization/distillation.py
index 1d7e1f42..4996930a 100644
--- a/examples/summarization/distillation.py
+++ b/examples/summarization/distillation.py
@@ -21,7 +21,7 @@ try:
         any_requires_grad,
     )
     from .initialization_utils import init_student, copy_layers
-    from .utils import SummarizationDataset, pickle_load
+    from .utils import use_task_specific_params, SummarizationDataset, pickle_load
     from .finetune import main as ft_main
 except ImportError:
     from finetune import (
@@ -32,7 +32,7 @@ except ImportError:
     )
     from finetune import main as ft_main
     from initialization_utils import init_student, copy_layers
-    from utils import SummarizationDataset, pickle_load
+    from utils import use_task_specific_params, SummarizationDataset, pickle_load
 
 
 class SummarizationDistiller(SummarizationTrainer):
@@ -45,6 +45,7 @@ class SummarizationDistiller(SummarizationTrainer):
 
         super().__init__(hparams, model=student, config=student_cfg)
         self.teacher = teacher
+        use_task_specific_params(self.teacher, "summarization")
         freeze_params(self.teacher)
         assert len(self.model.model.decoder.layers) == len(d_layers_to_copy)
         self.sanity_check_gradients()
diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index d9f920c7..87bea47f 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -16,14 +16,14 @@ from torch import nn
 from torch.utils.data import DataLoader
 
 from lightning_base import BaseTransformer, add_generic_args, generic_train
-from transformers import AutoModelWithLMHead, get_linear_schedule_with_warmup
+from transformers import get_linear_schedule_with_warmup
 
 
 try:
-    from .utils import SummarizationDataset, lmap, flatten_list, pickle_save
+    from .utils import use_task_specific_params, SummarizationDataset, lmap, flatten_list, pickle_save
     from .callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 except ImportError:
-    from utils import SummarizationDataset, lmap, flatten_list, pickle_save
+    from utils import use_task_specific_params, SummarizationDataset, lmap, flatten_list, pickle_save
     from callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 
 logger = logging.getLogger(__name__)
@@ -105,13 +105,13 @@ def assert_not_all_frozen(model):
 
 
 class SummarizationTrainer(BaseTransformer):
-    mode = "language-modeling"
+    mode = "summarization"
     loss_names = ["loss"]
 
     def __init__(self, hparams, **kwargs):
         super().__init__(hparams, num_labels=None, mode=self.mode, **kwargs)
+        use_task_specific_params(self.model, "summarization")
         save_git_info(self.hparams.output_dir)
-        self.model: AutoModelWithLMHead
         self.metrics_save_path = Path(self.output_dir) / "metrics.pkl"
         self.hparams_save_path = Path(self.output_dir) / "hparams.pkl"
         self.step_count = 0
diff --git a/examples/summarization/finetune_bart.sh b/examples/summarization/finetune.sh
similarity index 71%
rename from examples/summarization/finetune_bart.sh
rename to examples/summarization/finetune.sh
index d3a86ba1..f4e6de20 100755
--- a/examples/summarization/finetune_bart.sh
+++ b/examples/summarization/finetune.sh
@@ -9,10 +9,12 @@ export PYTHONPATH="../":"${PYTHONPATH}"
 
 python -m pdb -c continue finetune.py \
 --data_dir=$CNN_DIR \
---model_name_or_path=facebook/bart-large \
+--model_name_or_path=facebook/bart-large \   # --model_name_or_path=t5-base for t5, and smaller batch_size
 --learning_rate=3e-5 \
---train_batch_size=4 \
---eval_batch_size=4 \
+--train_batch_size=2 \   # roughly 13GB
+--eval_batch_size=2 \
+--fp16 \
+--gpus 1 \
 --output_dir=$OUTPUT_DIR \
 --do_train \
 --do_predict \
diff --git a/examples/summarization/evaluate_cnn.py b/examples/summarization/run_eval.py
similarity index 69%
rename from examples/summarization/evaluate_cnn.py
rename to examples/summarization/run_eval.py
index f2b57f09..e2dcf16c 100644
--- a/examples/summarization/evaluate_cnn.py
+++ b/examples/summarization/run_eval.py
@@ -1,13 +1,18 @@
 import argparse
+import json
 from pathlib import Path
 
 import torch
-from rouge_score import rouge_scorer, scoring
 from tqdm import tqdm
 
 from transformers import AutoModelWithLMHead, AutoTokenizer
 
 
+try:
+    from .finetune import calculate_rouge, use_task_specific_params
+except ImportError:
+    from finetune import calculate_rouge, use_task_specific_params
+
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -26,9 +31,7 @@ def generate_summaries(
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     # update config with summarization specific params
-    task_specific_params = model.config.task_specific_params
-    if task_specific_params is not None:
-        model.config.update(task_specific_params.get("summarization", {}))
+    use_task_specific_params(model, "summarization")
 
     for batch in tqdm(list(chunks(examples, batch_size))):
         if "t5" in model_name:
@@ -44,27 +47,10 @@ def generate_summaries(
             fout.flush()
 
 
-def calculate_rouge(output_lns, reference_lns, score_path):
-    score_file = Path(score_path).open("w")
-    scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
-    aggregator = scoring.BootstrapAggregator()
-
-    for reference_ln, output_ln in zip(reference_lns, output_lns):
-        scores = scorer.score(reference_ln, output_ln)
-        aggregator.add_scores(scores)
-
-    result = aggregator.aggregate()
-    score_file.write(
-        "ROUGE_1: \n{} \n\n ROUGE_2: \n{} \n\n ROUGE_L: \n{} \n\n".format(
-            result["rouge1"], result["rouge2"], result["rougeL"]
-        )
-    )
-
-
 def run_generate():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "input_path", type=str, help="like cnn_dm/test.source or cnn_dm/test_articles_input.txt",
+        "input_path", type=str, help="like cnn_dm/test.source",
     )
     parser.add_argument(
         "output_path", type=str, help="where to save summaries",
@@ -73,11 +59,11 @@ def run_generate():
         "model_name",
         type=str,
         default="facebook/bart-large-cnn",
-        help="like bart-large-cnn,'t5-small', 't5-base', 't5-large', 't5-3b', 't5-11b",
+        help="like facebook/bart-large-cnn,'t5-small', 't5-base', 't5-large', 't5-3b', 't5-11b",
     )
     parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test_reference_summaries.txt")
     parser.add_argument(
-        "--score_path", type=str, required=False, help="where to save the rouge score",
+        "--score_path", type=str, required=False, help="where to save the rouge score in json format",
     )
     parser.add_argument(
         "--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.",
@@ -93,7 +79,9 @@ def run_generate():
         output_lns = [x.rstrip() for x in open(args.output_path).readlines()]
         reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()]
 
-        calculate_rouge(output_lns, reference_lns, args.score_path)
+        rouge: dict = calculate_rouge(output_lns, reference_lns)
+
+        json.dump(rouge, open("score_path", "w+"))
 
 
 if __name__ == "__main__":
diff --git a/examples/summarization/test_summarization_examples.py b/examples/summarization/test_summarization_examples.py
index 3234eb46..64b19f69 100644
--- a/examples/summarization/test_summarization_examples.py
+++ b/examples/summarization/test_summarization_examples.py
@@ -13,8 +13,8 @@ from torch.utils.data import DataLoader
 from transformers import BartTokenizer
 
 from .distillation import distill_main, evaluate_checkpoint
-from .evaluate_cnn import generate_summaries, run_generate
 from .finetune import main
+from .run_eval import generate_summaries, run_generate
 from .utils import SummarizationDataset, lmap, pickle_load
 
 
@@ -268,7 +268,7 @@ class TestBartExamples(unittest.TestCase):
         output_file_name = Path(tempfile.gettempdir()) / "utest_output_bart_sum.hypo"
         articles = [" New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County."]
         _dump_articles(tmp, articles)
-        testargs = ["evaluate_cnn.py", str(tmp), str(output_file_name), "sshleifer/bart-tiny-random"]
+        testargs = ["run_eval.py", str(tmp), str(output_file_name), "sshleifer/bart-tiny-random"]
         with patch.object(sys, "argv", testargs):
             run_generate()
             self.assertTrue(Path(output_file_name).exists())
diff --git a/examples/summarization/utils.py b/examples/summarization/utils.py
index 43923d21..e3254078 100644
--- a/examples/summarization/utils.py
+++ b/examples/summarization/utils.py
@@ -158,6 +158,13 @@ class SortishSampler(Sampler):
         return iter(sort_idx)
 
 
+def use_task_specific_params(model, task):
+    # update config with summarization specific params
+    task_specific_params = model.config.task_specific_params
+    if task_specific_params is not None:
+        model.config.update(task_specific_params.get(task, {}))
+
+
 def pickle_load(path):
     """pickle.load(path)"""
     with open(path, "rb") as f:
diff --git a/src/transformers/modeling_bart.py b/src/transformers/modeling_bart.py
index 09a52eae..01b92ed2 100644
--- a/src/transformers/modeling_bart.py
+++ b/src/transformers/modeling_bart.py
@@ -59,7 +59,7 @@ BART_GENERATION_EXAMPLE = r"""
     Examples::
 
         from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig
-        # see ``examples/summarization/bart/evaluate_cnn.py`` for a longer example
+        # see ``examples/summarization/bart/run_eval.py`` for a longer example
         model = BartForConditionalGeneration.from_pretrained('bart-large-cnn')
         tokenizer = BartTokenizer.from_pretrained('bart-large-cnn')
         ARTICLE_TO_SUMMARIZE = "My friends are cool but they eat too many carbs."

From 969c271bd7faeaa6b37e24902fba70f77e627f38 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 20:18:43 +0000
Subject: [PATCH 11/18] more honest docs

---
 examples/summarization/README.md   |  5 +++++
 examples/summarization/finetune.sh | 35 ++++++++++++++++--------------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index 721c42b1..b2dda578 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -27,3 +27,8 @@ The default batch size, 4, fits in 16GB GPU memory, but may need to be adjusted
 
 ### Training
 Run/modify `finetune.sh`
+
+Tips:
+- 1 epoch at batch size 1 for bart-large takes 24 hours, requires 13GB GPU RAM with fp16.
+- try --freeze_encoder or --freeze_embeds for faster training/larger batch size.
+- fp16 opt level O1 is best.
diff --git a/examples/summarization/finetune.sh b/examples/summarization/finetune.sh
index f4e6de20..709fe38b 100755
--- a/examples/summarization/finetune.sh
+++ b/examples/summarization/finetune.sh
@@ -1,4 +1,3 @@
-export OUTPUT_DIR_NAME=bart_cnn_finetune
 export OUTPUT_DIR=bart_cnn_finetune
 
 # Make output directory if it doesn't exist
@@ -7,19 +6,23 @@ mkdir -p $OUTPUT_DIR
 # Add parent directory to python path to access lightning_base.py
 export PYTHONPATH="../":"${PYTHONPATH}"
 
+
+# --model_name_or_path=t5-base for t5
+
 python -m pdb -c continue finetune.py \
---data_dir=$CNN_DIR \
---model_name_or_path=facebook/bart-large \   # --model_name_or_path=t5-base for t5, and smaller batch_size
---learning_rate=3e-5 \
---train_batch_size=2 \   # roughly 13GB
---eval_batch_size=2 \
---fp16 \
---gpus 1 \
---output_dir=$OUTPUT_DIR \
---do_train \
---do_predict \
---n_val 1000 \
---val_check_interval 0.1 \
---sortish_sampler \
---num_train_epochs 6 \
-$@
+    --data_dir=$CNN_DIR \
+    --model_name_or_path=facebook/bart-large \
+    --learning_rate=3e-5 \
+    --train_batch_size=1\
+    --eval_batch_size=1 \
+    --fp16 \
+    --gpus 1 \
+    --output_dir $OUTPUT_DIR \
+    --do_train \
+    --do_predict \
+    --n_val 1000 \
+    --val_check_interval 0.1 \
+    --sortish_sampler \
+    --num_train_epochs 6 \
+    --max_target_length=56 \
+    $@

From f179e7b37c32fd894d0d5d8f1798dbe7d88cf310 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 17:58:41 -0400
Subject: [PATCH 12/18] Allow wandb logger

---
 examples/summarization/finetune.py                | 15 +++++++++++++--
 .../summarization/test_summarization_examples.py  |  1 +
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 87bea47f..0266b239 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -11,6 +11,7 @@ import git
 import numpy as np
 import pytorch_lightning as pl
 import torch
+from pytorch_lightning.loggers import WandbLogger
 from rouge_score import rouge_scorer, scoring
 from torch import nn
 from torch.utils.data import DataLoader
@@ -323,7 +324,6 @@ class SummarizationTrainer(BaseTransformer):
             help="The maximum total input sequence length after tokenization. Sequences longer "
             "than this will be truncated, sequences shorter will be padded.",
         )
-
         parser.add_argument(
             "--data_dir",
             default=None,
@@ -337,25 +337,36 @@ class SummarizationTrainer(BaseTransformer):
         parser.add_argument(
             "--freeze_embeds", action="store_true",
         )
+        parser.add_argument("--sortish_sampler", action="store_true", default=False)
+        parser.add_argument("--logger", type="str", choices=["default", "wandb", "wandb_shared"], default="default")
         parser.add_argument("--n_train", type=int, default=-1, required=False)
         parser.add_argument("--n_val", type=int, default=500, required=False)
         parser.add_argument("--n_test", type=int, default=-1, required=False)
-        parser.add_argument("--sortish_sampler", action="store_true", default=False)
+
         return parser
 
 
+
+
 def main(args, model=None):
     Path(args.output_dir).mkdir(exist_ok=True)
     if len(os.listdir(args.output_dir)) > 3 and args.do_train:
         raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
     if model is None:
         model: BaseTransformer = SummarizationTrainer(args)
+    if args.logger == "default":
+        logger = True
+    elif args.logger == "wandb":
+        logger = WandbLogger()
+    elif args.logger == "wandb_shared":
+        logger = WandbLogger(name=args.output_dir, project="hf_summarization")
     trainer: pl.Trainer = generic_train(
         model,
         args,
         early_stopping_callback=True,
         logging_callback=Seq2SeqLoggingCallback(),
         checkpoint_callback=get_rouge2_checkpoint_callback(args.output_dir),
+        logger=logger,
     )
     if not args.do_predict:
         return model
diff --git a/examples/summarization/test_summarization_examples.py b/examples/summarization/test_summarization_examples.py
index 64b19f69..f81364dc 100644
--- a/examples/summarization/test_summarization_examples.py
+++ b/examples/summarization/test_summarization_examples.py
@@ -23,6 +23,7 @@ logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger()
 FP16_EVER = False
 CHEAP_ARGS = {
+    "logger": "default",
     "alpha_hid": 0,
     "freeze_embeds": True,
     "enc_only": False,

From c34b886cd141452fd25d5a6bb399bfdd39253fcd Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 18:11:22 -0400
Subject: [PATCH 13/18] Wandb logger

---
 examples/summarization/README.md   | 34 ++++++++++++++++++++++++++++--
 examples/summarization/finetune.py |  4 +---
 examples/summarization/finetune.sh |  7 +-----
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index b2dda578..ba748d76 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -2,8 +2,10 @@
 
 CNN/DailyMail data
 ```bash
+cd examples/summarization
 wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/cnn_dm.tgz
 tar -xzvf cnn_dm.tgz
+export CNN_DIR=${PWD}/cnn_dm
 ```
 
 this should make a directory called cnn_dm/ with files like `test.source`.
@@ -11,8 +13,10 @@ To use your own data, copy that files format. Each article to be summarized is o
 
 XSUM Data:
 ```bash
+cd examples/summarization
 wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/xsum.tar.gz
 tar -xzvf xsum.tar.gz
+export XSUM_DIR=${PWD}/xsum
 ```
 
 
@@ -28,7 +32,33 @@ The default batch size, 4, fits in 16GB GPU memory, but may need to be adjusted
 ### Training
 Run/modify `finetune.sh`
 
+The following command should work on a 16GB GPU:
+```bash
+export myhandle='sshleifer'
+./finetune.sh \
+    --data_dir $XSUM_DIR \
+    --train_batch_size=1 \
+    --eval_batch_size=1 \
+    --output_dir="$myhandle"_xsum_results \
+    --num_train_epochs 1
+```
+
 Tips:
 - 1 epoch at batch size 1 for bart-large takes 24 hours, requires 13GB GPU RAM with fp16.
-- try --freeze_encoder or --freeze_embeds for faster training/larger batch size.
-- fp16 opt level O1 is best.
+- try --freeze_encoder or --freeze_embeds for faster training/larger batch size.  (3hr/epoch with bs=8, see below)
+- fp16 opt level O1 is best
+
+### Shared Task
+Compare XSUM results with others by using `--logger wandb_shared`. This requires `wandb` registration.
+Here is an example command
+```bash
+export myhandle='sshleifer'
+./finetune.sh
+    --data_dir $XSUM_DIR \
+    --output_dir "$myhandle"xsum_frozen_embs \
+    --logger wandb_shared \
+    --train_batch_size 8 -eval_batch_size 16 --freeze_embeds --freeze_encoder \
+    --num_train_epochs 6
+```
+
+Results can be viewed [here](https://app.wandb.ai/sshleifer)
diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 0266b239..88e7d2a1 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -346,15 +346,13 @@ class SummarizationTrainer(BaseTransformer):
         return parser
 
 
-
-
 def main(args, model=None):
     Path(args.output_dir).mkdir(exist_ok=True)
     if len(os.listdir(args.output_dir)) > 3 and args.do_train:
         raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
     if model is None:
         model: BaseTransformer = SummarizationTrainer(args)
-    if args.logger == "default":
+    if args.logger == "default" or args.fast_dev_run:
         logger = True
     elif args.logger == "wandb":
         logger = WandbLogger()
diff --git a/examples/summarization/finetune.sh b/examples/summarization/finetune.sh
index 709fe38b..ead55c58 100755
--- a/examples/summarization/finetune.sh
+++ b/examples/summarization/finetune.sh
@@ -9,20 +9,15 @@ export PYTHONPATH="../":"${PYTHONPATH}"
 
 # --model_name_or_path=t5-base for t5
 
-python -m pdb -c continue finetune.py \
-    --data_dir=$CNN_DIR \
+python finetune.py \
     --model_name_or_path=facebook/bart-large \
     --learning_rate=3e-5 \
-    --train_batch_size=1\
-    --eval_batch_size=1 \
     --fp16 \
     --gpus 1 \
-    --output_dir $OUTPUT_DIR \
     --do_train \
     --do_predict \
     --n_val 1000 \
     --val_check_interval 0.1 \
     --sortish_sampler \
-    --num_train_epochs 6 \
     --max_target_length=56 \
     $@

From c9597fef425b9a15bbd0fdee093b0a8301801f2a Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 18:15:28 -0400
Subject: [PATCH 14/18] docs

---
 examples/summarization/README.md   | 12 ++++++------
 examples/summarization/finetune.py |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index ba748d76..fac0b1e2 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -34,12 +34,12 @@ Run/modify `finetune.sh`
 
 The following command should work on a 16GB GPU:
 ```bash
-export myhandle='sshleifer'
+export me=`git config user.name`
 ./finetune.sh \
     --data_dir $XSUM_DIR \
     --train_batch_size=1 \
     --eval_batch_size=1 \
-    --output_dir="$myhandle"_xsum_results \
+    --output_dir="$me"_xsum_results \
     --num_train_epochs 1
 ```
 
@@ -52,12 +52,12 @@ Tips:
 Compare XSUM results with others by using `--logger wandb_shared`. This requires `wandb` registration.
 Here is an example command
 ```bash
-export myhandle='sshleifer'
-./finetune.sh
+export me=`git config user.name`
+./finetune.sh \
     --data_dir $XSUM_DIR \
-    --output_dir "$myhandle"xsum_frozen_embs \
+    --output_dir "$me"_xsum_frozen_embs \
     --logger wandb_shared \
-    --train_batch_size 8 -eval_batch_size 16 --freeze_embeds --freeze_encoder \
+    --train_batch_size 16 -eval_batch_size 32 --freeze_embeds --freeze_encoder \
     --num_train_epochs 6
 ```
 
diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 88e7d2a1..3a2aa298 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -338,7 +338,7 @@ class SummarizationTrainer(BaseTransformer):
             "--freeze_embeds", action="store_true",
         )
         parser.add_argument("--sortish_sampler", action="store_true", default=False)
-        parser.add_argument("--logger", type="str", choices=["default", "wandb", "wandb_shared"], default="default")
+        parser.add_argument("--logger", type=str, choices=["default", "wandb", "wandb_shared"], default="default")
         parser.add_argument("--n_train", type=int, default=-1, required=False)
         parser.add_argument("--n_val", type=int, default=500, required=False)
         parser.add_argument("--n_test", type=int, default=-1, required=False)

From 9e954298434c954b72920d02e5bb8c8a58d9c0a9 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 18:20:02 -0400
Subject: [PATCH 15/18] docs

---
 examples/summarization/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index fac0b1e2..bbe9e0d6 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -57,7 +57,7 @@ export me=`git config user.name`
     --data_dir $XSUM_DIR \
     --output_dir "$me"_xsum_frozen_embs \
     --logger wandb_shared \
-    --train_batch_size 16 -eval_batch_size 32 --freeze_embeds --freeze_encoder \
+    --train_batch_size 16 --eval_batch_size 16 --freeze_embeds --freeze_encoder \
     --num_train_epochs 6
 ```
 

From 99de2c3c3ce40f0a9a5543cbb0ccc49dcb3bf53f Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 19:01:15 -0400
Subject: [PATCH 16/18] pass through logger

---
 examples/lightning_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/lightning_base.py b/examples/lightning_base.py
index d4f241c0..e221747e 100644
--- a/examples/lightning_base.py
+++ b/examples/lightning_base.py
@@ -309,7 +309,7 @@ def generic_train(
         train_params["distributed_backend"] = "ddp"
 
     trainer = pl.Trainer(
-        logger=True,
+        logger=logger,
         accumulate_grad_batches=args.gradient_accumulation_steps,
         gpus=args.gpus,
         max_epochs=args.num_train_epochs,

From 6bdfb14664a35c6c57a37007f1f5bdc0c58b04aa Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 19:07:15 -0400
Subject: [PATCH 17/18] boom boom

---
 examples/summarization/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/summarization/README.md b/examples/summarization/README.md
index bbe9e0d6..3e3d798c 100644
--- a/examples/summarization/README.md
+++ b/examples/summarization/README.md
@@ -61,4 +61,4 @@ export me=`git config user.name`
     --num_train_epochs 6
 ```
 
-Results can be viewed [here](https://app.wandb.ai/sshleifer)
+Results can be viewed [here](https://app.wandb.ai/sshleifer/hf_summarization/table?workspace=user-)

From 5283878c9fe5607eac829c232d8ab17c9f98ed16 Mon Sep 17 00:00:00 2001
From: sshleifer <sshleifer@gmail.com>
Date: Sun, 14 Jun 2020 19:17:31 -0400
Subject: [PATCH 18/18] Move stuff to utils

---
 examples/summarization/distillation.py |  22 ++---
 examples/summarization/finetune.py     | 107 ++++++-------------------
 examples/summarization/utils.py        |  69 +++++++++++++++-
 3 files changed, 105 insertions(+), 93 deletions(-)

diff --git a/examples/summarization/distillation.py b/examples/summarization/distillation.py
index 4996930a..bc7c616a 100644
--- a/examples/summarization/distillation.py
+++ b/examples/summarization/distillation.py
@@ -14,25 +14,29 @@ from transformers import AdamW, BartConfig, BartForConditionalGeneration, T5Conf
 
 
 try:
-    from .finetune import (
-        SummarizationTrainer,
+    from .finetune import SummarizationTrainer
+    from .initialization_utils import init_student, copy_layers
+    from .utils import (
+        use_task_specific_params,
+        SummarizationDataset,
+        pickle_load,
         freeze_params,
         assert_all_frozen,
         any_requires_grad,
     )
-    from .initialization_utils import init_student, copy_layers
-    from .utils import use_task_specific_params, SummarizationDataset, pickle_load
     from .finetune import main as ft_main
 except ImportError:
-    from finetune import (
-        SummarizationTrainer,
+    from finetune import SummarizationTrainer
+    from finetune import main as ft_main
+    from initialization_utils import init_student, copy_layers
+    from utils import (
+        use_task_specific_params,
+        SummarizationDataset,
+        pickle_load,
         freeze_params,
         assert_all_frozen,
         any_requires_grad,
     )
-    from finetune import main as ft_main
-    from initialization_utils import init_student, copy_layers
-    from utils import use_task_specific_params, SummarizationDataset, pickle_load
 
 
 class SummarizationDistiller(SummarizationTrainer):
diff --git a/examples/summarization/finetune.py b/examples/summarization/finetune.py
index 3a2aa298..1e9ccda6 100644
--- a/examples/summarization/finetune.py
+++ b/examples/summarization/finetune.py
@@ -1,19 +1,15 @@
 import argparse
 import glob
-import json
 import logging
 import os
 import time
 from pathlib import Path
-from typing import Dict, Iterable, List, Tuple
+from typing import Dict, List, Tuple
 
-import git
 import numpy as np
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.loggers import WandbLogger
-from rouge_score import rouge_scorer, scoring
-from torch import nn
 from torch.utils.data import DataLoader
 
 from lightning_base import BaseTransformer, add_generic_args, generic_train
@@ -21,89 +17,36 @@ from transformers import get_linear_schedule_with_warmup
 
 
 try:
-    from .utils import use_task_specific_params, SummarizationDataset, lmap, flatten_list, pickle_save
+    from .utils import (
+        use_task_specific_params,
+        SummarizationDataset,
+        lmap,
+        flatten_list,
+        pickle_save,
+        save_git_info,
+        freeze_params,
+        calculate_rouge,
+        get_git_info,
+        ROUGE_KEYS,
+    )
     from .callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 except ImportError:
-    from utils import use_task_specific_params, SummarizationDataset, lmap, flatten_list, pickle_save
+    from .utils import (
+        use_task_specific_params,
+        SummarizationDataset,
+        lmap,
+        flatten_list,
+        pickle_save,
+        save_git_info,
+        freeze_params,
+        calculate_rouge,
+        get_git_info,
+        ROUGE_KEYS,
+    )
     from callbacks import Seq2SeqLoggingCallback, get_rouge2_checkpoint_callback
 
 logger = logging.getLogger(__name__)
 
-ROUGE_KEYS = ["rouge1", "rouge2", "rougeL"]
-
-
-def save_git_info(folder_path: str):
-    """
-    Log commit info.
-    """
-    repo_infos = get_git_info()
-
-    with open(os.path.join(folder_path, "git_log.json"), "w") as f:
-        json.dump(repo_infos, f, indent=4)
-
-
-def get_git_info():
-    repo = git.Repo(search_parent_directories=True)
-    repo_infos = {
-        "repo_id": str(repo),
-        "repo_sha": str(repo.head.object.hexsha),
-        "repo_branch": str(repo.active_branch),
-    }
-    return repo_infos
-
-
-def calculate_rouge(output_lns: List[str], reference_lns: List[str]) -> Dict:
-    scorer = rouge_scorer.RougeScorer(ROUGE_KEYS, use_stemmer=True)
-    aggregator = scoring.BootstrapAggregator()
-
-    for reference_ln, output_ln in zip(reference_lns, output_lns):
-        scores = scorer.score(reference_ln, output_ln)
-        aggregator.add_scores(scores)
-
-    result = aggregator.aggregate()
-    return {k: v.mid.fmeasure for k, v in result.items()}
-
-
-def dictify(rouge_obj) -> List:
-    records = []
-    for k, rouge_measurement in rouge_obj.items():
-        if k == "rouge1":
-            continue
-        for k1 in ["low", "mid", "high"]:
-            if k1 != "mid":
-                continue
-            v1 = getattr(rouge_measurement, k1)
-            for k2 in ["precision", "recall", "fmeasure"]:
-                records.append([k, k1, k2, getattr(v1, k2)])
-
-    return records
-
-
-def freeze_params(model: nn.Module):
-    for par in model.parameters():
-        par.requires_grad = False
-
-
-def grad_status(model: nn.Module) -> Iterable:
-    return (par.requires_grad for par in model.parameters())
-
-
-def any_requires_grad(model: nn.Module) -> bool:
-    return any(grad_status(model))
-
-
-def assert_all_frozen(model):
-    model_grads: List[bool] = list(grad_status(model))
-    n_require_grad = sum(lmap(int, model_grads))
-    npars = len(model_grads)
-    assert not any(model_grads), f"{n_require_grad/npars:.1%} of {npars} weights require grad"
-
-
-def assert_not_all_frozen(model):
-    model_grads: List[bool] = list(grad_status(model))
-    npars = len(model_grads)
-    assert any(model_grads), f"none of {npars} weights require grad"
-
 
 class SummarizationTrainer(BaseTransformer):
     mode = "summarization"
diff --git a/examples/summarization/utils.py b/examples/summarization/utils.py
index e3254078..f8a2c910 100644
--- a/examples/summarization/utils.py
+++ b/examples/summarization/utils.py
@@ -1,11 +1,15 @@
 import itertools
+import json
 import os
 import pickle
 from pathlib import Path
-from typing import List
+from typing import Dict, Iterable, List
 
+import git
 import numpy as np
 import torch
+from rouge_score import rouge_scorer, scoring
+from torch import nn
 from torch.utils.data import Dataset, Sampler
 from tqdm import tqdm
 
@@ -55,7 +59,7 @@ def lmap(f, x):
     return list(map(f, x))
 
 
-T5_PREFIX = "summarize: "
+T5_PREFIX = "summarize: "  # HACK, fixme
 
 
 class SummarizationDataset(Dataset):
@@ -179,3 +183,64 @@ def pickle_save(obj, path):
 
 def flatten_list(summary_ids: List[List]):
     return [x for x in itertools.chain.from_iterable(summary_ids)]
+
+
+def save_git_info(folder_path: str):
+    """
+    Log commit info.
+    """
+    repo_infos = get_git_info()
+
+    with open(os.path.join(folder_path, "git_log.json"), "w") as f:
+        json.dump(repo_infos, f, indent=4)
+
+
+def get_git_info():
+    repo = git.Repo(search_parent_directories=True)
+    repo_infos = {
+        "repo_id": str(repo),
+        "repo_sha": str(repo.head.object.hexsha),
+        "repo_branch": str(repo.active_branch),
+    }
+    return repo_infos
+
+
+ROUGE_KEYS = ["rouge1", "rouge2", "rougeL"]
+
+
+def calculate_rouge(output_lns: List[str], reference_lns: List[str]) -> Dict:
+    scorer = rouge_scorer.RougeScorer(ROUGE_KEYS, use_stemmer=True)
+    aggregator = scoring.BootstrapAggregator()
+
+    for reference_ln, output_ln in zip(reference_lns, output_lns):
+        scores = scorer.score(reference_ln, output_ln)
+        aggregator.add_scores(scores)
+
+    result = aggregator.aggregate()
+    return {k: v.mid.fmeasure for k, v in result.items()}
+
+
+def freeze_params(model: nn.Module):
+    for par in model.parameters():
+        par.requires_grad = False
+
+
+def grad_status(model: nn.Module) -> Iterable:
+    return (par.requires_grad for par in model.parameters())
+
+
+def any_requires_grad(model: nn.Module) -> bool:
+    return any(grad_status(model))
+
+
+def assert_all_frozen(model):
+    model_grads: List[bool] = list(grad_status(model))
+    n_require_grad = sum(lmap(int, model_grads))
+    npars = len(model_grads)
+    assert not any(model_grads), f"{n_require_grad/npars:.1%} of {npars} weights require grad"
+
+
+def assert_not_all_frozen(model):
+    model_grads: List[bool] = list(grad_status(model))
+    npars = len(model_grads)
+    assert any(model_grads), f"none of {npars} weights require grad"