[examples] SummarizationDataset cleanup (#3451)

examples/summarization/bart/run_bart_sum.py

@@ -19,30 +19,20 @@ class BartSystem(BaseTransformer):
     mode = "language-modeling"
 
     def __init__(self, hparams):
-        super(BartSystem, self).__init__(hparams, num_labels=None, mode=self.mode)
+        super().__init__(hparams, num_labels=None, mode=self.mode, output_past=False)
 
-    def forward(
-        self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, lm_labels=None
-    ):
+    def forward(self, input_ids, attention_mask=None, decoder_input_ids=None, lm_labels=None):
         return self.model(
-            input_ids,
-            attention_mask=attention_mask,
-            decoder_input_ids=decoder_input_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            lm_labels=lm_labels,
+            input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, lm_labels=lm_labels,
         )
 
     def _step(self, batch):
-        y = batch["target_ids"]
+        pad_token_id = self.tokenizer.pad_token_id
+        source_ids, source_mask, y = batch["source_ids"], batch["source_mask"], batch["target_ids"]
         y_ids = y[:, :-1].contiguous()
         lm_labels = y[:, 1:].clone()
-        lm_labels[y[:, 1:] == self.tokenizer.pad_token_id] = -100
-        outputs = self(
-            input_ids=batch["source_ids"],
-            attention_mask=batch["source_mask"],
-            decoder_input_ids=y_ids,
-            lm_labels=lm_labels,
-        )
+        lm_labels[y[:, 1:] == pad_token_id] = -100
+        outputs = self(source_ids, attention_mask=source_mask, decoder_input_ids=y_ids, lm_labels=lm_labels,)
 
         loss = outputs[0]
 
@@ -64,9 +54,13 @@ class BartSystem(BaseTransformer):
         return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
 
     def test_step(self, batch, batch_idx):
+        # NOTE: this generation will not use the cache.
+        pad_token_id = self.tokenizer.pad_token_id
+        source_ids, source_mask, y = SummarizationDataset.trim_seq2seq_batch(batch, pad_token_id)
+        # NOTE: these kwargs get more speed and lower quality summaries than those in evaluate_cnn.py.
         generated_ids = self.model.generate(
-            batch["source_ids"],
-            attention_mask=batch["source_mask"],
+            source_ids,
+            source_mask,
             num_beams=1,
             max_length=80,
             repetition_penalty=2.5,
@@ -77,10 +71,7 @@ class BartSystem(BaseTransformer):
             self.tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True)
             for g in generated_ids
         ]
-        target = [
-            self.tokenizer.decode(t, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-            for t in batch["target_ids"]
-        ]
+        target = [self.tokenizer.decode(t, skip_special_tokens=True, clean_up_tokenization_spaces=True) for t in y]
         loss = self._step(batch)
 
         return {"val_loss": loss, "preds": preds, "target": target}
@@ -101,11 +92,21 @@ class BartSystem(BaseTransformer):
 
         return self.test_end(outputs)
 
-    def train_dataloader(self):
-        train_dataset = SummarizationDataset(
-            self.tokenizer, data_dir=self.hparams.data_dir, type_path="train", block_size=self.hparams.max_seq_length
+    @property
+    def dataset_kwargs(self):
+        return dict(
+            data_dir=self.hparams.data_dir,
+            max_source_length=self.hparams.max_source_length,
+            max_target_length=self.hparams.max_target_length,
         )
-        dataloader = DataLoader(train_dataset, batch_size=self.hparams.train_batch_size)
+
+    def get_dataloader(self, type_path: str, batch_size: int) -> DataLoader:
+        dataset = SummarizationDataset(self.tokenizer, type_path=type_path, **self.dataset_kwargs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=dataset.collate_fn)
+        return dataloader
+
+    def train_dataloader(self) -> DataLoader:
+        dataloader = self.get_dataloader("train", batch_size=self.hparams.train_batch_size)
         t_total = (
             (len(dataloader.dataset) // (self.hparams.train_batch_size * max(1, self.hparams.n_gpu)))
             // self.hparams.gradient_accumulation_steps
@@ -117,29 +118,30 @@ class BartSystem(BaseTransformer):
         self.lr_scheduler = scheduler
         return dataloader
 
-    def val_dataloader(self):
-        val_dataset = SummarizationDataset(
-            self.tokenizer, data_dir=self.hparams.data_dir, type_path="val", block_size=self.hparams.max_seq_length
-        )
-        return DataLoader(val_dataset, batch_size=self.hparams.eval_batch_size)
+    def val_dataloader(self) -> DataLoader:
+        return self.get_dataloader("val", batch_size=self.hparams.eval_batch_size)
 
-    def test_dataloader(self):
-        test_dataset = SummarizationDataset(
-            self.tokenizer, data_dir=self.hparams.data_dir, type_path="test", block_size=self.hparams.max_seq_length
-        )
-        return DataLoader(test_dataset, batch_size=self.hparams.eval_batch_size)
+    def test_dataloader(self) -> DataLoader:
+        return self.get_dataloader("test", batch_size=self.hparams.eval_batch_size)
 
     @staticmethod
     def add_model_specific_args(parser, root_dir):
         BaseTransformer.add_model_specific_args(parser, root_dir)
         # Add BART specific options
         parser.add_argument(
-            "--max_seq_length",
+            "--max_source_length",
             default=1024,
             type=int,
             help="The maximum total input sequence length after tokenization. Sequences longer "
             "than this will be truncated, sequences shorter will be padded.",
         )
+        parser.add_argument(
+            "--max_target_length",
+            default=56,
+            type=int,
+            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",
@@ -158,7 +160,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # If output_dir not provided, a folder will be generated in pwd
-    if args.output_dir is None:
+    if not args.output_dir:
         args.output_dir = os.path.join("./results", f"{args.task}_{args.model_type}_{time.strftime('%Y%m%d_%H%M%S')}",)
         os.makedirs(args.output_dir)