examples/seq2seq supports translation (huggingface#5202)

examples/seq2seq/README.md

@@ -0,0 +1,169 @@
+This directory contains examples for finetuning and evaluating transformers on summarization and translation tasks.
+Summarization support is more mature than translation support.
+Please tag @sshleifer with any issues/unexpected behaviors, or send a PR!
+For `bertabs` instructions, see `bertabs/README.md`. 
+
+### Data
+
+CNN/DailyMail data
+```bash
+cd examples/seq2seq
+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`.
+To use your own data, copy that files format. Each article to be summarized is on its own line.
+
+XSUM Data:
+```bash
+cd examples/seq2seq
+wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/xsum.tar.gz
+tar -xzvf xsum.tar.gz
+export XSUM_DIR=${PWD}/xsum
+```
+
+
+WMT16 English-Romanian Translation Data:
+```bash
+cd examples/seq2seq
+wget https://s3.amazonaws.com/datasets.huggingface.co/translation/wmt_en_ro.tar.gz
+tar -xzvf wmt_en_ro.tar.gz
+export ENRO_DIR=${PWD}/wmt_en_ro
+```
+
+If you are using your own data, it must be formatted as one directory with 6 files: train.source, train.target, val.source, val.target, test.source, test.target.  
+The `.source` files are the input, the `.target` files are the desired output.
+
+### Evaluation
+
+To create summaries for each article in dataset, run:
+```bash
+python run_eval.py <path_to_test.source> test_generations.txt <model-name>  --score_path rouge_scores.txt
+```
+The default batch size, 4, fits in 16GB GPU memory, but may need to be adjusted to fit your system.
+
+
+### Summarization Finetuning
+Run/modify `finetune.sh`
+
+The following command should work on a 16GB GPU:
+```bash
+./finetune.sh \
+    --data_dir $XSUM_DIR \
+    --train_batch_size=1 \
+    --eval_batch_size=1 \
+    --output_dir=xsum_results \
+    --num_train_epochs 1 \
+    --model_name_or_path facebook/bart-large
+```
+
+*Note*: The following tips mostly apply to summarization finetuning.
+
+Tips:
+- 1 epoch at batch size 1 for bart-large takes 24 hours and requires 13GB GPU RAM with fp16 on an NVIDIA-V100. 
+- try `bart-base`, `--freeze_encoder` or `--freeze_embeds` for faster training/larger batch size.  (3hr/epoch with bs=8, see the "xsum_shared_task" command below)
+- `fp16_opt_level=O1` (the default works best).
+- If you are finetuning on your own dataset, start from `distilbart-cnn-12-6` if you want long summaries and `distilbart-xsum-12-6` if you want short summaries.
+(It rarely makes sense to start from `bart-large` unless you are a researching finetuning methods).
+- In addition to the pytorch-lightning .ckpt checkpoint, a transformers checkpoint will be saved.
+Load it with `BartForConditionalGeneration.from_pretrained(f'{output_dir}/best_tfmr)`.
+- At the moment, `--do_predict` does not work in a multi-gpu setting. You need to use `evaluate_checkpoint` or the `run_eval.py` code.
+- If you want to run experiments on improving the summarization finetuning process, try the XSUM Shared Task (below). It's faster to train than CNNDM because the summaries are shorter.
+- For CNN/DailyMail, the default `val_max_target_length` and `test_max_target_length` will truncate the ground truth labels, resulting in slightly higher rouge scores. To get accurate rouge scores, you should rerun calculate_rouge on the `{output_dir}/test_generations.txt` file saved by `trainer.test()`
+- `--max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 ` is a reasonable setting for XSUM.
+- `wandb` can be used by specifying `--logger wandb_shared` or `--logger wandb`. It is useful for reproducibility. 
+- This warning can be safely ignored: 
+    > "Some weights of BartForConditionalGeneration were not initialized from the model checkpoint at facebook/bart-large-xsum and are newly initialized: ['final_logits_bias']"
+- Both finetuning and eval are 30% faster with `--fp16`. For that you need to [install apex](https://github.com/NVIDIA/apex#quick-start).
+
+#### Finetuning Outputs 
+As you train, `output_dir` will be filled with files, that look kind of like this (comments are mine). 
+Some of them are metrics, some of them are checkpoints, some of them are metadata. Here is a quick tour:
+
+```bash
+output_dir
+├── best_tfmr  # this is a huggingface checkpoint generated by save_pretrained. It is the same model as the PL .ckpt file below
+│   ├── config.json
+│   ├── merges.txt
+│   ├── pytorch_model.bin
+│   ├── special_tokens_map.json
+│   ├── tokenizer_config.json
+│   └── vocab.json
+├── git_log.json   # repo, branch, and commit hash
+├── val_avg_rouge2=0.1984-step_count=11.ckpt  # this is a pytorch lightning checkpoint associated with the best val score.
+├── metrics.json  # new validation metrics will continually be appended to this
+├── student  # this is a huggingface checkpoint generated by SummarizationDistiller. It is the student before it gets finetuned.
+│   ├── config.json
+│   └── pytorch_model.bin
+├── test_generations.txt 
+# ^^ are the summaries or translations produced by your best checkpoint on the test data. Populated when training is done 
+├── test_results.txt  # a convenience file with the test set metrics. This data is also in metrics.json['test']
+├── hparams.pkl  # the command line args passed after some light preprocessing. Should be saved fairly quickly.
+```
+After training, you can recover the best checkpoint by running
+```python
+from transformers import AutoModelForSeq2SeqLM
+model = AutoModelForSeq2SeqLM.from_pretrained(f'{output_dir}/best_tfmr')
+```
+
+
+### XSUM Shared Task
+Compare XSUM results with others by using `--logger wandb_shared`. This requires `wandb` registration.
+
+Here is an example command, but you can do whatever you want. Hopefully this will make debugging and collaboration easier!
+```bash
+./finetune.sh \
+    --data_dir $XSUM_DIR \
+    --output_dir xsum_frozen_embs \
+    --model_name_or_path facebook/bart-large \
+    --logger wandb_shared \
+    --train_batch_size 16 --eval_batch_size 16 --freeze_embeds --freeze_encoder \
+    --num_train_epochs 6 \
+    --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100
+```
+
+Results can be viewed [here](https://app.wandb.ai/sshleifer/hf_xsum?workspace=user-)
+
+
+### Distilbart
+
+#### No Teacher Distillation
+To run the simpler distilbart-cnn style distillation all you need is data, a GPU, and a properly initialized student.
+You don't even need `distillation.py`.
+
+Some [un-finetuned students](https://huggingface.co/models?search=sshleifer%2Fstudent) are available for replication purposes.
+They are initialized by copying layers from the associated `bart-large-{cnn|xsum}` teacher using `--init_strategy alternate`. (You can read about that in `initialization_utils.py`)
+The command that produced `sshleifer/distilbart-cnn-12-6` is
+```bash
+./train_distilbart_cnn.sh
+```  
+runtime: 6H on NVIDIA RTX 24GB GPU
+
+*Note*: You can get the same simple distillation logic by using `./run_distiller.sh --no_teacher` followed by identical arguments as the ones in `train_distilbart_cnn.sh`.
+If you are using `wandb` and comparing the two distillation methods, using this entry point will make your logs consistent,
+because you will have the same hyperparameters logged in every run.
+
+#### With a teacher
+*Note* only BART variants are supported
+
+In this method, we use try to enforce that the student and teacher produce similar encoder_outputs, logits, and hidden_states using `BartSummarizationDistiller`.
+This is how `sshleifer/distilbart-xsum*` checkpoints were produced.
+
+The command that produced `sshleifer/distilbart-xsum-12-6` is:
+
+```bash
+./train_distilbart_xsum.sh  
+```
+
+runtime: 13H on V-100 16GB GPU. 
+
+### Contributing
+- follow the standard contributing guidelines and code of conduct.
+- add tests to `test_seq2seq_examples.py`
+- To run only the seq2seq tests, you must be in the root of the repository and run:
+```bash
+pytest examples/seq2seq/  
+```

examples/summarization/bertabs/README.md → examples/seq2seq/bertabs/README.md

@@ -12,7 +12,7 @@ The model is loaded with the pre-trained weights for the abstractive summarizati
 git clone https://github.com/huggingface/transformers && cd transformers
 pip install .
 pip install nltk py-rouge
-cd examples/summarization
+cd examples/seq2seq/bertabs
 ```
 
 ## Reproduce the authors'  ROUGE score

examples/summarization/callbacks.py → examples/seq2seq/callbacks.py

@@ -32,9 +32,12 @@ def _write_logs(
             results_file = od / "test_results.txt"
             generations_file = od / "test_generations.txt"
         else:
-            results_file = od / f"{type_path}_results_{trainer.global_step:05d}.txt"
-            generations_file = od / f"{type_path}_generations_{trainer.global_step:05d}.txt"
-
+            # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
+            # If people want this it will be easy enough to add back.
+            results_file = od / f"{type_path}_results/{trainer.global_step:05d}.txt"
+            generations_file = od / f"{type_path}_generations/{trainer.global_step:05d}.txt"
+            results_file.parent.mkdir(exist_ok=True)
+            generations_file.parent.mkdir(exist_ok=True)
         with open(results_file, "a+") as writer:
             for key in sorted(metrics):
                 if key in ["log", "progress_bar", "preds"]:
@@ -63,20 +66,25 @@ def on_train_start(self, trainer, pl_module):
         # mp stands for million parameters
         trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6})
 
-    @rank_zero_only
-    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
-        return self._write_logs(trainer, pl_module, "val")
-
     @rank_zero_only
     def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
         return self._write_logs(trainer, pl_module, "test")
 
 
-def get_rouge2_checkpoint_callback(output_dir):
+def get_checkpoint_callback(output_dir, metric):
     """Saves the best model by validation ROUGE2 score."""
+    if metric == "rouge2":
+        exp = "{val_avg_rouge2:.4f}-{step_count}"
+    elif metric == "bleu":
+        exp = "{val_avg_bleu:.4f}-{step_count}"
+    else:
+        raise NotImplementedError(
+            f"seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this function."
+        )
+
     checkpoint_callback = ModelCheckpoint(
-        filepath=os.path.join(output_dir, "{val_avg_rouge2:.4f}-{step_count}"),
-        monitor="val_rouge",
+        filepath=os.path.join(output_dir, exp),
+        monitor=f"val_{metric}",
         mode="max",
         save_top_k=1,
         period=0,  # maybe save a checkpoint every time val is run, not just end of epoch.

examples/summarization/distillation.py → examples/seq2seq/distillation.py

@@ -39,13 +39,12 @@
     )
 
 
-class SummarizationDistiller(SummarizationModule):
+class BartSummarizationDistiller(SummarizationModule):
     loss_names = ["loss", "ce_loss", "mlm_loss", "enc_mse_loss", "hid_loss_enc", "hid_loss_dec"]
 
     def __init__(self, hparams):
         assert Path(hparams.data_dir).exists()
-
-        d_layers_to_copy, student, student_cfg, teacher = self.pre_init(hparams)
+        student, student_cfg, teacher = self.pre_init(hparams)
 
         super().__init__(hparams, model=student, config=student_cfg)
         self.teacher = teacher
@@ -73,12 +72,15 @@ def sanity_check_gradients(self):
             del self.teacher.model.encoder
 
     def pre_init(self, hparams):
-        # Dump empty student model at a path, then call from_pretrained on it
+        self.output_dir = Path(hparams.output_dir)
+        self.output_dir.mkdir(exist_ok=True)
         teacher = BartForConditionalGeneration.from_pretrained(hparams.teacher).eval()
         student_updates = {
             "decoder_layers": hparams.student_decoder_layers,
             "encoder_layers": hparams.student_encoder_layers,
         }
+        if hparams.length_penalty != -1:
+            student_updates["length_penalty"] = hparams.length_penalty
         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.d_layer_to_copy = d_layers_to_copy
@@ -89,9 +91,13 @@ def pre_init(self, hparams):
         student_cfg = BartConfig(**kw)
         student = BartForConditionalGeneration(student_cfg)
         student, _ = init_student(student, teacher)
+        save_dir = self.output_dir.joinpath("student")
+        save_dir.mkdir(exist_ok=True)
+
         self.copy_to_student(d_layers_to_copy, e_layers_to_copy, hparams, student, teacher)
-        Path(hparams.output_dir).mkdir(exist_ok=True)
-        return d_layers_to_copy, student, student_cfg, teacher
+        student.save_pretrained(save_dir)
+        hparams.model_name_or_path = str(save_dir)
+        return student, student_cfg, teacher
 
     def copy_to_student(self, d_layers_to_copy, e_layers_to_copy, hparams, student, teacher):
         if teacher.config.model_type == "t5":
@@ -154,7 +160,6 @@ def calc_ce_loss(self, mask, s_logits, t_logits):
 
     def configure_optimizers(self):
         "Prepare optimizer and schedule (linear warmup and decay)"
-
         model = self.model
         no_decay = ["bias", "LayerNorm.weight"]
         optimizer_grouped_parameters = [
@@ -180,18 +185,11 @@ def add_model_specific_args(parser, root_dir):
         # parser.add_argument("--alpha_cos", default=0.0, type=float)
         parser.add_argument("--alpha_encoder_loss", default=0.0, type=float)
         parser.add_argument("--alpha_hid", default=0.0, type=float, required=False)
-        parser.add_argument(
-            "--student_decoder_layers", default=12, type=int, required=False,
-        )
-        parser.add_argument(
-            "--student_encoder_layers", default=12, type=int, required=False,
-        )
-        parser.add_argument(
-            "--no_teacher", action="store_true", default=False,
-        )
-        parser.add_argument(  # TODO: remove
-            "--enc_only", action="store_true", default=False,
-        )
+        parser.add_argument("--student_decoder_layers", default=12, type=int, required=False)
+        parser.add_argument("--student_encoder_layers", default=12, type=int, required=False)
+        parser.add_argument("--no_teacher", action="store_true", default=False)
+        parser.add_argument("--length_penalty", type=float, default=-1)
+
         return parser
 
     def _step(self, batch):
@@ -269,12 +267,14 @@ def calc_hidden_loss(self, attention_mask, hidden_states, hidden_states_T, match
         return sum(hidden_losses)
 
 
-class T5SummarizationDistiller(SummarizationDistiller):
+class T5SummarizationDistiller(BartSummarizationDistiller):
     def pre_init(self, hparams):
         raise NotImplementedError("T5 Distillation does not work yet")
+        self.output_dir = Path(hparams.output_dir)
+        self.output_dir.mkdir(exist_ok=True)
         teacher = T5ForConditionalGeneration.from_pretrained(hparams.teacher)
         n_layer = hparams.student_decoder_layers
-        assert n_layer == hparams.student_encoder_layers  # TODO(SS): relax this
+        assert n_layer == hparams.student_encoder_layers  # TODO(SS): relax this constraint so that we can do 12-6.
         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}
@@ -291,8 +291,13 @@ def pre_init(self, hparams):
         Path(hparams.output_dir).mkdir(exist_ok=True)
         task_specific_params = student.config.task_specific_params
         if task_specific_params is not None:
-            student.config.update(task_specific_params.get("summarization", {}))
-        return d_layers_to_copy, student, student_cfg, teacher
+            student.config.update(task_specific_params.get("summarization", {}))  # TODO: dont hardcode
+        save_dir = self.output_dir.joinpath("student")
+        save_dir.mkdir(exist_ok=True)
+
+        student.save_pretrained(save_dir)
+        hparams.model_name_or_path = str(save_dir)
+        return student, student_cfg, teacher
 
     def freeze_embeds(self):
         freeze_params(self.model.shared)
@@ -386,7 +391,7 @@ def create_module(args):
     elif args.enc_only:
         raise ValueError("Deleted that")
     else:
-        module_cls = SummarizationDistiller
+        module_cls = BartSummarizationDistiller
     args.setup_cls: str = module_cls.__name__
     model = module_cls(args)
     return model
@@ -418,18 +423,18 @@ def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
 def get_layers_to_copy(n_to_get, tot):
     all_layers = list(range(tot))
     if tot == 12:  # Alternating for special cases
-        layers_to_copy = {  # maps # layers in student -> which teacher layers to copy
-            6: [0, 2, 4, 7, 9, 11],
-            1: [11],
+        layers_to_copy = {  # maps  num layers in student -> which teacher layers to copy
+            1: [0],
+            2: [0, 6],
             3: [0, 6, 11],
-            2: [0, 11],
             4: [0, 4, 8, 11],
+            6: [0, 2, 4, 7, 9, 11],
             9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
             12: all_layers,
         }
         return layers_to_copy[n_to_get]
     else:
-        return all_layers[:n_to_get]
+        return all_layers[:n_to_get]  # TODO: better version on theseus-bart branch
 
 
 def distill_main(args):
@@ -443,7 +448,7 @@ def distill_main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser = SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
+    parser = BartSummarizationDistiller.add_model_specific_args(parser, os.getcwd())
     args = parser.parse_args()
 
     distill_main(args)