U &c@sdZddlZddlZddlZddlZddlZddlmZddl m Z m Z m Z m Z ddlmZmZmZeeZddd d d d ZGd ddejjjZGdddejjjZGdddejjjZGdddejjjZGdddejjjZGdddejjjZGdddejjjZGddde Z!GdddejjjZ"Gdd d eZ#d!Z$d"Z%e d#e$Gd$d%d%e#Z&e d&e$Gd'd(d(e#Z'dS))z TF 2.0 T5 model. N)T5Config) DUMMY_INPUTS DUMMY_MASKadd_start_docstrings add_start_docstrings_to_callable)TFPreTrainedModelTFSharedEmbeddings shape_listz/https://cdn.huggingface.co/t5-small-tf_model.h5z.https://cdn.huggingface.co/t5-base-tf_model.h5z/https://cdn.huggingface.co/t5-large-tf_model.h5z,https://cdn.huggingface.co/t5-3b-tf_model.h5z-https://cdn.huggingface.co/t5-11b-tf_model.h5)zt5-smallzt5-basezt5-largezt5-3bzt5-11bcs2eZdZdfdd ZfddZddZZS) TFT5LayerNormư>c stjf|||_dS)zg Construct a layernorm module in the T5 style No bias and no substraction of mean. N)super__init__variance_epsilon)selfepsilonkwargs __class__?/tmp/pip-unpacked-wheel-ymerj3tt/transformers/modeling_tf_t5.pyr1szTFT5LayerNorm.__init__cs(|jd|dfdd|_t|dS)z"Build shared word embedding layer weightZones)shapeZ initializerN)Z add_weightrr build)r input_shaperrrr8szTFT5LayerNorm.buildcCs:tjjtj|ddd}|tj||j}|j|S)NrT)axisZkeepdims)tfmathZ reduce_meanZsquareZrsqrtrr)rxZvariancerrrcall=szTFT5LayerNorm.call)r )__name__ __module__ __qualname__rrr __classcell__rrrrr 0s r cs&eZdZfddZdddZZS)TFT5DenseReluDensec s`tjf|tjjj|jddd|_tjjj|jddd|_ tjj |j |_ tjj j|_dS)NFwiZuse_biasnamewo)r rrkeraslayersDenseZd_ffr&d_modelr)Dropout dropout_ratedropoutZ activationsZreluactrconfigrrrrrDs zTFT5DenseReluDense.__init__FcCs0||}||}|j||d}||}|SNtraining)r&r1r0r))r hidden_statesr6hrrrr Ks    zTFT5DenseReluDense.call)Fr!r"r#rr r$rrrrr%Cs r%cs&eZdZfddZdddZZS) TFT5LayerFFc sBtjf|t|dd|_t|jdd|_tjj |j |_ dS)NDenseReluDenser( layer_normrr() r rr%r;r layer_norm_epsilonr=rr*r+r.r/r0r2rrrrTszTFT5LayerFF.__init__FcCs.||}|j||d}||j||d}|Sr4)r=r;r0)rr7r6norm_xy layer_outputrrrr Zs zTFT5LayerFF.call)Fr9rrrrr:Ss r:c sNeZdZeZdfdd ZddZedd d Z d d Z dddZ Z S) TFT5AttentionFc stjf|ttj|_|j|_||_|j|_|j |_ |j |_ |j |_ |j |_ |j |j |_tjjj|jddd|_tjjj|jddd|_tjjj|jddd|_tjjj|j ddd|_tjj|j|_|jrtjjj|j |j dd|_t|_dS) NFqr'kvorelative_attention_biasr<)r rnextrCNEW_IDZlayer_id is_decoderhas_relative_attention_biasoutput_attentionsrelative_attention_num_bucketsr-d_kvZ num_headsn_heads inner_dimrr*r+r,rDrErFrGr.r/r0Z EmbeddingrHsetZ pruned_headsrr3rLrrrrrds, zTFT5Attention.__init__cCstdSNNotImplementedError)rZheadsrrr prune_heads~szTFT5Attention.prune_headsT c Csd}| }|rF|d}|tjtj|dtj|7}tj|}ntj|d}|d}tj||}|tjtjtj|tj |t||||tj}tj ||d}|t |||7}|S)a Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 Translate relative position to a bucket number for relative attention. The relative position is defined as memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should allow for more graceful generalization to longer sequences than the model has been trained on. Args: relative_position: an int32 Tensor bidirectional: a boolean - whether the attention is bidirectional num_buckets: an integer max_distance: an integer Returns: a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) rr) rZdtypescastrZlessZint32absmaximumlogfloat32Zminimumwhere) relative_position bidirectional num_bucketsZ max_distanceretnZ max_exactZis_smallZ val_if_largerrr_relative_position_buckets(" z'TFT5Attention._relative_position_bucketcCstt|dddf}t|dddf}||}|j||j |jd}||}tjt|dddgdd}|S)z' Compute binned relative position bias N)rbrcrZrrr)rrangerfrKrNrHZ expand_dims transpose)rqlenklenZcontext_positionZmemory_positionraZ rp_bucketvaluesrrr compute_biass zTFT5Attention.compute_biasNc st|\} } |dk rhjdks(tdt|dksFtdt||dkrb| t|ddn|} n| } |dkrz| }n t|d}fdd }fd d }||}|dkr||}||}n.|dkr|}}||}||}|dk rN|dkrF|\}}tj ||gdd }tj ||gdd }n|\}}t | rxt | d rtt | } nd} jr| dkr||ff}nd}td||}|dkr jstd| |}|dk r|ddddddddf}|dk r ||}||7}tjj|dd }j|| d}|dk rB||}t||}||}|}|f|}jr|||f}jr||f}|S)zd Self-attention (if kv is None) or attention over source sentence (provided by kv). NTz*Encoder cannot cache past key value statesrZzSpast_key_value_state should have 2 past states: keys and values. Got {} past statesrrcs"tjt|djjfddS)z projection rrrZrperm)rrireshaperPrOrbsrrrrsz!TFT5Attention.call..shapecsttj|dddjfS)z compute context rnrpr)rrrrirQrsrtrrunshapesz#TFT5Attention.call..unshapergnumpyrTzbnqd,bnkd->bnqkzANo position_bias provided and no weights to compute position_biasrr5)r rKAssertionErrorlenformatrDrErFrconcatZ is_tensorhasattrboolrwZeinsumrL ValueErrorrmnnZsoftmaxr0matmulrGrM)rinputmaskkv position_biascachepast_key_value_state head_mask query_length use_cacher6rjZdimZ real_qlenrkrrvrDrErFZk_Zv_present_key_value_stateZscoresweightscontextoutputsrrtrr sx "                  zTFT5Attention.call)F)TrXrY) NNNNNNNFF) r!r"r# itertoolscountrJrrW staticmethodrfrmr r$rrrrrCas  -rCcs(eZdZdfdd ZdddZZS) TFT5LayerSelfAttentionFc sDtjf|t||dd|_t|jdd|_tjj |j |_ dS)N SelfAttentionrLr(r=r>) r rrCrr r?r=rr*r+r.r/r0rSrrrr.szTFT5LayerSelfAttention.__init__Nc CsR||}|j|||||||d} | d} ||j| |d} | f| dd} | S)N)rrrrrr6rr5r)r=rr0) rr7attention_maskrrrrr6r@attention_outputrArBrrrrr 6s  zTFT5LayerSelfAttention.call)F)NNNNFFr9rrrrr-s rcs(eZdZdfdd ZdddZZS) TFT5LayerCrossAttentionFc sDtjf|t||dd|_t|jdd|_tjj |j |_ dS)NEncDecAttentionrr=r>) r rrCrr r?r=rr*r+r.r/r0rSrrrrQsz TFT5LayerCrossAttention.__init__Nc  CsV||} |j| |||||||| d } | d} ||j| | d} | f| dd}|S)N)rrrrrrrr6rr5r)r=rr0)rr7rrrrrrrr6r@rrArBrrrrr Ys  zTFT5LayerCrossAttention.call)F)NNNNNFFr9rrrrrPs rc s(eZdZdfdd ZdddZZS) TFT5BlockFc srtjf||j|_g|_|jt||dd|jrN|jt||dd|jt|dt |jddS)Nz layer_._0rz layer_._1z layer_._{}r<) r rrKlayerappendrrr:rzryrSrrrrxs zTFT5Block.__init__Nc  CsZ|dk rr|jstd|dkr"dnd} d| | dkr8dndt|} t|| ksXt| |dd} |dd}nd\} }|jd||||| | | d }|dd\}}|dd}|jr0|dk r0|dk rt|dd}nd}|jd |||||||| | d }|d}|dk r ||d }||dd}|jd || d }|f}||f|}|S)Nz,Only decoder can use `past_key_value_states`rZzcThere should be {} past states. 2 (past / key) for self attention.{} Got {} past key / value statesz"2 (past / key) for cross attention)NNr)rrrrrr6r)rrrrrrrr6rr5)rKrxrzryrr )rr7rrencoder_hidden_statesencoder_attention_maskencoder_decoder_position_biasrrrr6Z"expected_num_past_key_value_states error_messageZself_attn_past_key_value_stateZcross_attn_past_key_value_stateZself_attention_outputsrZattention_outputsrZcross_attention_outputsrrrrr s\      zTFT5Block.call)F) NNNNNNNFFr9rrrrrwsrc@s.eZdZdZd ddZd ddZd dd ZdS) _NoLayerEmbedTokensa this class wraps a the TFSharedEmbeddingTokens layer into a python 'no-keras-layer' class to avoid problem with weight restoring. Also it makes sure that the layer is called from the correct scope to avoid problem with saving/storing the correct weights NcCs||_||_dSrT)_layer_abs_scope_name)rrabs_scope_namerrrrsz_NoLayerEmbedTokens.__init__ embeddingc Csv|jdkr|j||Stjjj|jdd@}t|j(|j||W5QRW5QRSQRXW5QRXdSNF)Zauxiliary_name_scope) rrr rcompatv1variable_scope name_scopeoriginal_name_scoperinputsmoderrrrr s  z_NoLayerEmbedTokens.callc Csr|jdkr|||Stjjj|jdd>}t|j&|||W5QRW5QRSQRXW5QRXdSr)rrrrrrrrrrrr__call__s   z_NoLayerEmbedTokens.__call__)N)r)r)r!r"r#__doc__rr rrrrrrs  rcsPeZdZdfdd ZddZddZdd Zd d Zd d ZdddZ Z S) TFT5MainLayerNc sztjf|j|_j|_||_j|_|_j|_fddt jD|_ t j dd|_ tjjj|_dS)Ncs(g|] }tt|dkd|dqS)rz block_._{}r)rr}rz).0ir3rr sz*TFT5MainLayer.__init__..final_layer_normr>)r rrMoutput_hidden_states embed_tokensrKr3Z num_layersnum_hidden_layersrhblockr r?rrr*r+r.r/r0)rr3rrrrrrs zTFT5MainLayer.__init__cCs|jSrTrrrrrget_input_embeddings sz"TFT5MainLayer.get_input_embeddingscCs|jSrTrrrrrget_output_embeddingssz#TFT5MainLayer.get_output_embeddingscCs ||_dSrTr)rrrrrset_embed_tokensszTFT5MainLayer.set_embed_tokenscCstdSrTrU)rZnew_num_tokensrrr_resize_token_embeddingssz&TFT5MainLayer._resize_token_embeddingscCstdSrTrU)rZheads_to_prunerrr _prune_headsszTFT5MainLayer._prune_headsFc ! Csv|dk r|dk rtdnH|dk r@t|} t|d| df}n"|dk rZt|dd} ntd|dkr|jdk s|td||}| \} } |dk r| dkstd| | dft|ddd| } n| } |dkrt| | fd}|jr |dkr |dk r t|d}t| |fd}|dkr:dgt |j }tj |tj d }t t|}|d kr~|dddddddf}n|dkr\|jrDt | }tt|ddddf| | df|ddddf}tj |tj d }|dddddddf|ddddddf}|ddk r\|ddddddddf}n|ddddddf}d |d }|jr|dk rtj |tj d }t t|}|d kr|dddddddf}|dkr|ddddddf}d |d }nd}|dk rtn dg|j}d }d }d }d}d}|j|| d}tt|j |D]\}\}}|jr\||f}|||||||||||| d }|dd\}}|dkr||jrd nd}|jr|dk r||jrdnd }||f}|jr>||df}q>||}|j|| d}|jr||f}|f} |dkrN|jsDtd|| |f} |jr`| |f} |jrr| |f} | S)NzDYou cannot specify both input_ids and inputs_embeds at the same timerz5You have to specify either input_ids or inputs_embedsz;You have to intialize the model with valid token embeddingsrzCInput shape is {}, but should be {} when using past_key_value_satesrrZ)Zdtyperog?gerr5) rrrrrrrrr6rTz@`use_cache` can only be set to `True` if {} is used as a decoder)r~r rrrrrxrzfillrKryrr[r_rhZ less_equalZtilerVrr0 enumerateziprrMr)!r input_idsrrr inputs_embedsrpast_key_value_statesrr6rZ batch_sizeZ seq_lengthZmask_seq_lengthZencoder_seq_lengthZnum_dims_attention_maskZextended_attention_maskZseq_idsZ causal_maskZnum_dims_encoder_attention_maskZencoder_extended_attention_maskZpresent_key_value_statesZall_hidden_statesZall_attentionsrrr7rZ layer_modulerZ layer_outputsrrrrrr s         4"               zTFT5MainLayer.call)N)NNNNNNFF) r!r"r#rrrrrrr r$rrrrrsrc@s(eZdZdZeZeZdZe ddZ dS)TFT5PreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. Z transformercCs$tt}tt}|||d}|S)N)rdecoder_input_idsdecoder_attention_mask)rZconstantrr)rrZ input_mask dummy_inputsrrrrs  z TFT5PreTrainedModel.dummy_inputsN) r!r"r#rrZ config_class"TF_T5_PRETRAINED_MODEL_ARCHIVE_MAPZpretrained_model_archive_mapZbase_model_prefixpropertyrrrrrrs ra The T5 model was proposed in `Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer`_ by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. It's an encoder decoder transformer pre-trained in a text-to-text denoising generative setting. This model is a tf.keras.Model `tf.keras.Model`_ sub-class. 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. .. _`Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer`: https://arxiv.org/abs/1910.10683 .. _`tf.keras.Model`: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Model Note on the model inputs: TF 2.0 models accepts two formats as inputs: - having all inputs as keyword arguments (like PyTorch models), or - having all inputs as a list, tuple or dict in the first positional arguments. This second option is usefull when using `tf.keras.Model.fit()` method which currently requires having all the tensors in the first argument of the model call function: `model(inputs)`. If you choose this second option, there are three possibilities you can use to gather all the input Tensors in the first positional argument : - a single Tensor with input_ids only and nothing else: `model(inputs_ids) - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])` - a dictionary with one or several input Tensors associaed to the input names given in the docstring: `model({'input_ids': input_ids, 'token_type_ids': token_type_ids})` Parameters: config (:class:`~transformers.T5Config`): 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 :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights. a Args: inputs are usually used as a `dict` (see T5 description above for more information) containing all the following. inputs (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on the right or the left. Indices can be obtained using :class:`transformers.T5Tokenizer`. To know more on how to prepare :obj:`input_ids` for pre-training take a look at `T5 Training <./t5.html#training>`_ . See :func:`transformers.PreTrainedTokenizer.encode` and :func:`transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details. decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`, defaults to :obj:`None`): Provide for sequence to sequence training. T5 uses the pad_token_id as the starting token for decoder_input_ids generation. If `decoder_past_key_value_states` is used, optionally only the last `decoder_input_ids` have to be input (see `decoder_past_key_value_states`). attention_mask (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`): Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens. encoder_outputs (:obj:`tuple(tuple(tf.FloatTensor)`, `optional`, defaults to :obj:`None`): Tuple consists of (`last_hidden_state`, `optional`: `hidden_states`, `optional`: `attentions`) `last_hidden_state` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`) is a sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(batch_size, tgt_seq_len)`, `optional`, defaults to :obj:`None`): Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also be used by default. decoder_past_key_value_states (:obj:`tuple(tuple(tf.Tensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains pre-computed key and value hidden-states of the attention blocks. Can be used to speed up decoding. If `decoder_past_key_value_states` 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 :obj:`(batch_size, 1)` use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): If `use_cache` is True, `decoder_past_key_value_states` are returned and can be used to speed up decoding (see `decoder_past_key_value_states`). inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`): Optionally, instead of passing :obj:`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 (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`): Optionally, instead of passing :obj:`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. To know more on how to prepare :obj:`decoder_input_ids` for pre-training take a look at `T5 Training <./t5.html#training>`_ . head_mask: (:obj:`tf.Tensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`): Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**. z[The bare T5 Model transformer outputting raw hidden-stateswithout any specific head on top.csLeZdZfddZddZddZddZd d Zee d d Z Z S) TFT5Modelc stj|f||t|j|jdd|_tjj d}W5QRXt |j|d}t |}t ||dd|_t |}d|_t ||dd|_dSNsharedr<)rencoderTdecoder)r rr vocab_sizer-rrrrrrcopydeepcopyrrrKrrr3rrZshared_abs_scope_namerZencoder_configZdecoder_configrrrr=s   zTFT5Model.__init__cCs|jSrTrrrrrrNszTFT5Model.get_input_embeddingscCs|jSrTrrrrrrQszTFT5Model.get_output_embeddingscCs|jSrTrrrrr get_encoderTszTFT5Model.get_encodercCs|jSrTrrrrr get_decoderWszTFT5Model.get_decoderc KsNt|tr||n||d<|dd}|dd}|dd}|dd}|dd}|dd}|dd} |d d} |d d } |d d} |dkr|j|||| d }|d} | dk r|dk r|ddddf}| dk r| ddddf} |j||| | | || | d}| d krF||dff}|dd||dd}||S)aT Return: :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.T5Config`) and inputs. last_hidden_state (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. If `decoder_past_key_value_states` is used only the last hidden-state of the sequences of shape :obj:`(batch_size, 1, hidden_size)` is output. decoder_past_key_value_states (:obj:`tuple(tuple(tf.Tensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length, embed_size_per_head)`, `optional`, returned when ``use_cache=True``): Contains pre-computed key and value hidden-states of the attention blocks. Can be used to speed up sequential decoding (see `decoder_past_key_value_states` input). Note that when using `decoder_past_key_value_states`, the model only outputs the last `hidden-state` of the sequence of shape :obj:`(batch_size, 1, config.vocab_size)`. hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``): Tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape :obj:`(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``): Tuple of :obj:`tf.Tensor` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Examples:: from transformers import T5Tokenizer, TFT5Model tokenizer = T5Tokenizer.from_pretrained('t5-small') model = TFT5Model.from_pretrained('t5-small') input_ids = tokenizer.encode("Hello, my dog is cute", return_tensors="tf") # Batch size 1 outputs = model(input_ids, decoder_input_ids=input_ids) last_hidden_states = outputs[0] # The last hidden-state is the first element of the output tuple rNrrencoder_outputsrrdecoder_inputs_embedsdecoder_past_key_value_statesrTrrrrrrrrrrrrrrrZ) isinstancedictupdategetrr)rrrrrrrrrrrrrr7decoder_outputspastrrrr ZsP$             zTFT5Model.call) r!r"r#rrrrrrT5_INPUTS_DOCSTRINGr r$rrrrr8s rz1T5 Model with a `language modeling` head on top. cs\eZdZfddZddZddZddZd d Zee d d Z d dZ ddZ Z S)TFT5ForConditionalGenerationc stj|f|||j|_t|j|jdd|_tjj d}W5QRXt |j|d}t |}t||dd|_t |}d|_t||dd|_dSr)r rr- model_dimr rrrrrrrrrrrrKrrrrrrs   z%TFT5ForConditionalGeneration.__init__cCs|jSrTrrrrrrsz1TFT5ForConditionalGeneration.get_input_embeddingscCs|jSrTrrrrrrsz2TFT5ForConditionalGeneration.get_output_embeddingscCs|jSrTrrrrrrsz(TFT5ForConditionalGeneration.get_encodercCs|jSrTrrrrrrsz(TFT5ForConditionalGeneration.get_decoderc Kst|tr||n||d<|dd}|dd}|dd}|dd}|dd}|dd}|dd } |d d} |d d} |d d} |dkr|j||| | d }|d} |dk r|dk r|ddddf}| dk r| ddddf} |j||| || || | d}| d krF||dff}|dd||dd}|d|jd}|}||dd}|f|dd}||S)a^ Return: :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.T5Config`) and inputs. loss (:obj:`tf.Tensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`lm_label` is provided): Classification loss (cross entropy). prediction_scores (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`) Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). decoder_past_key_value_states (:obj:`tuple(tuple(tf.Tensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length, embed_size_per_head)`, `optional`, returned when ``use_cache=True``): Contains pre-computed key and value hidden-states of the attention blocks. Can be used to speed up sequential decoding (see `decoder_past_key_value_states` input). Note that when using `decoder_past_key_value_states`, the model only outputs the last `prediction_score` of the sequence of shape :obj:`(batch_size, 1, config.vocab_size)`. hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``): Tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape :obj:`(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``): Tuple of :obj:`tf.Tensor` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention. Examples:: from transformers import T5Tokenizer, TFT5ForConditionalGeneration tokenizer = T5Tokenizer.from_pretrained('t5-small') model = TFT5ForConditionalGeneration.from_pretrained('t5-small') input_ids = tokenizer.encode("Hello, my dog is cute", return_tensors="tf") # Batch size 1 outputs = model(input_ids, decoder_input_ids=input_ids) prediction_scores = outputs[0] tokenizer = T5Tokenizer.from_pretrained('t5-small') model = TFT5ForConditionalGeneration.from_pretrained('t5-small') input_ids = tokenizer.encode("summarize: Hello, my dog is cute", return_tensors="tf") # Batch size 1 model.generate(input_ids) rNrrrrrrTrrrrrrrrrZgZlinear)r)rrrrrrrr)rrrrrrrrrrrrrr7rrZsequence_outputrZ lm_logitsrrrr sX)              z!TFT5ForConditionalGeneration.callcKsL|dk stdt|dkr(|d}}n|d|d}}d|||||dS)Nz*past has to be defined for encoder_outputsrZrr)rrrrrr)rxry)rrrrrrrrrrrprepare_inputs_for_generation7s  z:TFT5ForConditionalGeneration.prepare_inputs_for_generationcCst|dkrtd|S|d}|df}d}|D]^}d}|D]}|t||f}q@t|dt|dksttt|t|kst||f}q4||fS)NrZzHYou might want to consider setting `use_cache=True` to speed up decodingrrr)ryloggerwarningrZgatherr rx)rrZbeam_idxZ decoder_pastZreordered_decoder_pastZlayer_past_statesZreordered_layer_past_statesZlayer_past_staterrr_reorder_cacheIs    z+TFT5ForConditionalGeneration._reorder_cache)r!r"r#rrrrrrrr rrr$rrrrrs  cr)(rrrloggingrZ tensorflowrZconfiguration_t5rZ file_utilsrrrrZmodeling_tf_utilsrr r getLoggerr!rrr*r+ZLayerr r%r:rCrrrobjectrrrZT5_START_DOCSTRINGrrrrrrrsH  M#'_"V$2v