Fix ProphetNet forward to handle tuple encoder_outputs

src/transformers/models/prophetnet/modeling_prophetnet.py

@@ -25,13 +25,16 @@
 from torch import Tensor, nn
 from torch.nn import LayerNorm
 
+from transformers.utils.generic import check_model_inputs
+
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
 from ...generation import GenerationMixin
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutput
 from ...modeling_utils import PreTrainedModel
-from ...utils import ModelOutput, auto_docstring, logging
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging
 from ...utils.deprecation import deprecate_kwarg
 from .configuration_prophetnet import ProphetNetConfig
 
@@ -1048,16 +1051,15 @@ def get_input_embeddings(self):
     def set_input_embeddings(self, value):
         self.word_embeddings = value
 
+    @check_model_inputs
     @auto_docstring
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         head_mask: Optional[torch.Tensor] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, BaseModelOutput]:
         r"""
         Example:
@@ -1073,12 +1075,14 @@ def forward(
 
         >>> last_hidden_states = outputs.last_hidden_state
         ```"""
+        # Extract parameters from kwargs or fall back to config
+        output_attentions = kwargs.get("output_attentions")
+        if output_attentions is None:
+            output_attentions = getattr(self.config, "output_attentions", False)
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = kwargs.get("output_hidden_states")
+        if output_hidden_states is None:
+            output_hidden_states = getattr(self.config, "output_hidden_states", False)
 
         if input_ids is None and inputs_embeds is None:
             raise ValueError("Either input_ids or inputs_embeds has to be passed.")
@@ -1129,8 +1133,6 @@ def forward(
         if output_hidden_states:
             encoder_hidden_states = encoder_hidden_states + (hidden_states,)
 
-        if not return_dict:
-            return tuple(v for v in [hidden_states, encoder_hidden_states, all_attentions] if v is not None)
         return BaseModelOutput(
             last_hidden_state=hidden_states, hidden_states=encoder_hidden_states, attentions=all_attentions
         )
@@ -1179,6 +1181,7 @@ def get_input_embeddings(self):
     def set_input_embeddings(self, value):
         self.word_embeddings = value
 
+    @check_model_inputs
     @auto_docstring
     def forward(
         self,
@@ -1191,10 +1194,8 @@ def forward(
         past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, ProphetNetDecoderModelOutput]:
         r"""
         cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
@@ -1216,12 +1217,16 @@ def forward(
 
         >>> last_hidden_states = outputs.last_hidden_state
         ```"""
+        # Extract parameters from kwargs or fall back to config
+        output_attentions = kwargs.get("output_attentions")
+        if output_attentions is None:
+            output_attentions = getattr(self.config, "output_attentions", False)
+
+        output_hidden_states = kwargs.get("output_hidden_states")
+        if output_hidden_states is None:
+            output_hidden_states = getattr(self.config, "output_hidden_states", False)
+
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if input_ids is None and inputs_embeds is None:
             raise ValueError("Either `decoder_input_ids` or `decoder_inputs_embeds` has to be passed.")
@@ -1366,21 +1371,6 @@ def forward(
         last_hidden_state = hidden_states[:, :sequence_length]
         last_hidden_state_ngram = hidden_states[:, sequence_length:] if self.config.ngram > 0 else None
 
-        if not return_dict:
-            return tuple(
-                v
-                for v in [
-                    last_hidden_state,
-                    last_hidden_state_ngram,
-                    past_key_values,
-                    all_main_stream_hidden_states,
-                    all_ngram_stream_hidden_states,
-                    all_main_stream_attns,
-                    all_ngram_stream_attns,
-                    all_cross_attns,
-                ]
-                if v is not None
-            )
         return ProphetNetDecoderModelOutput(
             last_hidden_state=last_hidden_state,
             last_hidden_state_ngram=last_hidden_state_ngram,
@@ -1511,6 +1501,7 @@ def _tie_weights(self):
     def get_encoder(self):
         return self.encoder
 
+    @check_model_inputs
     @auto_docstring
     def forward(
         self,
@@ -1521,15 +1512,13 @@ def forward(
         head_mask: Optional[torch.Tensor] = None,
         decoder_head_mask: Optional[torch.Tensor] = None,
         cross_attn_head_mask: Optional[torch.Tensor] = None,
-        encoder_outputs: Optional[tuple] = None,
+        encoder_outputs: Optional[BaseModelOutput] = None,
         past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         decoder_inputs_embeds: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, ProphetNetSeq2SeqModelOutput]:
         r"""
         decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
@@ -1570,21 +1559,14 @@ def forward(
         >>> last_hidden_states_ngram = outputs.last_hidden_state_ngram  # predict hidden states
         ```"""
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if encoder_outputs is None:
             encoder_outputs = self.encoder(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
                 head_mask=head_mask,
                 inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
+                **kwargs,
             )
 
         # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn)
@@ -1597,15 +1579,11 @@ def forward(
             cross_attn_head_mask=cross_attn_head_mask,
             past_key_values=past_key_values,
             inputs_embeds=decoder_inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
             use_cache=use_cache,
-            return_dict=return_dict,
             cache_position=cache_position,
+            **kwargs,
         )
 
-        if not return_dict:
-            return decoder_outputs + encoder_outputs
         return ProphetNetSeq2SeqModelOutput(
             last_hidden_state=decoder_outputs.last_hidden_state,
             last_hidden_state_ngram=decoder_outputs.last_hidden_state_ngram,
@@ -1647,6 +1625,7 @@ def _tie_weights(self):
     def get_input_embeddings(self):
         return self.prophetnet.word_embeddings
 
+    @check_model_inputs
     @auto_docstring
     def forward(
         self,
@@ -1663,10 +1642,8 @@ def forward(
         decoder_inputs_embeds: Optional[torch.Tensor] = None,
         labels: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, ProphetNetSeq2SeqLMOutput]:
         r"""
         decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
@@ -1710,7 +1687,6 @@ def forward(
         >>> logits_next_token = outputs.logits  # logits to predict next token as usual
         >>> logits_ngram_next_tokens = outputs.logits_ngram  # logits to predict 2nd, 3rd, ... next tokens
         ```"""
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
             # get decoder inputs from shifting lm labels to the right
@@ -1729,10 +1705,8 @@ def forward(
             inputs_embeds=inputs_embeds,
             decoder_inputs_embeds=decoder_inputs_embeds,
             use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
             cache_position=cache_position,
+            **kwargs,
         )
         batch_size, sequence_length = (
             decoder_input_ids.shape if decoder_input_ids is not None else decoder_inputs_embeds.shape[:2]
@@ -1752,24 +1726,20 @@ def forward(
         if labels is not None:
             loss = self._compute_loss(predict_logits, labels)
 
-        if not return_dict:
-            all_logits = tuple(v for v in [logits, logits_ngram] if v is not None)
-            return (loss,) + all_logits + outputs[2:] if loss is not None else all_logits + outputs[2:]
-        else:
-            return ProphetNetSeq2SeqLMOutput(
-                loss=loss,
-                logits=logits,
-                logits_ngram=logits_ngram,
-                past_key_values=outputs.past_key_values,
-                decoder_hidden_states=outputs.decoder_hidden_states,
-                decoder_ngram_hidden_states=outputs.decoder_ngram_hidden_states,
-                decoder_attentions=outputs.decoder_attentions,
-                decoder_ngram_attentions=outputs.decoder_ngram_attentions,
-                cross_attentions=outputs.cross_attentions,
-                encoder_last_hidden_state=outputs.encoder_last_hidden_state,
-                encoder_hidden_states=outputs.encoder_hidden_states,
-                encoder_attentions=outputs.encoder_attentions,
-            )
+        return ProphetNetSeq2SeqLMOutput(
+            loss=loss,
+            logits=logits,
+            logits_ngram=logits_ngram,
+            past_key_values=outputs.past_key_values,
+            decoder_hidden_states=outputs.decoder_hidden_states,
+            decoder_ngram_hidden_states=outputs.decoder_ngram_hidden_states,
+            decoder_attentions=outputs.decoder_attentions,
+            decoder_ngram_attentions=outputs.decoder_ngram_attentions,
+            cross_attentions=outputs.cross_attentions,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            encoder_hidden_states=outputs.encoder_hidden_states,
+            encoder_attentions=outputs.encoder_attentions,
+        )
 
     def _compute_loss(self, logits, labels, ignore_index=-100):
         expend_targets = labels.new_zeros(self.config.ngram, labels.size(0), labels.size(1)).fill_(ignore_index)
@@ -1853,6 +1823,7 @@ def set_decoder(self, decoder):
     def get_decoder(self):
         return self.prophetnet.decoder
 
+    @check_model_inputs
     @auto_docstring
     def forward(
         self,
@@ -1866,9 +1837,7 @@ def forward(
         inputs_embeds: Optional[torch.Tensor] = None,
         labels: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, ProphetNetDecoderLMOutput]:
         r"""
         cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
@@ -1918,7 +1887,6 @@ def forward(
 
         >>> loss = outputs.loss
         ```"""
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn)
         outputs = self.prophetnet.decoder(
@@ -1931,9 +1899,7 @@ def forward(
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
+            **kwargs,
         )
 
         batch_size, sequence_length = input_ids.shape if input_ids is not None else inputs_embeds.shape[:2]
@@ -1948,21 +1914,17 @@ def forward(
         if labels is not None:
             loss = self._compute_loss(predict_logits, labels)
 
-        if not return_dict:
-            all_logits = tuple(v for v in [logits, logits_ngram] if v is not None)
-            return (loss,) + all_logits + outputs[2:] if loss is not None else all_logits + outputs[2:]
-        else:
-            return ProphetNetDecoderLMOutput(
-                loss=loss,
-                logits=logits,
-                logits_ngram=logits_ngram,
-                past_key_values=outputs.past_key_values,
-                hidden_states=outputs.hidden_states,
-                hidden_states_ngram=outputs.hidden_states_ngram,
-                attentions=outputs.attentions,
-                ngram_attentions=outputs.ngram_attentions,
-                cross_attentions=outputs.cross_attentions,
-            )
+        return ProphetNetDecoderLMOutput(
+            loss=loss,
+            logits=logits,
+            logits_ngram=logits_ngram,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            hidden_states_ngram=outputs.hidden_states_ngram,
+            attentions=outputs.attentions,
+            ngram_attentions=outputs.ngram_attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
 
     def _compute_loss(self, logits, labels, ignore_index=-100):
         expend_targets = labels.new_zeros(self.config.ngram, labels.size(0), labels.size(1)).fill_(ignore_index)

src/transformers/utils/generic.py

@@ -972,7 +972,7 @@ def wrapper(self, *args, **kwargs):
             for k, v in all_args["kwargs"].items():
                 all_args[k] = v
 
-        capture_flags = _CAN_RECORD_REGISTRY.get(str(self.__class__), {})  # there is a weak ref for executorch
+        capture_flags = _CAN_RECORD_REGISTRY.get(str(self.__class__)) or {}  # there is a weak ref for executorch
         recordable_keys = {
             f"output_{k}": all_args.get(
                 f"output_{k}",