[Seq2Seq] Allow EncoderDecoderModels to be trained with Seq2Seq

examples/seq2seq/seq2seq_trainer.py

@@ -41,12 +41,13 @@
 
 
 class Seq2SeqTrainer(Trainer):
-    def __init__(self, config, data_args, *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.config = config
-        self.data_args = data_args
-        self.max_gen_length = data_args.val_max_target_length
-        self.vocab_size = self.config.tgt_vocab_size if isinstance(self.config, FSMTConfig) else self.config.vocab_size
+        self.vocab_size = (
+            self.model.config.tgt_vocab_size
+            if isinstance(self.model.config, FSMTConfig)
+            else self.model.config.vocab_size
+        )
 
     def create_optimizer_and_scheduler(self, num_training_steps: int):
         """
@@ -114,23 +115,22 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
                 else DistributedSampler(self.train_dataset)
             )
 
-    def compute_loss(self, model, inputs):
-        labels = inputs.pop("labels")
-        outputs = model(**inputs, use_cache=False)
-        logits = outputs[0]
-        return self._compute_loss(logits, labels)
-
-    def _compute_loss(self, logits, labels):
+    def _compute_loss(self, model, inputs):
         if self.args.label_smoothing == 0:
-            # Same behavior as modeling_bart.py
-            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
-            assert logits.shape[-1] == self.vocab_size
-            loss = loss_fct(logits.view(-1, logits.shape[-1]), labels.view(-1))
+            # compute usual loss via models
+            loss, logits = model(**inputs)[:2]
         else:
+            # compute label smoothed loss
+            labels = inputs.pop("labels")
+            logits = model(**inputs, use_cache=False)[0]
             lprobs = torch.nn.functional.log_softmax(logits, dim=-1)
-            loss, nll_loss = label_smoothed_nll_loss(
-                lprobs, labels, self.args.label_smoothing, ignore_index=self.config.pad_token_id
+            loss, _ = label_smoothed_nll_loss(
+                lprobs, labels, self.args.label_smoothing, ignore_index=self.model.config.pad_token_id
             )
+        return loss, logits
+
+    def compute_loss(self, model, inputs):
+        loss, _ = self._compute_loss(model, inputs)
         return loss
 
     def prediction_step(
@@ -158,34 +158,29 @@ def prediction_step(
         """
         inputs = self._prepare_inputs(inputs)
 
+        if self.args.predict_with_generate and not self.args.prediction_loss_only:
+            generated_tokens = model.generate(
+                inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+            )
+            # in case the batch is shorter than max length, the output should be padded
+            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, self.model.config.max_length)
+
+            # compute loss on predict data
         with torch.no_grad():
-            if self.args.predict_with_generate and not self.args.prediction_loss_only:
-                generated_tokens = model.generate(
-                    inputs["input_ids"],
-                    attention_mask=inputs["attention_mask"],
-                    use_cache=True,
-                    num_beams=self.data_args.eval_beams,
-                    max_length=self.max_gen_length,
-                )
-                # in case the batch is shorter than max length, the output should be padded
-                generated_tokens = self._pad_tensors_to_max_len(generated_tokens, self.max_gen_length)
+            loss, logits = self._compute_loss(model, inputs)
 
-            labels_out = inputs.get("labels")
-            # Call forward again to get loss # TODO: avoidable?
-            outputs = model(**inputs, use_cache=False)
-            loss = self._compute_loss(outputs[1], labels_out)
-            loss = loss.mean().detach()
-            if self.args.prediction_loss_only:
-                return (loss, None, None)
+        loss = loss.mean().detach()
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
 
-            logits = generated_tokens if self.args.predict_with_generate else outputs[1]
+        logits = generated_tokens if self.args.predict_with_generate else logits
+        labels = self._pad_tensors_to_max_len(inputs["labels"], self.model.config.max_length)
 
-        labels_out = labels_out.detach()
-        labels = self._pad_tensors_to_max_len(labels_out, self.max_gen_length)
-        return (loss, logits.detach(), labels)
+        return (loss, logits, labels)
 
     def _pad_tensors_to_max_len(self, tensor, max_length):
-        padded_tensor = self.config.pad_token_id * torch.ones(
+        padded_tensor = self.model.config.pad_token_id * torch.ones(
             (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device
         )
         padded_tensor[:, : tensor.shape[-1]] = tensor