Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""TF 2.0 BERT model."""17 18from __future__ import annotations19 20import math21import warnings22from dataclasses import dataclass23 24import numpy as np25import tensorflow as tf26 27from ...activations_tf import get_tf_activation28from ...modeling_tf_outputs import (29 TFBaseModelOutputWithPastAndCrossAttentions,30 TFBaseModelOutputWithPoolingAndCrossAttentions,31 TFCausalLMOutputWithCrossAttentions,32 TFMaskedLMOutput,33 TFMultipleChoiceModelOutput,34 TFNextSentencePredictorOutput,35 TFQuestionAnsweringModelOutput,36 TFSequenceClassifierOutput,37 TFTokenClassifierOutput,38)39from ...modeling_tf_utils import (40 TFCausalLanguageModelingLoss,41 TFMaskedLanguageModelingLoss,42 TFModelInputType,43 TFMultipleChoiceLoss,44 TFNextSentencePredictionLoss,45 TFPreTrainedModel,46 TFQuestionAnsweringLoss,47 TFSequenceClassificationLoss,48 TFTokenClassificationLoss,49 get_initializer,50 keras,51 keras_serializable,52 unpack_inputs,53)54from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax55from ...utils import (56 ModelOutput,57 add_code_sample_docstrings,58 add_start_docstrings,59 add_start_docstrings_to_model_forward,60 logging,61 replace_return_docstrings,62)63from .configuration_bert import BertConfig64 65 66logger = logging.get_logger(__name__)67 68_CHECKPOINT_FOR_DOC = "google-bert/bert-base-uncased"69_CONFIG_FOR_DOC = "BertConfig"70 71# TokenClassification docstring72_CHECKPOINT_FOR_TOKEN_CLASSIFICATION = "dbmdz/bert-large-cased-finetuned-conll03-english"73_TOKEN_CLASS_EXPECTED_OUTPUT = (74 "['O', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 'O', 'O', 'O', 'O', 'I-LOC', 'O', 'I-LOC', 'I-LOC'] "75)76_TOKEN_CLASS_EXPECTED_LOSS = 0.0177 78# QuestionAnswering docstring79_CHECKPOINT_FOR_QA = "ydshieh/bert-base-cased-squad2"80_QA_EXPECTED_OUTPUT = "'a nice puppet'"81_QA_EXPECTED_LOSS = 7.4182_QA_TARGET_START_INDEX = 1483_QA_TARGET_END_INDEX = 1584 85# SequenceClassification docstring86_CHECKPOINT_FOR_SEQUENCE_CLASSIFICATION = "ydshieh/bert-base-uncased-yelp-polarity"87_SEQ_CLASS_EXPECTED_OUTPUT = "'LABEL_1'"88_SEQ_CLASS_EXPECTED_LOSS = 0.0189 90 91class TFBertPreTrainingLoss:92 """93 Loss function suitable for BERT-like pretraining, that is, the task of pretraining a language model by combining94 NSP + MLM. .. note:: Any label of -100 will be ignored (along with the corresponding logits) in the loss95 computation.96 """97 98 def hf_compute_loss(self, labels: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:99 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)100 101 # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway102 unmasked_lm_losses = loss_fn(y_true=tf.nn.relu(labels["labels"]), y_pred=logits[0])103 # make sure only labels that are not equal to -100104 # are taken into account for the loss computation105 lm_loss_mask = tf.cast(labels["labels"] != -100, dtype=unmasked_lm_losses.dtype)106 masked_lm_losses = unmasked_lm_losses * lm_loss_mask107 reduced_masked_lm_loss = tf.reduce_sum(masked_lm_losses) / tf.reduce_sum(lm_loss_mask)108 109 # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway110 unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels["next_sentence_label"]), y_pred=logits[1])111 ns_loss_mask = tf.cast(labels["next_sentence_label"] != -100, dtype=unmasked_ns_loss.dtype)112 masked_ns_loss = unmasked_ns_loss * ns_loss_mask113 114 reduced_masked_ns_loss = tf.reduce_sum(masked_ns_loss) / tf.reduce_sum(ns_loss_mask)115 116 return tf.reshape(reduced_masked_lm_loss + reduced_masked_ns_loss, (1,))117 118 119class TFBertEmbeddings(keras.layers.Layer):120 """Construct the embeddings from word, position and token_type embeddings."""121 122 def __init__(self, config: BertConfig, **kwargs):123 super().__init__(**kwargs)124 125 self.config = config126 self.hidden_size = config.hidden_size127 self.max_position_embeddings = config.max_position_embeddings128 self.initializer_range = config.initializer_range129 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")130 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)131 132 def build(self, input_shape=None):133 with tf.name_scope("word_embeddings"):134 self.weight = self.add_weight(135 name="weight",136 shape=[self.config.vocab_size, self.hidden_size],137 initializer=get_initializer(self.initializer_range),138 )139 140 with tf.name_scope("token_type_embeddings"):141 self.token_type_embeddings = self.add_weight(142 name="embeddings",143 shape=[self.config.type_vocab_size, self.hidden_size],144 initializer=get_initializer(self.initializer_range),145 )146 147 with tf.name_scope("position_embeddings"):148 self.position_embeddings = self.add_weight(149 name="embeddings",150 shape=[self.max_position_embeddings, self.hidden_size],151 initializer=get_initializer(self.initializer_range),152 )153 154 if self.built:155 return156 self.built = True157 if getattr(self, "LayerNorm", None) is not None:158 with tf.name_scope(self.LayerNorm.name):159 self.LayerNorm.build([None, None, self.config.hidden_size])160 161 def call(162 self,163 input_ids: tf.Tensor | None = None,164 position_ids: tf.Tensor | None = None,165 token_type_ids: tf.Tensor | None = None,166 inputs_embeds: tf.Tensor | None = None,167 past_key_values_length=0,168 training: bool = False,169 ) -> tf.Tensor:170 """171 Applies embedding based on inputs tensor.172 173 Returns:174 final_embeddings (`tf.Tensor`): output embedding tensor.175 """176 if input_ids is None and inputs_embeds is None:177 raise ValueError("Need to provide either `input_ids` or `input_embeds`.")178 179 if input_ids is not None:180 check_embeddings_within_bounds(input_ids, self.config.vocab_size)181 inputs_embeds = tf.gather(params=self.weight, indices=input_ids)182 183 input_shape = shape_list(inputs_embeds)[:-1]184 185 if token_type_ids is None:186 token_type_ids = tf.fill(dims=input_shape, value=0)187 188 if position_ids is None:189 position_ids = tf.expand_dims(190 tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0191 )192 193 position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)194 token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)195 final_embeddings = inputs_embeds + position_embeds + token_type_embeds196 final_embeddings = self.LayerNorm(inputs=final_embeddings)197 final_embeddings = self.dropout(inputs=final_embeddings, training=training)198 199 return final_embeddings200 201 202class TFBertSelfAttention(keras.layers.Layer):203 def __init__(self, config: BertConfig, **kwargs):204 super().__init__(**kwargs)205 206 if config.hidden_size % config.num_attention_heads != 0:207 raise ValueError(208 f"The hidden size ({config.hidden_size}) is not a multiple of the number "209 f"of attention heads ({config.num_attention_heads})"210 )211 212 self.num_attention_heads = config.num_attention_heads213 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)214 self.all_head_size = self.num_attention_heads * self.attention_head_size215 self.sqrt_att_head_size = math.sqrt(self.attention_head_size)216 217 self.query = keras.layers.Dense(218 units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"219 )220 self.key = keras.layers.Dense(221 units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"222 )223 self.value = keras.layers.Dense(224 units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"225 )226 self.dropout = keras.layers.Dropout(rate=config.attention_probs_dropout_prob)227 228 self.is_decoder = config.is_decoder229 self.config = config230 231 def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:232 # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]233 tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))234 235 # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]236 return tf.transpose(tensor, perm=[0, 2, 1, 3])237 238 def call(239 self,240 hidden_states: tf.Tensor,241 attention_mask: tf.Tensor,242 head_mask: tf.Tensor,243 encoder_hidden_states: tf.Tensor,244 encoder_attention_mask: tf.Tensor,245 past_key_value: tuple[tf.Tensor],246 output_attentions: bool,247 training: bool = False,248 ) -> tuple[tf.Tensor]:249 batch_size = shape_list(hidden_states)[0]250 mixed_query_layer = self.query(inputs=hidden_states)251 252 # If this is instantiated as a cross-attention module, the keys253 # and values come from an encoder; the attention mask needs to be254 # such that the encoder's padding tokens are not attended to.255 is_cross_attention = encoder_hidden_states is not None256 257 if is_cross_attention and past_key_value is not None:258 # reuse k,v, cross_attentions259 key_layer = past_key_value[0]260 value_layer = past_key_value[1]261 attention_mask = encoder_attention_mask262 elif is_cross_attention:263 key_layer = self.transpose_for_scores(self.key(inputs=encoder_hidden_states), batch_size)264 value_layer = self.transpose_for_scores(self.value(inputs=encoder_hidden_states), batch_size)265 attention_mask = encoder_attention_mask266 elif past_key_value is not None:267 key_layer = self.transpose_for_scores(self.key(inputs=hidden_states), batch_size)268 value_layer = self.transpose_for_scores(self.value(inputs=hidden_states), batch_size)269 key_layer = tf.concat([past_key_value[0], key_layer], axis=2)270 value_layer = tf.concat([past_key_value[1], value_layer], axis=2)271 else:272 key_layer = self.transpose_for_scores(self.key(inputs=hidden_states), batch_size)273 value_layer = self.transpose_for_scores(self.value(inputs=hidden_states), batch_size)274 275 query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)276 277 if self.is_decoder:278 # if cross_attention save Tuple(tf.Tensor, tf.Tensor) of all cross attention key/value_states.279 # Further calls to cross_attention layer can then reuse all cross-attention280 # key/value_states (first "if" case)281 # if uni-directional self-attention (decoder) save Tuple(tf.Tensor, tf.Tensor) of282 # all previous decoder key/value_states. Further calls to uni-directional self-attention283 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)284 # if encoder bi-directional self-attention `past_key_value` is always `None`285 past_key_value = (key_layer, value_layer)286 287 # Take the dot product between "query" and "key" to get the raw attention scores.288 # (batch size, num_heads, seq_len_q, seq_len_k)289 attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)290 dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)291 attention_scores = tf.divide(attention_scores, dk)292 293 if attention_mask is not None:294 # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)295 attention_scores = tf.add(attention_scores, attention_mask)296 297 # Normalize the attention scores to probabilities.298 attention_probs = stable_softmax(logits=attention_scores, axis=-1)299 300 # This is actually dropping out entire tokens to attend to, which might301 # seem a bit unusual, but is taken from the original Transformer paper.302 attention_probs = self.dropout(inputs=attention_probs, training=training)303 304 # Mask heads if we want to305 if head_mask is not None:306 attention_probs = tf.multiply(attention_probs, head_mask)307 308 attention_output = tf.matmul(attention_probs, value_layer)309 attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])310 311 # (batch_size, seq_len_q, all_head_size)312 attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))313 outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)314 315 if self.is_decoder:316 outputs = outputs + (past_key_value,)317 return outputs318 319 def build(self, input_shape=None):320 if self.built:321 return322 self.built = True323 if getattr(self, "query", None) is not None:324 with tf.name_scope(self.query.name):325 self.query.build([None, None, self.config.hidden_size])326 if getattr(self, "key", None) is not None:327 with tf.name_scope(self.key.name):328 self.key.build([None, None, self.config.hidden_size])329 if getattr(self, "value", None) is not None:330 with tf.name_scope(self.value.name):331 self.value.build([None, None, self.config.hidden_size])332 333 334class TFBertSelfOutput(keras.layers.Layer):335 def __init__(self, config: BertConfig, **kwargs):336 super().__init__(**kwargs)337 338 self.dense = keras.layers.Dense(339 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"340 )341 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")342 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)343 self.config = config344 345 def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:346 hidden_states = self.dense(inputs=hidden_states)347 hidden_states = self.dropout(inputs=hidden_states, training=training)348 hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)349 350 return hidden_states351 352 def build(self, input_shape=None):353 if self.built:354 return355 self.built = True356 if getattr(self, "dense", None) is not None:357 with tf.name_scope(self.dense.name):358 self.dense.build([None, None, self.config.hidden_size])359 if getattr(self, "LayerNorm", None) is not None:360 with tf.name_scope(self.LayerNorm.name):361 self.LayerNorm.build([None, None, self.config.hidden_size])362 363 364class TFBertAttention(keras.layers.Layer):365 def __init__(self, config: BertConfig, **kwargs):366 super().__init__(**kwargs)367 368 self.self_attention = TFBertSelfAttention(config, name="self")369 self.dense_output = TFBertSelfOutput(config, name="output")370 371 def prune_heads(self, heads):372 raise NotImplementedError373 374 def call(375 self,376 input_tensor: tf.Tensor,377 attention_mask: tf.Tensor,378 head_mask: tf.Tensor,379 encoder_hidden_states: tf.Tensor,380 encoder_attention_mask: tf.Tensor,381 past_key_value: tuple[tf.Tensor],382 output_attentions: bool,383 training: bool = False,384 ) -> tuple[tf.Tensor]:385 self_outputs = self.self_attention(386 hidden_states=input_tensor,387 attention_mask=attention_mask,388 head_mask=head_mask,389 encoder_hidden_states=encoder_hidden_states,390 encoder_attention_mask=encoder_attention_mask,391 past_key_value=past_key_value,392 output_attentions=output_attentions,393 training=training,394 )395 attention_output = self.dense_output(396 hidden_states=self_outputs[0], input_tensor=input_tensor, training=training397 )398 # add attentions (possibly with past_key_value) if we output them399 outputs = (attention_output,) + self_outputs[1:]400 401 return outputs402 403 def build(self, input_shape=None):404 if self.built:405 return406 self.built = True407 if getattr(self, "self_attention", None) is not None:408 with tf.name_scope(self.self_attention.name):409 self.self_attention.build(None)410 if getattr(self, "dense_output", None) is not None:411 with tf.name_scope(self.dense_output.name):412 self.dense_output.build(None)413 414 415class TFBertIntermediate(keras.layers.Layer):416 def __init__(self, config: BertConfig, **kwargs):417 super().__init__(**kwargs)418 419 self.dense = keras.layers.Dense(420 units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"421 )422 423 if isinstance(config.hidden_act, str):424 self.intermediate_act_fn = get_tf_activation(config.hidden_act)425 else:426 self.intermediate_act_fn = config.hidden_act427 self.config = config428 429 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:430 hidden_states = self.dense(inputs=hidden_states)431 hidden_states = self.intermediate_act_fn(hidden_states)432 433 return hidden_states434 435 def build(self, input_shape=None):436 if self.built:437 return438 self.built = True439 if getattr(self, "dense", None) is not None:440 with tf.name_scope(self.dense.name):441 self.dense.build([None, None, self.config.hidden_size])442 443 444class TFBertOutput(keras.layers.Layer):445 def __init__(self, config: BertConfig, **kwargs):446 super().__init__(**kwargs)447 448 self.dense = keras.layers.Dense(449 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"450 )451 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")452 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)453 self.config = config454 455 def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:456 hidden_states = self.dense(inputs=hidden_states)457 hidden_states = self.dropout(inputs=hidden_states, training=training)458 hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)459 460 return hidden_states461 462 def build(self, input_shape=None):463 if self.built:464 return465 self.built = True466 if getattr(self, "dense", None) is not None:467 with tf.name_scope(self.dense.name):468 self.dense.build([None, None, self.config.intermediate_size])469 if getattr(self, "LayerNorm", None) is not None:470 with tf.name_scope(self.LayerNorm.name):471 self.LayerNorm.build([None, None, self.config.hidden_size])472 473 474class TFBertLayer(keras.layers.Layer):475 def __init__(self, config: BertConfig, **kwargs):476 super().__init__(**kwargs)477 478 self.attention = TFBertAttention(config, name="attention")479 self.is_decoder = config.is_decoder480 self.add_cross_attention = config.add_cross_attention481 if self.add_cross_attention:482 if not self.is_decoder:483 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")484 self.crossattention = TFBertAttention(config, name="crossattention")485 self.intermediate = TFBertIntermediate(config, name="intermediate")486 self.bert_output = TFBertOutput(config, name="output")487 488 def call(489 self,490 hidden_states: tf.Tensor,491 attention_mask: tf.Tensor,492 head_mask: tf.Tensor,493 encoder_hidden_states: tf.Tensor | None,494 encoder_attention_mask: tf.Tensor | None,495 past_key_value: tuple[tf.Tensor] | None,496 output_attentions: bool,497 training: bool = False,498 ) -> tuple[tf.Tensor]:499 # decoder uni-directional self-attention cached key/values tuple is at positions 1,2500 self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None501 self_attention_outputs = self.attention(502 input_tensor=hidden_states,503 attention_mask=attention_mask,504 head_mask=head_mask,505 encoder_hidden_states=None,506 encoder_attention_mask=None,507 past_key_value=self_attn_past_key_value,508 output_attentions=output_attentions,509 training=training,510 )511 attention_output = self_attention_outputs[0]512 513 # if decoder, the last output is tuple of self-attn cache514 if self.is_decoder:515 outputs = self_attention_outputs[1:-1]516 present_key_value = self_attention_outputs[-1]517 else:518 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights519 520 cross_attn_present_key_value = None521 if self.is_decoder and encoder_hidden_states is not None:522 if not hasattr(self, "crossattention"):523 raise ValueError(524 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"525 " by setting `config.add_cross_attention=True`"526 )527 528 # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple529 cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None530 cross_attention_outputs = self.crossattention(531 input_tensor=attention_output,532 attention_mask=attention_mask,533 head_mask=head_mask,534 encoder_hidden_states=encoder_hidden_states,535 encoder_attention_mask=encoder_attention_mask,536 past_key_value=cross_attn_past_key_value,537 output_attentions=output_attentions,538 training=training,539 )540 attention_output = cross_attention_outputs[0]541 outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights542 543 # add cross-attn cache to positions 3,4 of present_key_value tuple544 cross_attn_present_key_value = cross_attention_outputs[-1]545 present_key_value = present_key_value + cross_attn_present_key_value546 547 intermediate_output = self.intermediate(hidden_states=attention_output)548 layer_output = self.bert_output(549 hidden_states=intermediate_output, input_tensor=attention_output, training=training550 )551 outputs = (layer_output,) + outputs # add attentions if we output them552 553 # if decoder, return the attn key/values as the last output554 if self.is_decoder:555 outputs = outputs + (present_key_value,)556 557 return outputs558 559 def build(self, input_shape=None):560 if self.built:561 return562 self.built = True563 if getattr(self, "attention", None) is not None:564 with tf.name_scope(self.attention.name):565 self.attention.build(None)566 if getattr(self, "intermediate", None) is not None:567 with tf.name_scope(self.intermediate.name):568 self.intermediate.build(None)569 if getattr(self, "bert_output", None) is not None:570 with tf.name_scope(self.bert_output.name):571 self.bert_output.build(None)572 if getattr(self, "crossattention", None) is not None:573 with tf.name_scope(self.crossattention.name):574 self.crossattention.build(None)575 576 577class TFBertEncoder(keras.layers.Layer):578 def __init__(self, config: BertConfig, **kwargs):579 super().__init__(**kwargs)580 self.config = config581 self.layer = [TFBertLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]582 583 def call(584 self,585 hidden_states: tf.Tensor,586 attention_mask: tf.Tensor,587 head_mask: tf.Tensor,588 encoder_hidden_states: tf.Tensor | None,589 encoder_attention_mask: tf.Tensor | None,590 past_key_values: tuple[tuple[tf.Tensor]] | None,591 use_cache: bool | None,592 output_attentions: bool,593 output_hidden_states: bool,594 return_dict: bool,595 training: bool = False,596 ) -> TFBaseModelOutputWithPastAndCrossAttentions | tuple[tf.Tensor]:597 all_hidden_states = () if output_hidden_states else None598 all_attentions = () if output_attentions else None599 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None600 601 next_decoder_cache = () if use_cache else None602 for i, layer_module in enumerate(self.layer):603 if output_hidden_states:604 all_hidden_states = all_hidden_states + (hidden_states,)605 606 past_key_value = past_key_values[i] if past_key_values is not None else None607 608 layer_outputs = layer_module(609 hidden_states=hidden_states,610 attention_mask=attention_mask,611 head_mask=head_mask[i],612 encoder_hidden_states=encoder_hidden_states,613 encoder_attention_mask=encoder_attention_mask,614 past_key_value=past_key_value,615 output_attentions=output_attentions,616 training=training,617 )618 hidden_states = layer_outputs[0]619 620 if use_cache:621 next_decoder_cache += (layer_outputs[-1],)622 623 if output_attentions:624 all_attentions = all_attentions + (layer_outputs[1],)625 if self.config.add_cross_attention and encoder_hidden_states is not None:626 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)627 628 # Add last layer629 if output_hidden_states:630 all_hidden_states = all_hidden_states + (hidden_states,)631 632 if not return_dict:633 return tuple(634 v for v in [hidden_states, all_hidden_states, all_attentions, all_cross_attentions] if v is not None635 )636 637 return TFBaseModelOutputWithPastAndCrossAttentions(638 last_hidden_state=hidden_states,639 past_key_values=next_decoder_cache,640 hidden_states=all_hidden_states,641 attentions=all_attentions,642 cross_attentions=all_cross_attentions,643 )644 645 def build(self, input_shape=None):646 if self.built:647 return648 self.built = True649 if getattr(self, "layer", None) is not None:650 for layer in self.layer:651 with tf.name_scope(layer.name):652 layer.build(None)653 654 655class TFBertPooler(keras.layers.Layer):656 def __init__(self, config: BertConfig, **kwargs):657 super().__init__(**kwargs)658 659 self.dense = keras.layers.Dense(660 units=config.hidden_size,661 kernel_initializer=get_initializer(config.initializer_range),662 activation="tanh",663 name="dense",664 )665 self.config = config666 667 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:668 # We "pool" the model by simply taking the hidden state corresponding669 # to the first token.670 first_token_tensor = hidden_states[:, 0]671 pooled_output = self.dense(inputs=first_token_tensor)672 673 return pooled_output674 675 def build(self, input_shape=None):676 if self.built:677 return678 self.built = True679 if getattr(self, "dense", None) is not None:680 with tf.name_scope(self.dense.name):681 self.dense.build([None, None, self.config.hidden_size])682 683 684class TFBertPredictionHeadTransform(keras.layers.Layer):685 def __init__(self, config: BertConfig, **kwargs):686 super().__init__(**kwargs)687 688 self.dense = keras.layers.Dense(689 units=config.hidden_size,690 kernel_initializer=get_initializer(config.initializer_range),691 name="dense",692 )693 694 if isinstance(config.hidden_act, str):695 self.transform_act_fn = get_tf_activation(config.hidden_act)696 else:697 self.transform_act_fn = config.hidden_act698 699 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")700 self.config = config701 702 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:703 hidden_states = self.dense(inputs=hidden_states)704 hidden_states = self.transform_act_fn(hidden_states)705 hidden_states = self.LayerNorm(inputs=hidden_states)706 707 return hidden_states708 709 def build(self, input_shape=None):710 if self.built:711 return712 self.built = True713 if getattr(self, "dense", None) is not None:714 with tf.name_scope(self.dense.name):715 self.dense.build([None, None, self.config.hidden_size])716 if getattr(self, "LayerNorm", None) is not None:717 with tf.name_scope(self.LayerNorm.name):718 self.LayerNorm.build([None, None, self.config.hidden_size])719 720 721class TFBertLMPredictionHead(keras.layers.Layer):722 def __init__(self, config: BertConfig, input_embeddings: keras.layers.Layer, **kwargs):723 super().__init__(**kwargs)724 725 self.config = config726 self.hidden_size = config.hidden_size727 728 self.transform = TFBertPredictionHeadTransform(config, name="transform")729 730 # The output weights are the same as the input embeddings, but there is731 # an output-only bias for each token.732 self.input_embeddings = input_embeddings733 734 def build(self, input_shape=None):735 self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")736 737 if self.built:738 return739 self.built = True740 if getattr(self, "transform", None) is not None:741 with tf.name_scope(self.transform.name):742 self.transform.build(None)743 744 def get_output_embeddings(self) -> keras.layers.Layer:745 return self.input_embeddings746 747 def set_output_embeddings(self, value: tf.Variable):748 self.input_embeddings.weight = value749 self.input_embeddings.vocab_size = shape_list(value)[0]750 751 def get_bias(self) -> dict[str, tf.Variable]:752 return {"bias": self.bias}753 754 def set_bias(self, value: tf.Variable):755 self.bias = value["bias"]756 self.config.vocab_size = shape_list(value["bias"])[0]757 758 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:759 hidden_states = self.transform(hidden_states=hidden_states)760 seq_length = shape_list(hidden_states)[1]761 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.hidden_size])762 hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)763 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])764 hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)765 766 return hidden_states767 768 769class TFBertMLMHead(keras.layers.Layer):770 def __init__(self, config: BertConfig, input_embeddings: keras.layers.Layer, **kwargs):771 super().__init__(**kwargs)772 773 self.predictions = TFBertLMPredictionHead(config, input_embeddings, name="predictions")774 775 def call(self, sequence_output: tf.Tensor) -> tf.Tensor:776 prediction_scores = self.predictions(hidden_states=sequence_output)777 778 return prediction_scores779 780 def build(self, input_shape=None):781 if self.built:782 return783 self.built = True784 if getattr(self, "predictions", None) is not None:785 with tf.name_scope(self.predictions.name):786 self.predictions.build(None)787 788 789class TFBertNSPHead(keras.layers.Layer):790 def __init__(self, config: BertConfig, **kwargs):791 super().__init__(**kwargs)792 793 self.seq_relationship = keras.layers.Dense(794 units=2,795 kernel_initializer=get_initializer(config.initializer_range),796 name="seq_relationship",797 )798 self.config = config799 800 def call(self, pooled_output: tf.Tensor) -> tf.Tensor:801 seq_relationship_score = self.seq_relationship(inputs=pooled_output)802 803 return seq_relationship_score804 805 def build(self, input_shape=None):806 if self.built:807 return808 self.built = True809 if getattr(self, "seq_relationship", None) is not None:810 with tf.name_scope(self.seq_relationship.name):811 self.seq_relationship.build([None, None, self.config.hidden_size])812 813 814@keras_serializable815class TFBertMainLayer(keras.layers.Layer):816 config_class = BertConfig817 818 def __init__(self, config: BertConfig, add_pooling_layer: bool = True, **kwargs):819 super().__init__(**kwargs)820 821 self.config = config822 self.is_decoder = config.is_decoder823 824 self.embeddings = TFBertEmbeddings(config, name="embeddings")825 self.encoder = TFBertEncoder(config, name="encoder")826 self.pooler = TFBertPooler(config, name="pooler") if add_pooling_layer else None827 828 def get_input_embeddings(self) -> keras.layers.Layer:829 return self.embeddings830 831 def set_input_embeddings(self, value: tf.Variable):832 self.embeddings.weight = value833 self.embeddings.vocab_size = shape_list(value)[0]834 835 def _prune_heads(self, heads_to_prune):836 """837 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base838 class PreTrainedModel839 """840 raise NotImplementedError841 842 @unpack_inputs843 def call(844 self,845 input_ids: TFModelInputType | None = None,846 attention_mask: np.ndarray | tf.Tensor | None = None,847 token_type_ids: np.ndarray | tf.Tensor | None = None,848 position_ids: np.ndarray | tf.Tensor | None = None,849 head_mask: np.ndarray | tf.Tensor | None = None,850 inputs_embeds: np.ndarray | tf.Tensor | None = None,851 encoder_hidden_states: np.ndarray | tf.Tensor | None = None,852 encoder_attention_mask: np.ndarray | tf.Tensor | None = None,853 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,854 use_cache: bool | None = None,855 output_attentions: bool | None = None,856 output_hidden_states: bool | None = None,857 return_dict: bool | None = None,858 training: bool = False,859 ) -> TFBaseModelOutputWithPoolingAndCrossAttentions | tuple[tf.Tensor]:860 if not self.config.is_decoder:861 use_cache = False862 863 if input_ids is not None and inputs_embeds is not None:864 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")865 elif input_ids is not None:866 input_shape = shape_list(input_ids)867 elif inputs_embeds is not None:868 input_shape = shape_list(inputs_embeds)[:-1]869 else:870 raise ValueError("You have to specify either input_ids or inputs_embeds")871 872 batch_size, seq_length = input_shape873 874 if past_key_values is None:875 past_key_values_length = 0876 past_key_values = [None] * len(self.encoder.layer)877 else:878 past_key_values_length = shape_list(past_key_values[0][0])[-2]879 880 if attention_mask is None:881 attention_mask = tf.fill(dims=(batch_size, seq_length + past_key_values_length), value=1)882 883 if token_type_ids is None:884 token_type_ids = tf.fill(dims=input_shape, value=0)885 886 embedding_output = self.embeddings(887 input_ids=input_ids,888 position_ids=position_ids,889 token_type_ids=token_type_ids,890 inputs_embeds=inputs_embeds,891 past_key_values_length=past_key_values_length,892 training=training,893 )894 895 # We create a 3D attention mask from a 2D tensor mask.896 # Sizes are [batch_size, 1, 1, to_seq_length]897 # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]898 # this attention mask is more simple than the triangular masking of causal attention899 # used in OpenAI GPT, we just need to prepare the broadcast dimension here.900 attention_mask_shape = shape_list(attention_mask)901 902 mask_seq_length = seq_length + past_key_values_length903 # Copied from `modeling_tf_t5.py`904 # Provided a padding mask of dimensions [batch_size, mask_seq_length]905 # - if the model is a decoder, apply a causal mask in addition to the padding mask906 # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]907 if self.is_decoder:908 seq_ids = tf.range(mask_seq_length)909 causal_mask = tf.less_equal(910 tf.tile(seq_ids[None, None, :], (batch_size, mask_seq_length, 1)),911 seq_ids[None, :, None],912 )913 causal_mask = tf.cast(causal_mask, dtype=attention_mask.dtype)914 extended_attention_mask = causal_mask * attention_mask[:, None, :]915 attention_mask_shape = shape_list(extended_attention_mask)916 extended_attention_mask = tf.reshape(917 extended_attention_mask, (attention_mask_shape[0], 1, attention_mask_shape[1], attention_mask_shape[2])918 )919 if past_key_values[0] is not None:920 # attention_mask needs to be sliced to the shape `[batch_size, 1, from_seq_length - cached_seq_length, to_seq_length]921 extended_attention_mask = extended_attention_mask[:, :, -seq_length:, :]922 else:923 extended_attention_mask = tf.reshape(924 attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1])925 )926 927 # Since attention_mask is 1.0 for positions we want to attend and 0.0 for928 # masked positions, this operation will create a tensor which is 0.0 for929 # positions we want to attend and -10000.0 for masked positions.930 # Since we are adding it to the raw scores before the softmax, this is931 # effectively the same as removing these entirely.932 extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)933 one_cst = tf.constant(1.0, dtype=embedding_output.dtype)934 ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)935 extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)936 937 # Copied from `modeling_tf_t5.py` with -1e9 -> -10000938 if self.is_decoder and encoder_attention_mask is not None:939 # If a 2D ou 3D attention mask is provided for the cross-attention940 # we need to make broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]941 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]942 encoder_attention_mask = tf.cast(encoder_attention_mask, dtype=extended_attention_mask.dtype)943 num_dims_encoder_attention_mask = len(shape_list(encoder_attention_mask))944 if num_dims_encoder_attention_mask == 3:945 encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]946 if num_dims_encoder_attention_mask == 2:947 encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]948 949 # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition950 # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270951 # encoder_extended_attention_mask = tf.math.equal(encoder_extended_attention_mask,952 # tf.transpose(encoder_extended_attention_mask, perm=(-1, -2)))953 954 encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -10000.0955 else:956 encoder_extended_attention_mask = None957 958 # Prepare head mask if needed959 # 1.0 in head_mask indicate we keep the head960 # attention_probs has shape bsz x n_heads x N x N961 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]962 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]963 if head_mask is not None:964 raise NotImplementedError965 else:966 head_mask = [None] * self.config.num_hidden_layers967 968 encoder_outputs = self.encoder(969 hidden_states=embedding_output,970 attention_mask=extended_attention_mask,971 head_mask=head_mask,972 encoder_hidden_states=encoder_hidden_states,973 encoder_attention_mask=encoder_extended_attention_mask,974 past_key_values=past_key_values,975 use_cache=use_cache,976 output_attentions=output_attentions,977 output_hidden_states=output_hidden_states,978 return_dict=return_dict,979 training=training,980 )981 982 sequence_output = encoder_outputs[0]983 pooled_output = self.pooler(hidden_states=sequence_output) if self.pooler is not None else None984 985 if not return_dict:986 return (987 sequence_output,988 pooled_output,989 ) + encoder_outputs[1:]990 991 return TFBaseModelOutputWithPoolingAndCrossAttentions(992 last_hidden_state=sequence_output,993 pooler_output=pooled_output,994 past_key_values=encoder_outputs.past_key_values,995 hidden_states=encoder_outputs.hidden_states,996 attentions=encoder_outputs.attentions,997 cross_attentions=encoder_outputs.cross_attentions,998 )999 1000 def build(self, input_shape=None):1001 if self.built:1002 return1003 self.built = True1004 if getattr(self, "embeddings", None) is not None:1005 with tf.name_scope(self.embeddings.name):1006 self.embeddings.build(None)1007 if getattr(self, "encoder", None) is not None:1008 with tf.name_scope(self.encoder.name):1009 self.encoder.build(None)1010 if getattr(self, "pooler", None) is not None:1011 with tf.name_scope(self.pooler.name):1012 self.pooler.build(None)1013 1014 1015class TFBertPreTrainedModel(TFPreTrainedModel):1016 """1017 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained1018 models.1019 """1020 1021 config_class = BertConfig1022 base_model_prefix = "bert"1023 1024 1025@dataclass1026class TFBertForPreTrainingOutput(ModelOutput):1027 """1028 Output type of [`TFBertForPreTraining`].1029 1030 Args:1031 prediction_logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):1032 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).1033 seq_relationship_logits (`tf.Tensor` of shape `(batch_size, 2)`):1034 Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation1035 before SoftMax).1036 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):1037 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape1038 `(batch_size, sequence_length, hidden_size)`.1039 1040 Hidden-states of the model at the output of each layer plus the initial embedding outputs.1041 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):1042 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,1043 sequence_length)`.1044 1045 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention1046 heads.1047 """1048 1049 loss: tf.Tensor | None = None1050 prediction_logits: tf.Tensor | None = None1051 seq_relationship_logits: tf.Tensor | None = None1052 hidden_states: tuple[tf.Tensor] | tf.Tensor | None = None1053 attentions: tuple[tf.Tensor] | tf.Tensor | None = None1054 1055 1056BERT_START_DOCSTRING = r"""1057 1058 This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the1059 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads1060 etc.)1061 1062 This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it1063 as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and1064 behavior.1065 1066 <Tip>1067 1068 TensorFlow models and layers in `transformers` accept two formats as input:1069 1070 - having all inputs as keyword arguments (like PyTorch models), or1071 - having all inputs as a list, tuple or dict in the first positional argument.1072 1073 The reason the second format is supported is that Keras methods prefer this format when passing inputs to models1074 and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just1075 pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second1076 format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with1077 the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first1078 positional argument:1079 1080 - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`1081 - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:1082 `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`1083 - a dictionary with one or several input Tensors associated to the input names given in the docstring:1084 `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`1085 1086 Note that when creating models and layers with1087 [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry1088 about any of this, as you can just pass inputs like you would to any other Python function!1089 1090 </Tip>1091 1092 Args:1093 config ([`BertConfig`]): Model configuration class with all the parameters of the model.1094 Initializing with a config file does not load the weights associated with the model, only the1095 configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.1096"""1097 1098BERT_INPUTS_DOCSTRING = r"""1099 Args:1100 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})`):1101 Indices of input sequence tokens in the vocabulary.1102 1103 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and1104 [`PreTrainedTokenizer.encode`] for details.1105 1106 [What are input IDs?](../glossary#input-ids)1107 attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):1108 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:1109 1110 - 1 for tokens that are **not masked**,1111 - 0 for tokens that are **masked**.1112 1113 [What are attention masks?](../glossary#attention-mask)1114 token_type_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):1115 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1116 1]`:1117 1118 - 0 corresponds to a *sentence A* token,1119 - 1 corresponds to a *sentence B* token.1120 1121 [What are token type IDs?](../glossary#token-type-ids)1122 position_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):1123 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,1124 config.max_position_embeddings - 1]`.1125 1126 [What are position IDs?](../glossary#position-ids)1127 head_mask (`np.ndarray` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):1128 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:1129 1130 - 1 indicates the head is **not masked**,1131 - 0 indicates the head is **masked**.1132 1133 inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `({0}, hidden_size)`, *optional*):1134 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This1135 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the1136 model's internal embedding lookup matrix.1137 output_attentions (`bool`, *optional*):1138 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned1139 tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the1140 config will be used instead.1141 output_hidden_states (`bool`, *optional*):1142 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for1143 more detail. This argument can be used only in eager mode, in graph mode the value in the config will be1144 used instead.1145 return_dict (`bool`, *optional*):1146 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in1147 eager mode, in graph mode the value will always be set to True.1148 training (`bool`, *optional*, defaults to `False``):1149 Whether or not to use the model in training mode (some modules like dropout modules have different1150 behaviors between training and evaluation).1151"""1152 1153 1154@add_start_docstrings(1155 "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",1156 BERT_START_DOCSTRING,1157)1158class TFBertModel(TFBertPreTrainedModel):1159 def __init__(self, config: BertConfig, add_pooling_layer: bool = True, *inputs, **kwargs):1160 super().__init__(config, *inputs, **kwargs)1161 1162 self.bert = TFBertMainLayer(config, add_pooling_layer, name="bert")1163 1164 @unpack_inputs1165 @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))1166 @add_code_sample_docstrings(1167 checkpoint=_CHECKPOINT_FOR_DOC,1168 output_type=TFBaseModelOutputWithPoolingAndCrossAttentions,1169 config_class=_CONFIG_FOR_DOC,1170 )1171 def call(1172 self,1173 input_ids: TFModelInputType | None = None,1174 attention_mask: np.ndarray | tf.Tensor | None = None,1175 token_type_ids: np.ndarray | tf.Tensor | None = None,1176 position_ids: np.ndarray | tf.Tensor | None = None,1177 head_mask: np.ndarray | tf.Tensor | None = None,1178 inputs_embeds: np.ndarray | tf.Tensor | None = None,1179 encoder_hidden_states: np.ndarray | tf.Tensor | None = None,1180 encoder_attention_mask: np.ndarray | tf.Tensor | None = None,1181 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,1182 use_cache: bool | None = None,1183 output_attentions: bool | None = None,1184 output_hidden_states: bool | None = None,1185 return_dict: bool | None = None,1186 training: bool | None = False,1187 ) -> TFBaseModelOutputWithPoolingAndCrossAttentions | tuple[tf.Tensor]:1188 r"""1189 encoder_hidden_states (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):1190 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if1191 the model is configured as a decoder.1192 encoder_attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1193 Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in1194 the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:1195 1196 - 1 for tokens that are **not masked**,1197 - 0 for tokens that are **masked**.1198 1199 past_key_values (`tuple[tuple[tf.Tensor]]` of length `config.n_layers`)1200 contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.