codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license classes to support tf encoderdecoder architectures from future import annotations import inspect import re import warnings from typing import optional tuple union import numpy as np import tensorflow as tf from configurationutils import pretrainedconfig from modelingtfoutputs import tfbasemodeloutput tfseq2seqlmoutput from modelingtfutils import tfcausallanguagemodelingloss tfmodelinputtype tfpretrainedmodel getinitializer unpackinputs from tfutils import shapelist from utils import modeloutput addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from auto configurationauto import autoconfig from auto modelingtfauto import tfautomodel tfautomodelforcausallm from configurationencoderdecoder import encoderdecoderconfig logger logging getloggername configfordoc encoderdecoderconfig deprecationwarning version v4 17 0 introduces a better way to train encoderdecoder models by computing the loss inside the encoderdecoder framework rather than in the decoder itself you may observe training discrepancies if finetuning a model trained with versions anterior to 4 17 0 the decoderinputids are now created based on the labels no need to pass them yourself anymore encoderdecoderstartdocstring r this class can be used to initialize a sequencetosequence model with any pretrained autoencoding model as the encoder and any pretrained autoregressive model as the decoder the encoder is loaded via tfautomodel frompretrained function and the decoder is loaded via tfautomodelforcausallm frompretrained function crossattention layers are automatically added to the decoder and should be finetuned on a downstream generative task like summarization the effectiveness of initializing sequencetosequence models with pretrained checkpoints for sequence generation tasks was shown in leveraging pretrained checkpoints for sequence generation taskshttps arxiv orgabs1907 12461 by sascha rothe shashi narayan aliaksei severyn michael matena yanqi zhou wei li peter j liu after such an encoder decoder model has been trainedfinetuned it can be savedloaded just like any other models see the examples for more information this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tf keras modelhttps www tensorflow orgapidocspythontfkerasmodel subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior parameters config encoderdecoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the tfpretrainedmodel frompretrained method to load the model weights def shifttokensrightinputids tf tensor padtokenid int decoderstarttokenid int if padtokenid is none raise valueerrormake sure to set the padtokenid attribute of the model s configuration padtokenid tf castpadtokenid inputids dtype if decoderstarttokenid is none raise valueerrormake sure to set the decoderstarttokenid attribute of the model s configuration decoderstarttokenid tf castdecoderstarttokenid inputids dtype starttokens tf fillshapelistinputids0 1 decoderstarttokenid shiftedinputids tf concatstarttokens inputids 1 1 replace possible 100 values in labels by padtokenid shiftedinputids tf where shiftedinputids 100 tf fillshapelistshiftedinputids padtokenid shiftedinputids verify that labels has only positive values and 100 assertgte0 tf debugging assertgreaterequalshiftedinputids tf constant0 dtypeinputids dtype make sure the assertion op is called by wrapping the result in an identity noop with tf controldependenciesassertgte0 shiftedinputids tf identityshiftedinputids return shiftedinputids addstartdocstringsencoderdecoderstartdocstring class tfencoderdecodermodeltfpretrainedmodel tfcausallanguagemodelingloss r tfencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with one of the base model classes of the library as encoder and another one as decoder when created with the tfautomodel frompretrained class method for the encoder and tfautomodelforcausallm frompretrained class method for the decoder initialize with config make sure that the individual model s config refers to the shared config so that the updates to the config will be synced encoder outputs might need to be projected to different dimension for decoder matt the tf and pt weights don t align because our tf base classes have an extra layer compared to pt models the main model stem is in the mainlayer class if we remove that layer then weight names sync up as normal however the name of that extra layer is the name of the mainlayer in the base model we make the assumption here that the config modeltype is the same as the name of the mainlayer i don t know of anywhere that s not the case and i wasn t sure how else to go from the config to the correct mainlayer name this override is only needed in the case where we re crossloading weights from pt however since weights are often safetensors now we don t know if we re going to be crossloading until we sniff the weights file therefore we specify tftoptweightrename anyway and let the super method figure out if it needs it or not remove encoder decoder kwargs from kwargs load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag isdecoder that we need to set correctly make sure these 2 tf keras model have fixed names so frompretrained could load model weights correctly instantiate config with corresponding kwargs let the user be responsible for the expected format add arguments to encoder from kwargsencoder handle the case where the inputs are passed as a single dict which contains labels the labels shouldn t be passed to self encoder below because it is a based model without this parameter otherwise an error occurs when inputprocessing is called inside self encoder call handle the init case where dummyinputs returns a dict containing decoderinputids handle the init case where dummyinputs returns a dict containing decoderinputids optionally project encoderhiddenstates add arguments to decoder from kwargsdecoder compute loss independent from decoder as some shift the logits inside them the starting index of the remaining elements in decoderoutputs todo joao the tfbasemodeloutput wrapper should not be needed after the generate refactor is complete apply decoder cache reordering here coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license classes to support tf encoder decoder architectures this class can be used to initialize a sequence to sequence model with any pretrained autoencoding model as the encoder and any pretrained autoregressive model as the decoder the encoder is loaded via tfautomodel from_pretrained function and the decoder is loaded via tfautomodelforcausallm from_pretrained function cross attention layers are automatically added to the decoder and should be fine tuned on a downstream generative task like summarization the effectiveness of initializing sequence to sequence models with pretrained checkpoints for sequence generation tasks was shown in leveraging pre trained checkpoints for sequence generation tasks https arxiv org abs 1907 12461 by sascha rothe shashi narayan aliaksei severyn michael matena yanqi zhou wei li peter j liu after such an encoder decoder model has been trained fine tuned it can be saved loaded just like any other models see the examples for more information this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tf keras model https www tensorflow org api_docs python tf keras model subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior parameters config encoderdecoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the tfpretrainedmodel from_pretrained method to load the model weights args 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 indices of input sequence tokens in the vocabulary indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask np ndarray or tf tensor of shape 0 optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask decoder_input_ids np ndarray or tf tensor of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values provide for sequence to sequence training to the decoder indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details decoder_attention_mask np ndarray or tf tensor of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default encoder_outputs tuple tuple tf tensor optional this tuple must consist of last_hidden_state optional hidden_states optional attentions last_hidden_state tf tensor of shape 0 hidden_size is a tensor of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder 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 contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape 0 inputs_embeds np ndarray or tf tensor of shape 0 hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix decoder_inputs_embeds np ndarray or tf tensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix labels np ndarray or tf tensor of shape 0 optional labels for computing the masked language modeling loss for the decoder indices should be in 100 0 config vocab_size see input_ids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocab_size use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional if set to true the model will return a utils seq2seqlmoutput instead of a plain tuple training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation kwargs optional remaining dictionary of keyword arguments keyword arguments come in two flavors without a prefix which will be input as encoder_kwargs for the encoder forward function with a decoder_ prefix which will be input as decoder_kwargs for the decoder forward function replace possible 100 values in labels by pad_token_id verify that labels has only positive values and 100 make sure the assertion op is called by wrapping the result in an identity no op tfencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with one of the base model classes of the library as encoder and another one as decoder when created with the tfautomodel from_pretrained class method for the encoder and tfautomodelforcausallm from_pretrained class method for the decoder initialize with config make sure that the individual model s config refers to the shared config so that the updates to the config will be synced encoder outputs might need to be projected to different dimension for decoder example python from transformers import tfencoderdecodermodel model tfencoderdecodermodel from_pretrained ydshieh bert2bert cnn_dailymail fp16 matt the tf and pt weights don t align because our tf base classes have an extra layer compared to pt models the main model stem is in the mainlayer class if we remove that layer then weight names sync up as normal however the name of that extra layer is the name of the mainlayer in the base model we make the assumption here that the config model_type is the same as the name of the mainlayer i don t know of anywhere that s not the case and i wasn t sure how else to go from the config to the correct mainlayer name this override is only needed in the case where we re crossloading weights from pt however since weights are often safetensors now we don t know if we re going to be crossloading until we sniff the weights file therefore we specify tf_to_pt_weight_rename anyway and let the super method figure out if it needs it or not instantiate an encoder and a decoder from one or two base classes of the library from pretrained model checkpoints params encoder_pretrained_model_name_or_path str optional information necessary to initiate the encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using tfpretrainedmodel save_pretrained e g my_model_directory a path or url to a pytorch index checkpoint file e g pt_model in this case encoder_from_pt should be set to true decoder_pretrained_model_name_or_path str optional defaults to none information necessary to initiate the decoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using tfpretrainedmodel save_pretrained e g my_model_directory a path or url to a pytorch checkpoint file e g pt_model in this case decoder_from_pt should be set to true model_args remaining positional arguments optional all remaning positional arguments will be passed to the underlying model s __init__ method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g output_attentions true to update the encoder configuration use the prefix encoder_ for each configuration parameter to update the decoder configuration use the prefix decoder_ for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import tfencoderdecodermodel initialize a bert2gpt2 from two pretrained bert models note that the cross attention layers will be randomly initialized model tfencoderdecodermodel from_encoder_decoder_pretrained bert base uncased gpt2 saving model after fine tuning model save_pretrained bert2gpt2 load fine tuned model model tfencoderdecodermodel from_pretrained bert2gpt2 remove encoder decoder kwargs from kwargs load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag is_decoder that we need to set correctly make sure these 2 tf keras model have fixed names so from_pretrained could load model weights correctly instantiate config with corresponding kwargs returns examples python from transformers import tfencoderdecodermodel berttokenizer initialize a bert2gpt2 from a pretrained bert and gpt2 models note that the cross attention layers will be randomly initialized model tfencoderdecodermodel from_encoder_decoder_pretrained bert base cased gpt2 tokenizer berttokenizer from_pretrained bert base cased forward input_ids tokenizer encode hello my dog is cute add_special_tokens true return_tensors tf batch size 1 outputs model input_ids input_ids decoder_input_ids input_ids training outputs model input_ids input_ids decoder_input_ids input_ids labels input_ids loss logits outputs loss outputs logits save and load from pretrained model save_pretrained bert2gpt2 model tfencoderdecodermodel from_pretrained bert2gpt2 generation generated model generate input_ids decoder_start_token_id model config decoder bos_token_id let the user be responsible for the expected format add arguments to encoder from kwargs_encoder handle the case where the inputs are passed as a single dict which contains labels the labels shouldn t be passed to self encoder below because it is a based model without this parameter otherwise an error occurs when input_processing is called inside self encoder call handle the init case where dummy_inputs returns a dict containing decoder_input_ids handle the init case where dummy_inputs returns a dict containing decoder_input_ids optionally project encoder_hidden_states add arguments to decoder from kwargs_decoder compute loss independent from decoder as some shift the logits inside them the starting index of the remaining elements in decoder_outputs e g on tf gpt2 needs to be passed to make keras layer __call__ happy todo joao the tfbasemodeloutput wrapper should not be needed after the generate refactor is complete apply decoder cache reordering here
from __future__ import annotations

import inspect
import re
import warnings
from typing import Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ...configuration_utils import PretrainedConfig
from ...modeling_tf_outputs import TFBaseModelOutput, TFSeq2SeqLMOutput
from ...modeling_tf_utils import (
    TFCausalLanguageModelingLoss,
    TFModelInputType,
    TFPreTrainedModel,
    get_initializer,
    unpack_inputs,
)
from ...tf_utils import shape_list
from ...utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ..auto.configuration_auto import AutoConfig
from ..auto.modeling_tf_auto import TFAutoModel, TFAutoModelForCausalLM
from .configuration_encoder_decoder import EncoderDecoderConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "EncoderDecoderConfig"

DEPRECATION_WARNING = (
    "Version v4.17.0 introduces a better way to train encoder-decoder models by computing the loss inside the"
    " encoder-decoder framework rather than in the decoder itself. You may observe training discrepancies if"
    " fine-tuning a model trained with versions anterior to 4.17.0. The decoder_input_ids are now created based on the"
    " labels, no need to pass them yourself anymore."
)

ENCODER_DECODER_START_DOCSTRING = r

ENCODER_DECODER_INPUTS_DOCSTRING = r


def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):
    if pad_token_id is None:
        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
    pad_token_id = tf.cast(pad_token_id, input_ids.dtype)

    if decoder_start_token_id is None:
        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
    decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)

    start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
    shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)
    
    shifted_input_ids = tf.where(
        shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids
    )

    
    assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))

    
    with tf.control_dependencies([assert_gte0]):
        shifted_input_ids = tf.identity(shifted_input_ids)

    return shifted_input_ids


@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)
class TFEncoderDecoderModel(TFPreTrainedModel, TFCausalLanguageModelingLoss):
    r

    config_class = EncoderDecoderConfig
    base_model_prefix = "encoder_decoder"
    load_weight_prefix = "tf_encoder_decoder_model"

    def __init__(
        self,
        config: Optional[PretrainedConfig] = None,
        encoder: Optional[TFPreTrainedModel] = None,
        decoder: Optional[TFPreTrainedModel] = None,
    ):
        if config is None and (encoder is None or decoder is None):
            raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
        if config is None:
            config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
        else:
            if not isinstance(config, self.config_class):
                raise ValueError(f"config: {config} has to be of type {self.config_class}")

        if config.decoder.cross_attention_hidden_size is not None:
            if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
                raise ValueError(
                    "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
                    f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
                    f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
                    " `config.encoder.hidden_size`."
                )

        
        super().__init__(config)

        if encoder is None:
            encoder = TFAutoModel.from_config(config.encoder, name="encoder")

        if decoder is None:
            decoder = TFAutoModelForCausalLM.from_config(config.decoder, name="decoder")

        self.encoder = encoder
        self.decoder = decoder

        if self.encoder.config.to_dict() != self.config.encoder.to_dict():
            logger.warning(
                f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
                f" {self.config.encoder}"
            )
        if self.decoder.config.to_dict() != self.config.decoder.to_dict():
            logger.warning(
                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
                f" {self.config.decoder}"
            )

        
        
        self.encoder.config = self.config.encoder
        self.decoder.config = self.config.decoder

        
        if (
            self.encoder.config.hidden_size != self.decoder.config.hidden_size
            and self.decoder.config.cross_attention_hidden_size is None
        ):
            self.enc_to_dec_proj = tf.keras.layers.Dense(
                units=self.decoder.config.hidden_size,
                kernel_initializer=get_initializer(config.encoder.initializer_range),
                name="enc_to_dec_proj",
            )

        if self.encoder.get_output_embeddings() is not None:
            raise ValueError(
                f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
            )

        decoder_signature = set(inspect.signature(self.decoder.call).parameters.keys())
        if "encoder_hidden_states" not in decoder_signature:
            raise ValueError(
                "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
                "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
            )

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def get_input_embeddings(self):
        return self.encoder.get_input_embeddings()

    def get_output_embeddings(self):
        return self.decoder.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        return self.decoder.set_output_embeddings(new_embeddings)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        r
        
        
        
        
        

        
        
        
        

        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
        encoder_model_type = config.encoder.model_type

        def tf_to_pt_weight_rename(tf_weight):
            if "encoder" in tf_weight and "decoder" not in tf_weight:
                return re.sub(rf"encoder\.{encoder_model_type}\.", "encoder.", tf_weight)
            else:
                return tf_weight

        kwargs["tf_to_pt_weight_rename"] = tf_to_pt_weight_rename
        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)

    @classmethod
    def from_encoder_decoder_pretrained(
        cls,
        encoder_pretrained_model_name_or_path: str = None,
        decoder_pretrained_model_name_or_path: str = None,
        *model_args,
        **kwargs,
    ) -> TFPreTrainedModel:
        r

        kwargs_encoder = {
            argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
        }

        kwargs_decoder = {
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

        
        for key in kwargs_encoder.keys():
            del kwargs["encoder_" + key]
        for key in kwargs_decoder.keys():
            del kwargs["decoder_" + key]

        
        
        
        encoder = kwargs_encoder.pop("model", None)
        if encoder is None:
            if encoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_encoder:
                encoder_config = AutoConfig.from_pretrained(encoder_pretrained_model_name_or_path)
                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
                    logger.info(
                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
                        "from a decoder model. Cross-attention and casual mask are disabled."
                    )
                    encoder_config.is_decoder = False
                    encoder_config.add_cross_attention = False

                kwargs_encoder["config"] = encoder_config

            kwargs_encoder["name"] = "encoder"
            kwargs_encoder["load_weight_prefix"] = cls.load_weight_prefix
            encoder = TFAutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)

        decoder = kwargs_decoder.pop("model", None)
        if decoder is None:
            if decoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_decoder:
                decoder_config = AutoConfig.from_pretrained(decoder_pretrained_model_name_or_path)
                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
                    logger.info(
                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
                        f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
                        f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
                    )
                    decoder_config.is_decoder = True
                    decoder_config.add_cross_attention = True

                kwargs_decoder["config"] = decoder_config

            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
                logger.warning(
                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
                    "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
                    "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
                )

            kwargs_decoder["name"] = "decoder"
            kwargs_decoder["load_weight_prefix"] = cls.load_weight_prefix
            decoder = TFAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)

        
        if encoder.name != "encoder":
            raise ValueError("encoder model must be created with the name `encoder`.")
        if decoder.name != "decoder":
            raise ValueError("decoder model must be created with the name `decoder`.")

        
        config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
        return cls(encoder=encoder, decoder=decoder, config=config)

    @unpack_inputs
    @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
    @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        input_ids: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        decoder_input_ids: np.ndarray | tf.Tensor | None = None,
        decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
        encoder_outputs: np.ndarray | tf.Tensor | None = None,
        past_key_values: Tuple[Tuple[tf.Tensor]] | None = None,
        inputs_embeds: np.ndarray | tf.Tensor | None = None,
        decoder_inputs_embeds: np.ndarray | tf.Tensor | None = None,
        labels: np.ndarray | tf.Tensor | None = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs,
    ) -> Union[TFSeq2SeqLMOutput, Tuple[tf.Tensor]]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}

        kwargs_decoder = {
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

        
        if encoder_outputs is not None:
            if return_dict and not isinstance(encoder_outputs, ModelOutput):
                raise ValueError(
                    "If `return_dict=True` and `encoder_outputs` is provided, it should be an instance of "
                    f"`ModelOutput`. Got an instance {type(encoder_outputs)} for `encoder_outputs`."
                )

        if encoder_outputs is None:
            encoder_inputs = {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "inputs_embeds": inputs_embeds,
                "output_attentions": output_attentions,
                "output_hidden_states": output_hidden_states,
                "return_dict": return_dict,
                "training": training,
            }

            
            encoder_inputs.update(kwargs_encoder)

            
            
            
            if "labels" in encoder_inputs:
                labels = encoder_inputs.pop("labels")

            
            if "decoder_input_ids" in encoder_inputs:
                decoder_input_ids = encoder_inputs.pop("decoder_input_ids")
            
            if "decoder_attention_mask" in encoder_inputs:
                decoder_attention_mask = encoder_inputs.pop("decoder_attention_mask")

            encoder_outputs = self.encoder(**encoder_inputs)

        encoder_hidden_states = encoder_outputs[0]

        
        if (
            self.encoder.config.hidden_size != self.decoder.config.hidden_size
            and self.decoder.config.cross_attention_hidden_size is None
        ):
            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)

        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
            decoder_input_ids = shift_tokens_right(
                labels, self.config.pad_token_id, self.config.decoder_start_token_id
            )

        decoder_inputs = {
            "input_ids": decoder_input_ids,
            "attention_mask": decoder_attention_mask,
            "encoder_hidden_states": encoder_hidden_states,
            "encoder_attention_mask": attention_mask,
            "inputs_embeds": decoder_inputs_embeds,
            "output_attentions": output_attentions,
            "output_hidden_states": output_hidden_states,
            "use_cache": use_cache,
            "past_key_values": past_key_values,
            "return_dict": return_dict,
            "training": training,
        }

        
        decoder_inputs.update(kwargs_decoder)

        decoder_outputs = self.decoder(**decoder_inputs)

        logits = decoder_outputs[0]

        
        loss = None
        if labels is not None:
            warnings.warn(DEPRECATION_WARNING, FutureWarning)
            loss = self.hf_compute_loss(labels, logits)

        if not return_dict:
            past_key_values = None
            if use_cache:
                past_key_values = decoder_outputs[1]
            
            start_index = sum([1 if x is not None else 0 for x in (loss, logits, past_key_values)])

            if not isinstance(encoder_outputs, tuple):
                encoder_outputs = encoder_outputs.to_tuple()
            output = (loss, logits, past_key_values) + decoder_outputs[start_index:] + encoder_outputs
            output = tuple([x for x in output if x is not None])
            return output

        return TFSeq2SeqLMOutput(
            loss=loss,
            logits=decoder_outputs.logits,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs
    ):
        decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)
        decoder_attention_mask = decoder_inputs["attention_mask"] if "attention_mask" in decoder_inputs else None
        past_key_values = decoder_inputs.get("past_key_values")
        if past_key_values is None:
            past_key_values = decoder_inputs.get("past")  
        input_dict = {
            "input_ids": None,  
            "attention_mask": attention_mask,
            "decoder_attention_mask": decoder_attention_mask,
            "decoder_input_ids": decoder_inputs["input_ids"],
            
            "encoder_outputs": TFBaseModelOutput(last_hidden_state=encoder_outputs[0]),
            "past_key_values": past_key_values,
            "use_cache": use_cache,
        }
        return input_dict

    def prepare_decoder_input_ids_from_labels(self, labels: tf.Tensor):
        return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)

    def resize_token_embeddings(self, *args, **kwargs):
        raise NotImplementedError(
            "Resizing the embedding layers via the TFEncoderDecoderModel directly is not supported.Please use the"
            " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
            " model.decoder.resize_token_embeddings(...))"
        )

    def _reorder_cache(self, past, beam_idx):
        
        return self.decoder._reorder_cache(past, beam_idx)