Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Classes to support TF Encoder-Decoder architectures"""16 17from __future__ import annotations18 19import inspect20import re21import warnings22 23import numpy as np24import tensorflow as tf25 26from ...configuration_utils import PretrainedConfig27from ...modeling_tf_outputs import TFBaseModelOutput, TFSeq2SeqLMOutput28from ...modeling_tf_utils import (29    TFCausalLanguageModelingLoss,30    TFModelInputType,31    TFPreTrainedModel,32    get_initializer,33    keras,34    unpack_inputs,35)36from ...tf_utils import shape_list37from ...utils import (38    ModelOutput,39    add_start_docstrings,40    add_start_docstrings_to_model_forward,41    logging,42    replace_return_docstrings,43)44from ..auto.configuration_auto import AutoConfig45from ..auto.modeling_tf_auto import TFAutoModel, TFAutoModelForCausalLM46from .configuration_encoder_decoder import EncoderDecoderConfig47 48 49logger = logging.get_logger(__name__)50 51_CONFIG_FOR_DOC = "EncoderDecoderConfig"52 53DEPRECATION_WARNING = (54    "Version v4.17.0 introduces a better way to train encoder-decoder models by computing the loss inside the"55    " encoder-decoder framework rather than in the decoder itself. You may observe training discrepancies if"56    " fine-tuning a model trained with versions anterior to 4.17.0. The decoder_input_ids are now created based on the"57    " labels, no need to pass them yourself anymore."58)59 60ENCODER_DECODER_START_DOCSTRING = r"""61    This class can be used to initialize a sequence-to-sequence model with any pretrained autoencoding model as the62    encoder and any pretrained autoregressive model as the decoder. The encoder is loaded via63    [`~TFAutoModel.from_pretrained`] function and the decoder is loaded via [`~TFAutoModelForCausalLM.from_pretrained`]64    function. Cross-attention layers are automatically added to the decoder and should be fine-tuned on a downstream65    generative task, like summarization.66 67    The effectiveness of initializing sequence-to-sequence models with pretrained checkpoints for sequence generation68    tasks was shown in [Leveraging Pre-trained Checkpoints for Sequence Generation69    Tasks](https://huggingface.co/papers/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. Michael Matena, Yanqi70    Zhou, Wei Li, Peter J. Liu.71 72    After such an Encoder Decoder model has been trained/fine-tuned, it can be saved/loaded just like any other models73    (see the examples for more information).74 75    This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the76    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads77    etc.)78 79    This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it80    as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and81    behavior.82 83    Parameters:84        config ([`EncoderDecoderConfig`]): Model configuration class with all the parameters of the model.85            Initializing with a config file does not load the weights associated with the model, only the86            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.87"""88 89ENCODER_DECODER_INPUTS_DOCSTRING = r"""90    Args:91        input_ids (`np.ndarray`, `tf.Tensor`, `list[tf.Tensor]` ``dict[str, tf.Tensor]` or `dict[str, np.ndarray]` and each example must have the shape `({0})`):92            Indices of input sequence tokens in the vocabulary.93 94            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and95            [`PreTrainedTokenizer.__call__`] for details.96 97            [What are input IDs?](../glossary#input-ids)98        attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):99            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:100 101            - 1 for tokens that are **not masked**,102            - 0 for tokens that are **masked**.103 104            [What are attention masks?](../glossary#attention-mask)105        decoder_input_ids (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*):106            Indices of decoder input sequence tokens in the vocabulary.107 108            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and109            [`PreTrainedTokenizer.__call__`] for details.110 111            [What are input IDs?](../glossary#input-ids)112 113            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see114            `past_key_values`).115 116            Provide for sequence to sequence training to the decoder. Indices can be obtained using117            [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for118            details.119        decoder_attention_mask (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*):120            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also121            be used by default.122        encoder_outputs (`tuple(tuple(tf.Tensor)`, *optional*):123            This tuple must consist of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)124            `last_hidden_state` (`tf.Tensor` of shape `({0}, hidden_size)`) is a tensor of hidden-states at the output125            of the last layer of the encoder. Used in the cross-attention of the decoder.126        past_key_values (`tuple(tuple(tf.Tensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):127            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.128 129            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that130            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all131            `decoder_input_ids` of shape `({0})`.132        inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `({0}, hidden_size)`, *optional*):133            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This134            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the135            model's internal embedding lookup matrix.136        decoder_inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):137            Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded138            representation. This is useful if you want more control over how to convert `decoder_input_ids` indices139            into associated vectors than the model's internal embedding lookup matrix.140        labels (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):141            Labels for computing the masked language modeling loss for the decoder. Indices should be in `[-100, 0,142            ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored143            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`144        use_cache (`bool`, *optional*):145            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see146            `past_key_values`).147        output_attentions (`bool`, *optional*):148            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned149            tensors for more detail.150        output_hidden_states (`bool`, *optional*):151            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for152            more detail.153        return_dict (`bool`, *optional*):154            If set to `True`, the model will return a [`~utils.Seq2SeqLMOutput`] instead of a plain tuple.155        training (`bool`, *optional*, defaults to `False`):156            Whether or not to use the model in training mode (some modules like dropout modules have different157            behaviors between training and evaluation).158        kwargs (*optional*): Remaining dictionary of keyword arguments. Keyword arguments come in two flavors:159 160            - Without a prefix which will be input as `**encoder_kwargs` for the encoder forward function.161            - With a *decoder_* prefix which will be input as `**decoder_kwargs`` for the decoder forward function.162"""163 164 165def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):166    if pad_token_id is None:167        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")168    pad_token_id = tf.cast(pad_token_id, input_ids.dtype)169 170    if decoder_start_token_id is None:171        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")172    decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)173 174    start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)175    shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)176    # replace possible -100 values in labels by `pad_token_id`177    shifted_input_ids = tf.where(178        shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids179    )180 181    # "Verify that `labels` has only positive values and -100"182    assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))183 184    # Make sure the assertion op is called by wrapping the result in an identity no-op185    with tf.control_dependencies([assert_gte0]):186        shifted_input_ids = tf.identity(shifted_input_ids)187 188    return shifted_input_ids189 190 191@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)192class TFEncoderDecoderModel(TFPreTrainedModel, TFCausalLanguageModelingLoss):193    r"""194    [`TFEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with one195    of the base model classes of the library as encoder and another one as decoder when created with the196    [`~TFAutoModel.from_pretrained`] class method for the encoder and [`~TFAutoModelForCausalLM.from_pretrained`] class197    method for the decoder.198    """199 200    config_class = EncoderDecoderConfig201    base_model_prefix = "encoder_decoder"202    load_weight_prefix = "tf_encoder_decoder_model"203 204    def __init__(205        self,206        config: PretrainedConfig | None = None,207        encoder: TFPreTrainedModel | None = None,208        decoder: TFPreTrainedModel | None = None,209    ):210        if config is None and (encoder is None or decoder is None):211            raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")212        if config is None:213            config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)214        else:215            if not isinstance(config, self.config_class):216                raise ValueError(f"config: {config} has to be of type {self.config_class}")217 218        if config.decoder.cross_attention_hidden_size is not None:219            if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:220                raise ValueError(221                    "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"222                    f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"223                    f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"224                    " `config.encoder.hidden_size`."225                )226 227        # initialize with config228        super().__init__(config)229 230        if encoder is None:231            encoder = TFAutoModel.from_config(config.encoder, name="encoder")232 233        if decoder is None:234            decoder = TFAutoModelForCausalLM.from_config(config.decoder, name="decoder")235 236        self.encoder = encoder237        self.decoder = decoder238 239        if self.encoder.config.to_dict() != self.config.encoder.to_dict():240            logger.warning(241                f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"242                f" {self.config.encoder}"243            )244        if self.decoder.config.to_dict() != self.config.decoder.to_dict():245            logger.warning(246                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"247                f" {self.config.decoder}"248            )249 250        # make sure that the individual model's config refers to the shared config251        # so that the updates to the config will be synced252        self.encoder.config = self.config.encoder253        self.decoder.config = self.config.decoder254 255        # encoder outputs might need to be projected to different dimension for decoder256        if (257            self.encoder.config.hidden_size != self.decoder.config.hidden_size258            and self.decoder.config.cross_attention_hidden_size is None259        ):260            self.enc_to_dec_proj = keras.layers.Dense(261                units=self.decoder.config.hidden_size,262                kernel_initializer=get_initializer(config.encoder.initializer_range),263                name="enc_to_dec_proj",264            )265 266        if self.encoder.get_output_embeddings() is not None:267            raise ValueError(268                f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"269            )270 271        decoder_signature = set(inspect.signature(self.decoder.call).parameters.keys())272        if "encoder_hidden_states" not in decoder_signature:273            raise ValueError(274                "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "275                "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"276            )277 278    def get_encoder(self):279        return self.encoder280 281    def get_input_embeddings(self):282        return self.encoder.get_input_embeddings()283 284    def get_output_embeddings(self):285        return self.decoder.get_output_embeddings()286 287    def set_output_embeddings(self, new_embeddings):288        return self.decoder.set_output_embeddings(new_embeddings)289 290    def tf_to_pt_weight_rename(self, tf_weight):291        # Matt: The TF and PT weights don't align because our TF base classes have an extra layer compared to PT models292        # (the main model stem is in the MainLayer class). If we remove that layer, then weight names sync up as normal.293        # However, the name of that extra layer is the name of the MainLayer in the base model. We make the assumption294        # here that the config model_type is the same as the name of the MainLayer. I don't know of anywhere that's295        # not the case, and I wasn't sure how else to go from the config to the correct MainLayer name!296 297        # This override is only needed in the case where we're crossloading weights from PT. However, since weights are298        # often safetensors now, we don't know if we're going to be crossloading until we sniff the weights file.299        # Therefore, we specify tf_to_pt_weight_rename anyway, and let the super method figure out if it needs it300        # or not.301        encoder_model_type = self.config.encoder.model_type302        if "encoder" in tf_weight and "decoder" not in tf_weight:303            return (re.sub(rf"encoder\.{encoder_model_type}\.", "encoder.", tf_weight),)304        else:305            return (tf_weight,)306 307    @classmethod308    def from_encoder_decoder_pretrained(309        cls,310        encoder_pretrained_model_name_or_path: str | None = None,311        decoder_pretrained_model_name_or_path: str | None = None,312        *model_args,313        **kwargs,314    ) -> TFPreTrainedModel:315        r"""316        Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model317        checkpoints.318 319 320        Params:321            encoder_pretrained_model_name_or_path (`str`, *optional*):322                Information necessary to initiate the encoder. Can be either:323 324                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.325                    - A path to a *directory* containing model weights saved using326                      [`~TFPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.327                    - A path or url to a *pytorch index checkpoint file* (e.g, `./pt_model/`). In this case,328                      `encoder_from_pt` should be set to `True`.329 330            decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):331                Information necessary to initiate the decoder. Can be either:332 333                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.334                    - A path to a *directory* containing model weights saved using335                      [`~TFPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.336                    - A path or url to a *pytorch checkpoint file* (e.g, `./pt_model/`). In this case,337                      `decoder_from_pt` should be set to `True`.338 339            model_args (remaining positional arguments, *optional*):340                All remaining positional arguments will be passed to the underlying model's `__init__` method.341 342            kwargs (remaining dictionary of keyword arguments, *optional*):343                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,344                `output_attentions=True`).345 346                - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.347                - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.348                - To update the parent model configuration, do not use a prefix for each configuration parameter.349 350                Behaves differently depending on whether a `config` is provided or automatically loaded.351 352        Example:353 354        ```python355        >>> from transformers import TFEncoderDecoderModel356 357        >>> # initialize a bert2gpt2 from two pretrained BERT models. Note that the cross-attention layers will be randomly initialized358        >>> model = TFEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-uncased", "openai-community/gpt2")359        >>> # saving model after fine-tuning360        >>> model.save_pretrained("./bert2gpt2")361        >>> # load fine-tuned model362        >>> model = TFEncoderDecoderModel.from_pretrained("./bert2gpt2")363        ```"""364 365        kwargs_encoder = {366            argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")367        }368 369        kwargs_decoder = {370            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")371        }372 373        # remove encoder, decoder kwargs from kwargs374        for key in kwargs_encoder:375            del kwargs["encoder_" + key]376        for key in kwargs_decoder:377            del kwargs["decoder_" + key]378 379        # Load and initialize the encoder and decoder380        # The distinction between encoder and decoder at the model level is made381        # by the value of the flag `is_decoder` that we need to set correctly.382        encoder = kwargs_encoder.pop("model", None)383        if encoder is None:384            if encoder_pretrained_model_name_or_path is None:385                raise ValueError(386                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "387                    "to be defined."388                )389 390            if "config" not in kwargs_encoder:391                encoder_config = AutoConfig.from_pretrained(encoder_pretrained_model_name_or_path)392                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:393                    logger.info(394                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "395                        "from a decoder model. Cross-attention and causal mask are disabled."396                    )397                    encoder_config.is_decoder = False398                    encoder_config.add_cross_attention = False399 400                kwargs_encoder["config"] = encoder_config401 402            kwargs_encoder["name"] = "encoder"403            kwargs_encoder["load_weight_prefix"] = cls.load_weight_prefix404            encoder = TFAutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)405 406        decoder = kwargs_decoder.pop("model", None)407        if decoder is None:408            if decoder_pretrained_model_name_or_path is None:409                raise ValueError(410                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "411                    "to be defined."412                )413 414            if "config" not in kwargs_decoder:415                decoder_config = AutoConfig.from_pretrained(decoder_pretrained_model_name_or_path)416                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:417                    logger.info(418                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"419                        f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"420                        f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."421                    )422                    decoder_config.is_decoder = True423                    decoder_config.add_cross_attention = True424 425                kwargs_decoder["config"] = decoder_config426 427            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:428                logger.warning(429                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "430                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "431                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "432                    "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "433                    "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"434                )435 436            kwargs_decoder["name"] = "decoder"437            kwargs_decoder["load_weight_prefix"] = cls.load_weight_prefix438            decoder = TFAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)439 440        # Make sure these 2 `keras.Model` have fixed names so `from_pretrained` could load model weights correctly.441        if encoder.name != "encoder":442            raise ValueError("encoder model must be created with the name `encoder`.")443        if decoder.name != "decoder":444            raise ValueError("decoder model must be created with the name `decoder`.")445 446        # instantiate config with corresponding kwargs447        config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)448        return cls(encoder=encoder, decoder=decoder, config=config)449 450    @unpack_inputs451    @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))452    @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)453    def call(454        self,455        input_ids: TFModelInputType | None = None,456        attention_mask: np.ndarray | tf.Tensor | None = None,457        decoder_input_ids: np.ndarray | tf.Tensor | None = None,458        decoder_attention_mask: np.ndarray | tf.Tensor | None = None,459        encoder_outputs: np.ndarray | tf.Tensor | None = None,460        past_key_values: tuple[tuple[tf.Tensor]] | None = None,461        inputs_embeds: np.ndarray | tf.Tensor | None = None,462        decoder_inputs_embeds: np.ndarray | tf.Tensor | None = None,463        labels: np.ndarray | tf.Tensor | None = None,464        use_cache: bool | None = None,465        output_attentions: bool | None = None,466        output_hidden_states: bool | None = None,467        return_dict: bool | None = None,468        training: bool = False,469        **kwargs,470    ) -> TFSeq2SeqLMOutput | tuple[tf.Tensor]:471        r"""472        Returns:473 474        Examples:475 476        ```python477        >>> from transformers import TFEncoderDecoderModel, BertTokenizer478 479        >>> # initialize a bert2gpt2 from a pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized480        >>> model = TFEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")481 482        >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")483 484        >>> # forward485        >>> input_ids = tokenizer.encode(486        ...     "Hello, my dog is cute", add_special_tokens=True, return_tensors="tf"487        ... )  # Batch size 1488        >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids)489 490        >>> # training491        >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, labels=input_ids)492        >>> loss, logits = outputs.loss, outputs.logits493 494        >>> # save and load from pretrained495        >>> model.save_pretrained("bert2gpt2")496        >>> model = TFEncoderDecoderModel.from_pretrained("bert2gpt2")497 498        >>> # generation499        >>> generated = model.generate(input_ids, decoder_start_token_id=model.config.decoder.bos_token_id)500        ```"""501        return_dict = return_dict if return_dict is not None else self.config.use_return_dict502 503        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}504 505        kwargs_decoder = {506            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")507        }508 509        # Let the user be responsible for the expected format.510        if encoder_outputs is not None:511            if return_dict and not isinstance(encoder_outputs, ModelOutput):512                raise ValueError(513                    "If `return_dict=True` and `encoder_outputs` is provided, it should be an instance of "514                    f"`ModelOutput`. Got an instance {type(encoder_outputs)} for `encoder_outputs`."515                )516 517        if encoder_outputs is None:518            encoder_inputs = {519                "input_ids": input_ids,520                "attention_mask": attention_mask,521                "inputs_embeds": inputs_embeds,522                "output_attentions": output_attentions,523                "output_hidden_states": output_hidden_states,524                "return_dict": return_dict,525                "training": training,526            }527 528            # Add arguments to encoder from `kwargs_encoder`529            encoder_inputs.update(kwargs_encoder)530 531            # Handle the case where the inputs are passed as a single dict which contains `labels`.532            # The `labels` shouldn't be passed to `self.encoder` below, because it is a based model without this533            # parameter (otherwise, an error occurs when `input_processing` is called inside `self.encoder.call()`).534            if "labels" in encoder_inputs:535                labels = encoder_inputs.pop("labels")536 537            # handle the init case where `dummy_inputs` returns a dict containing `decoder_input_ids`.538            if "decoder_input_ids" in encoder_inputs:539                decoder_input_ids = encoder_inputs.pop("decoder_input_ids")540            # handle the init case where `dummy_inputs` returns a dict containing `decoder_input_ids`.541            if "decoder_attention_mask" in encoder_inputs:542                decoder_attention_mask = encoder_inputs.pop("decoder_attention_mask")543 544            encoder_outputs = self.encoder(**encoder_inputs)545 546        encoder_hidden_states = encoder_outputs[0]547 548        # optionally project encoder_hidden_states549        if (550            self.encoder.config.hidden_size != self.decoder.config.hidden_size551            and self.decoder.config.cross_attention_hidden_size is None552        ):553            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)554 555        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):556            decoder_input_ids = shift_tokens_right(557                labels, self.config.pad_token_id, self.config.decoder_start_token_id558            )559 560        decoder_inputs = {561            "input_ids": decoder_input_ids,562            "attention_mask": decoder_attention_mask,563            "encoder_hidden_states": encoder_hidden_states,564            "encoder_attention_mask": attention_mask,565            "inputs_embeds": decoder_inputs_embeds,566            "output_attentions": output_attentions,567            "output_hidden_states": output_hidden_states,568            "use_cache": use_cache,569            "past_key_values": past_key_values,570            "return_dict": return_dict,571            "training": training,572        }573 574        # Add arguments to decoder from `kwargs_decoder`575        decoder_inputs.update(kwargs_decoder)576 577        decoder_outputs = self.decoder(**decoder_inputs)578 579        logits = decoder_outputs[0]580 581        # Compute loss independent from decoder (as some shift the logits inside them)582        loss = None583        if labels is not None:584            warnings.warn(DEPRECATION_WARNING, FutureWarning)585            loss = self.hf_compute_loss(labels, logits)586 587        if not return_dict:588            past_key_values = None589            if use_cache:590                past_key_values = decoder_outputs[1]591            # The starting index of the remaining elements in `decoder_outputs`592            start_index = sum([1 if x is not None else 0 for x in (loss, logits, past_key_values)])593 594            if not isinstance(encoder_outputs, tuple):595                encoder_outputs = encoder_outputs.to_tuple()596            output = (loss, logits, past_key_values) + decoder_outputs[start_index:] + encoder_outputs597            output = tuple(x for x in output if x is not None)598            return output599 600        return TFSeq2SeqLMOutput(601            loss=loss,602            logits=decoder_outputs.logits,603            past_key_values=decoder_outputs.past_key_values,604            decoder_hidden_states=decoder_outputs.hidden_states,605            decoder_attentions=decoder_outputs.attentions,606            cross_attentions=decoder_outputs.cross_attentions,607            encoder_last_hidden_state=encoder_outputs.last_hidden_state,608            encoder_hidden_states=encoder_outputs.hidden_states,609            encoder_attentions=encoder_outputs.attentions,610        )611 612    def prepare_inputs_for_generation(613        self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs614    ):615        decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)616        decoder_attention_mask = decoder_inputs.get("attention_mask", None)617        past_key_values = decoder_inputs.get("past_key_values")618        if past_key_values is None:619            past_key_values = decoder_inputs.get("past")  # e.g. on TF GPT2620        input_dict = {621            "input_ids": None,  # needs to be passed to make Keras.layer.__call__ happy622            "attention_mask": attention_mask,623            "decoder_attention_mask": decoder_attention_mask,624            "decoder_input_ids": decoder_inputs["input_ids"],625            # TODO (joao): the `TFBaseModelOutput` wrapper should not be needed after the generate refactor is complete626            "encoder_outputs": TFBaseModelOutput(last_hidden_state=encoder_outputs[0]),627            "past_key_values": past_key_values,628            "use_cache": use_cache,629        }630        return input_dict631 632    def prepare_decoder_input_ids_from_labels(self, labels: tf.Tensor):633        return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)634 635    def resize_token_embeddings(self, *args, **kwargs):636        raise NotImplementedError(637            "Resizing the embedding layers via the TFEncoderDecoderModel directly is not supported.Please use the"638            " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"639            " model.decoder.resize_token_embeddings(...))"640        )641 642    def _reorder_cache(self, past, beam_idx):643        # apply decoder cache reordering here644        return self.decoder._reorder_cache(past, beam_idx)645 646    def build(self, input_shape=None):647        if self.built:648            return649        self.built = True650        if getattr(self, "enc_to_dec_proj", None) is not None:651            with tf.name_scope(self.enc_to_dec_proj.name):652                self.enc_to_dec_proj.build([None, None, self.encoder.config.hidden_size])653        if getattr(self, "encoder", None) is not None:654            with tf.name_scope(self.encoder.name):655                self.encoder.build(None)656        if getattr(self, "decoder", None) is not None:657            with tf.name_scope(self.decoder.name):658                self.decoder.build(None)659 660 661__all__ = ["TFEncoderDecoderModel"]662