Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2020 The Allen Institute for AI team and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Tensorflow Longformer model."""16 17from __future__ import annotations18 19import warnings20from dataclasses import dataclass21 22import numpy as np23import tensorflow as tf24 25from ...activations_tf import get_tf_activation26from ...modeling_tf_utils import (27 TFMaskedLanguageModelingLoss,28 TFModelInputType,29 TFMultipleChoiceLoss,30 TFPreTrainedModel,31 TFQuestionAnsweringLoss,32 TFSequenceClassificationLoss,33 TFTokenClassificationLoss,34 get_initializer,35 keras,36 keras_serializable,37 unpack_inputs,38)39from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax40from ...utils import (41 ModelOutput,42 add_code_sample_docstrings,43 add_start_docstrings,44 add_start_docstrings_to_model_forward,45 logging,46)47from .configuration_longformer import LongformerConfig48 49 50logger = logging.get_logger(__name__)51 52_CHECKPOINT_FOR_DOC = "allenai/longformer-base-4096"53_CONFIG_FOR_DOC = "LongformerConfig"54 55LARGE_NEGATIVE = -1e856 57 58@dataclass59class TFLongformerBaseModelOutput(ModelOutput):60 """61 Base class for Longformer's outputs, with potential hidden states, local and global attentions.62 63 Args:64 last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):65 Sequence of hidden-states at the output of the last layer of the model.66 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):67 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape68 `(batch_size, sequence_length, hidden_size)`.69 70 Hidden-states of the model at the output of each layer plus the initial embedding outputs.71 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):72 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +73 attention_window + 1)`, where `x` is the number of tokens with global attention mask.74 75 Local attentions weights after the attention softmax, used to compute the weighted average in the76 self-attention heads. Those are the attention weights from every token in the sequence to every token with77 global attention (first `x` values) and to every token in the attention window (remaining `attention_window78 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the79 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a80 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding81 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.82 If the attention window contains a token with global attention, the attention weight at the corresponding83 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global84 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be85 accessed from `global_attentions`.86 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):87 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`88 is the number of tokens with global attention mask.89 90 Global attentions weights after the attention softmax, used to compute the weighted average in the91 self-attention heads. Those are the attention weights from every token with global attention to every token92 in the sequence.93 """94 95 last_hidden_state: tf.Tensor | None = None96 hidden_states: tuple[tf.Tensor, ...] | None = None97 attentions: tuple[tf.Tensor, ...] | None = None98 global_attentions: tuple[tf.Tensor, ...] | None = None99 100 101@dataclass102class TFLongformerBaseModelOutputWithPooling(ModelOutput):103 """104 Base class for Longformer's outputs that also contains a pooling of the last hidden states.105 106 Args:107 last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):108 Sequence of hidden-states at the output of the last layer of the model.109 pooler_output (`tf.Tensor` of shape `(batch_size, hidden_size)`):110 Last layer hidden-state of the first token of the sequence (classification token) further processed by a111 Linear layer and a Tanh activation function. The Linear layer weights are trained from the next sentence112 prediction (classification) objective during pretraining.113 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):114 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape115 `(batch_size, sequence_length, hidden_size)`.116 117 Hidden-states of the model at the output of each layer plus the initial embedding outputs.118 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):119 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +120 attention_window + 1)`, where `x` is the number of tokens with global attention mask.121 122 Local attentions weights after the attention softmax, used to compute the weighted average in the123 self-attention heads. Those are the attention weights from every token in the sequence to every token with124 global attention (first `x` values) and to every token in the attention window (remaining `attention_window125 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the126 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a127 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding128 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.129 If the attention window contains a token with global attention, the attention weight at the corresponding130 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global131 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be132 accessed from `global_attentions`.133 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):134 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`135 is the number of tokens with global attention mask.136 137 Global attentions weights after the attention softmax, used to compute the weighted average in the138 self-attention heads. Those are the attention weights from every token with global attention to every token139 in the sequence.140 """141 142 last_hidden_state: tf.Tensor | None = None143 pooler_output: tf.Tensor | None = None144 hidden_states: tuple[tf.Tensor, ...] | None = None145 attentions: tuple[tf.Tensor, ...] | None = None146 global_attentions: tuple[tf.Tensor, ...] | None = None147 148 149@dataclass150class TFLongformerMaskedLMOutput(ModelOutput):151 """152 Base class for masked language models outputs.153 154 Args:155 loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):156 Masked language modeling (MLM) loss.157 logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):158 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).159 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):160 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape161 `(batch_size, sequence_length, hidden_size)`.162 163 Hidden-states of the model at the output of each layer plus the initial embedding outputs.164 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):165 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +166 attention_window + 1)`, where `x` is the number of tokens with global attention mask.167 168 Local attentions weights after the attention softmax, used to compute the weighted average in the169 self-attention heads. Those are the attention weights from every token in the sequence to every token with170 global attention (first `x` values) and to every token in the attention window (remaining `attention_window171 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the172 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a173 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding174 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.175 If the attention window contains a token with global attention, the attention weight at the corresponding176 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global177 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be178 accessed from `global_attentions`.179 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):180 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`181 is the number of tokens with global attention mask.182 183 Global attentions weights after the attention softmax, used to compute the weighted average in the184 self-attention heads. Those are the attention weights from every token with global attention to every token185 in the sequence.186 """187 188 loss: tf.Tensor | None = None189 logits: tf.Tensor | None = None190 hidden_states: tuple[tf.Tensor, ...] | None = None191 attentions: tuple[tf.Tensor, ...] | None = None192 global_attentions: tuple[tf.Tensor, ...] | None = None193 194 195@dataclass196class TFLongformerQuestionAnsweringModelOutput(ModelOutput):197 """198 Base class for outputs of question answering Longformer models.199 200 Args:201 loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):202 Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.203 start_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):204 Span-start scores (before SoftMax).205 end_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):206 Span-end scores (before SoftMax).207 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):208 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape209 `(batch_size, sequence_length, hidden_size)`.210 211 Hidden-states of the model at the output of each layer plus the initial embedding outputs.212 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):213 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +214 attention_window + 1)`, where `x` is the number of tokens with global attention mask.215 216 Local attentions weights after the attention softmax, used to compute the weighted average in the217 self-attention heads. Those are the attention weights from every token in the sequence to every token with218 global attention (first `x` values) and to every token in the attention window (remaining `attention_window219 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the220 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a221 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding222 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.223 If the attention window contains a token with global attention, the attention weight at the corresponding224 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global225 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be226 accessed from `global_attentions`.227 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):228 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`229 is the number of tokens with global attention mask.230 231 Global attentions weights after the attention softmax, used to compute the weighted average in the232 self-attention heads. Those are the attention weights from every token with global attention to every token233 in the sequence.234 """235 236 loss: tf.Tensor | None = None237 start_logits: tf.Tensor | None = None238 end_logits: tf.Tensor | None = None239 hidden_states: tuple[tf.Tensor, ...] | None = None240 attentions: tuple[tf.Tensor, ...] | None = None241 global_attentions: tuple[tf.Tensor, ...] | None = None242 243 244@dataclass245class TFLongformerSequenceClassifierOutput(ModelOutput):246 """247 Base class for outputs of sentence classification models.248 249 Args:250 loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):251 Classification (or regression if config.num_labels==1) loss.252 logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):253 Classification (or regression if config.num_labels==1) scores (before SoftMax).254 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):255 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape256 `(batch_size, sequence_length, hidden_size)`.257 258 Hidden-states of the model at the output of each layer plus the initial embedding outputs.259 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):260 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +261 attention_window + 1)`, where `x` is the number of tokens with global attention mask.262 263 Local attentions weights after the attention softmax, used to compute the weighted average in the264 self-attention heads. Those are the attention weights from every token in the sequence to every token with265 global attention (first `x` values) and to every token in the attention window (remaining `attention_window266 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the267 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a268 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding269 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.270 If the attention window contains a token with global attention, the attention weight at the corresponding271 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global272 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be273 accessed from `global_attentions`.274 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):275 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`276 is the number of tokens with global attention mask.277 278 Global attentions weights after the attention softmax, used to compute the weighted average in the279 self-attention heads. Those are the attention weights from every token with global attention to every token280 in the sequence.281 """282 283 loss: tf.Tensor | None = None284 logits: tf.Tensor | None = None285 hidden_states: tuple[tf.Tensor, ...] | None = None286 attentions: tuple[tf.Tensor, ...] | None = None287 global_attentions: tuple[tf.Tensor, ...] | None = None288 289 290@dataclass291class TFLongformerMultipleChoiceModelOutput(ModelOutput):292 """293 Base class for outputs of multiple choice models.294 295 Args:296 loss (`tf.Tensor` of shape *(1,)*, *optional*, returned when `labels` is provided):297 Classification loss.298 logits (`tf.Tensor` of shape `(batch_size, num_choices)`):299 *num_choices* is the second dimension of the input tensors. (see *input_ids* above).300 301 Classification scores (before SoftMax).302 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):303 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape304 `(batch_size, sequence_length, hidden_size)`.305 306 Hidden-states of the model at the output of each layer plus the initial embedding outputs.307 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):308 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +309 attention_window + 1)`, where `x` is the number of tokens with global attention mask.310 311 Local attentions weights after the attention softmax, used to compute the weighted average in the312 self-attention heads. Those are the attention weights from every token in the sequence to every token with313 global attention (first `x` values) and to every token in the attention window (remaining `attention_window314 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the315 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a316 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding317 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.318 If the attention window contains a token with global attention, the attention weight at the corresponding319 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global320 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be321 accessed from `global_attentions`.322 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):323 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`324 is the number of tokens with global attention mask.325 326 Global attentions weights after the attention softmax, used to compute the weighted average in the327 self-attention heads. Those are the attention weights from every token with global attention to every token328 in the sequence.329 """330 331 loss: tf.Tensor | None = None332 logits: tf.Tensor | None = None333 hidden_states: tuple[tf.Tensor, ...] | None = None334 attentions: tuple[tf.Tensor, ...] | None = None335 global_attentions: tuple[tf.Tensor, ...] | None = None336 337 338@dataclass339class TFLongformerTokenClassifierOutput(ModelOutput):340 """341 Base class for outputs of token classification models.342 343 Args:344 loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :345 Classification loss.346 logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.num_labels)`):347 Classification scores (before SoftMax).348 hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):349 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape350 `(batch_size, sequence_length, hidden_size)`.351 352 Hidden-states of the model at the output of each layer plus the initial embedding outputs.353 attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):354 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x +355 attention_window + 1)`, where `x` is the number of tokens with global attention mask.356 357 Local attentions weights after the attention softmax, used to compute the weighted average in the358 self-attention heads. Those are the attention weights from every token in the sequence to every token with359 global attention (first `x` values) and to every token in the attention window (remaining `attention_window360 + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the361 remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a362 token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding363 (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens.364 If the attention window contains a token with global attention, the attention weight at the corresponding365 index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global366 attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be367 accessed from `global_attentions`.368 global_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):369 Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`, where `x`370 is the number of tokens with global attention mask.371 372 Global attentions weights after the attention softmax, used to compute the weighted average in the373 self-attention heads. Those are the attention weights from every token with global attention to every token374 in the sequence.375 """376 377 loss: tf.Tensor | None = None378 logits: tf.Tensor | None = None379 hidden_states: tuple[tf.Tensor, ...] | None = None380 attentions: tuple[tf.Tensor, ...] | None = None381 global_attentions: tuple[tf.Tensor, ...] | None = None382 383 384def _compute_global_attention_mask(input_ids_shape, sep_token_indices, before_sep_token=True):385 """386 Computes global attention mask by putting attention on all tokens before `sep_token_id` if `before_sep_token is387 True` else after `sep_token_id`.388 """389 assert shape_list(sep_token_indices)[1] == 2, "`input_ids` should have two dimensions"390 question_end_index = tf.reshape(sep_token_indices, (input_ids_shape[0], 3, 2))[:, 0, 1][:, None]391 # bool attention mask with True in locations of global attention392 attention_mask = tf.expand_dims(tf.range(input_ids_shape[1], dtype=tf.int64), axis=0)393 attention_mask = tf.tile(attention_mask, (input_ids_shape[0], 1))394 if before_sep_token is True:395 question_end_index = tf.tile(question_end_index, (1, input_ids_shape[1]))396 attention_mask = tf.cast(attention_mask < question_end_index, dtype=question_end_index.dtype)397 else:398 # last token is separation token and should not be counted and in the middle are two separation tokens399 question_end_index = tf.tile(question_end_index + 1, (1, input_ids_shape[1]))400 attention_mask = tf.cast(401 attention_mask > question_end_index,402 dtype=question_end_index.dtype,403 ) * tf.cast(attention_mask < input_ids_shape[-1], dtype=question_end_index.dtype)404 405 return attention_mask406 407 408# Copied from transformers.models.roberta.modeling_tf_roberta.TFRobertaLMHead with Roberta->Longformer409class TFLongformerLMHead(keras.layers.Layer):410 """Longformer Head for masked language modeling."""411 412 def __init__(self, config, input_embeddings, **kwargs):413 super().__init__(**kwargs)414 415 self.config = config416 self.hidden_size = config.hidden_size417 self.dense = keras.layers.Dense(418 config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"419 )420 self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")421 self.act = get_tf_activation("gelu")422 423 # The output weights are the same as the input embeddings, but there is424 # an output-only bias for each token.425 self.decoder = input_embeddings426 427 def build(self, input_shape=None):428 self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")429 430 if self.built:431 return432 self.built = True433 if getattr(self, "dense", None) is not None:434 with tf.name_scope(self.dense.name):435 self.dense.build([None, None, self.config.hidden_size])436 if getattr(self, "layer_norm", None) is not None:437 with tf.name_scope(self.layer_norm.name):438 self.layer_norm.build([None, None, self.config.hidden_size])439 440 def get_output_embeddings(self):441 return self.decoder442 443 def set_output_embeddings(self, value):444 self.decoder.weight = value445 self.decoder.vocab_size = shape_list(value)[0]446 447 def get_bias(self):448 return {"bias": self.bias}449 450 def set_bias(self, value):451 self.bias = value["bias"]452 self.config.vocab_size = shape_list(value["bias"])[0]453 454 def call(self, hidden_states):455 hidden_states = self.dense(hidden_states)456 hidden_states = self.act(hidden_states)457 hidden_states = self.layer_norm(hidden_states)458 459 # project back to size of vocabulary with bias460 seq_length = shape_list(tensor=hidden_states)[1]461 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.hidden_size])462 hidden_states = tf.matmul(a=hidden_states, b=self.decoder.weight, transpose_b=True)463 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])464 hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)465 466 return hidden_states467 468 469class TFLongformerEmbeddings(keras.layers.Layer):470 """471 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing and some extra casting.472 """473 474 def __init__(self, config, **kwargs):475 super().__init__(**kwargs)476 477 self.padding_idx = 1478 self.config = config479 self.hidden_size = config.hidden_size480 self.max_position_embeddings = config.max_position_embeddings481 self.initializer_range = config.initializer_range482 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")483 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)484 485 def build(self, input_shape=None):486 with tf.name_scope("word_embeddings"):487 self.weight = self.add_weight(488 name="weight",489 shape=[self.config.vocab_size, self.hidden_size],490 initializer=get_initializer(self.initializer_range),491 )492 493 with tf.name_scope("token_type_embeddings"):494 self.token_type_embeddings = self.add_weight(495 name="embeddings",496 shape=[self.config.type_vocab_size, self.hidden_size],497 initializer=get_initializer(self.initializer_range),498 )499 500 with tf.name_scope("position_embeddings"):501 self.position_embeddings = self.add_weight(502 name="embeddings",503 shape=[self.max_position_embeddings, self.hidden_size],504 initializer=get_initializer(self.initializer_range),505 )506 507 if self.built:508 return509 self.built = True510 if getattr(self, "LayerNorm", None) is not None:511 with tf.name_scope(self.LayerNorm.name):512 self.LayerNorm.build([None, None, self.config.hidden_size])513 514 def create_position_ids_from_input_ids(self, input_ids, past_key_values_length=0):515 """516 Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding517 symbols are ignored. This is modified from fairseq's `utils.make_positions`.518 519 Args:520 input_ids: tf.Tensor521 Returns: tf.Tensor522 """523 mask = tf.cast(tf.math.not_equal(input_ids, self.padding_idx), dtype=input_ids.dtype)524 incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask525 526 return incremental_indices + self.padding_idx527 528 def call(529 self,530 input_ids=None,531 position_ids=None,532 token_type_ids=None,533 inputs_embeds=None,534 past_key_values_length=0,535 training=False,536 ):537 """538 Applies embedding based on inputs tensor.539 540 Returns:541 final_embeddings (`tf.Tensor`): output embedding tensor.542 """543 assert not (input_ids is None and inputs_embeds is None)544 545 if input_ids is not None:546 check_embeddings_within_bounds(input_ids, self.config.vocab_size)547 inputs_embeds = tf.gather(params=self.weight, indices=input_ids)548 549 input_shape = shape_list(inputs_embeds)[:-1]550 551 if token_type_ids is None:552 token_type_ids = tf.cast(tf.fill(dims=input_shape, value=0), tf.int64)553 554 if position_ids is None:555 if input_ids is not None:556 # Create the position ids from the input token ids. Any padded tokens remain padded.557 position_ids = self.create_position_ids_from_input_ids(558 input_ids=input_ids, past_key_values_length=past_key_values_length559 )560 else:561 position_ids = tf.expand_dims(562 tf.range(start=self.padding_idx + 1, limit=input_shape[-1] + self.padding_idx + 1, dtype=tf.int64),563 axis=0,564 )565 566 position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)567 token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)568 final_embeddings = inputs_embeds + position_embeds + token_type_embeds569 final_embeddings = self.LayerNorm(inputs=final_embeddings)570 final_embeddings = self.dropout(inputs=final_embeddings, training=training)571 572 return final_embeddings573 574 575# Copied from transformers.models.bert.modeling_tf_bert.TFBertIntermediate with Bert->Longformer576class TFLongformerIntermediate(keras.layers.Layer):577 def __init__(self, config: LongformerConfig, **kwargs):578 super().__init__(**kwargs)579 580 self.dense = keras.layers.Dense(581 units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"582 )583 584 if isinstance(config.hidden_act, str):585 self.intermediate_act_fn = get_tf_activation(config.hidden_act)586 else:587 self.intermediate_act_fn = config.hidden_act588 self.config = config589 590 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:591 hidden_states = self.dense(inputs=hidden_states)592 hidden_states = self.intermediate_act_fn(hidden_states)593 594 return hidden_states595 596 def build(self, input_shape=None):597 if self.built:598 return599 self.built = True600 if getattr(self, "dense", None) is not None:601 with tf.name_scope(self.dense.name):602 self.dense.build([None, None, self.config.hidden_size])603 604 605# Copied from transformers.models.bert.modeling_tf_bert.TFBertOutput with Bert->Longformer606class TFLongformerOutput(keras.layers.Layer):607 def __init__(self, config: LongformerConfig, **kwargs):608 super().__init__(**kwargs)609 610 self.dense = keras.layers.Dense(611 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"612 )613 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")614 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)615 self.config = config616 617 def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:618 hidden_states = self.dense(inputs=hidden_states)619 hidden_states = self.dropout(inputs=hidden_states, training=training)620 hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)621 622 return hidden_states623 624 def build(self, input_shape=None):625 if self.built:626 return627 self.built = True628 if getattr(self, "dense", None) is not None:629 with tf.name_scope(self.dense.name):630 self.dense.build([None, None, self.config.intermediate_size])631 if getattr(self, "LayerNorm", None) is not None:632 with tf.name_scope(self.LayerNorm.name):633 self.LayerNorm.build([None, None, self.config.hidden_size])634 635 636# Copied from transformers.models.bert.modeling_tf_bert.TFBertPooler with Bert->Longformer637class TFLongformerPooler(keras.layers.Layer):638 def __init__(self, config: LongformerConfig, **kwargs):639 super().__init__(**kwargs)640 641 self.dense = keras.layers.Dense(642 units=config.hidden_size,643 kernel_initializer=get_initializer(config.initializer_range),644 activation="tanh",645 name="dense",646 )647 self.config = config648 649 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:650 # We "pool" the model by simply taking the hidden state corresponding651 # to the first token.652 first_token_tensor = hidden_states[:, 0]653 pooled_output = self.dense(inputs=first_token_tensor)654 655 return pooled_output656 657 def build(self, input_shape=None):658 if self.built:659 return660 self.built = True661 if getattr(self, "dense", None) is not None:662 with tf.name_scope(self.dense.name):663 self.dense.build([None, None, self.config.hidden_size])664 665 666# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput with Bert->Longformer667class TFLongformerSelfOutput(keras.layers.Layer):668 def __init__(self, config: LongformerConfig, **kwargs):669 super().__init__(**kwargs)670 671 self.dense = keras.layers.Dense(672 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"673 )674 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")675 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)676 self.config = config677 678 def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:679 hidden_states = self.dense(inputs=hidden_states)680 hidden_states = self.dropout(inputs=hidden_states, training=training)681 hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)682 683 return hidden_states684 685 def build(self, input_shape=None):686 if self.built:687 return688 self.built = True689 if getattr(self, "dense", None) is not None:690 with tf.name_scope(self.dense.name):691 self.dense.build([None, None, self.config.hidden_size])692 if getattr(self, "LayerNorm", None) is not None:693 with tf.name_scope(self.LayerNorm.name):694 self.LayerNorm.build([None, None, self.config.hidden_size])695 696 697class TFLongformerSelfAttention(keras.layers.Layer):698 def __init__(self, config, layer_id, **kwargs):699 super().__init__(**kwargs)700 self.config = config701 702 if config.hidden_size % config.num_attention_heads != 0:703 raise ValueError(704 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "705 f"heads ({config.num_attention_heads}"706 )707 708 self.num_heads = config.num_attention_heads709 self.head_dim = int(config.hidden_size / config.num_attention_heads)710 self.embed_dim = config.hidden_size711 self.query = keras.layers.Dense(712 self.embed_dim,713 kernel_initializer=get_initializer(config.initializer_range),714 name="query",715 )716 self.key = keras.layers.Dense(717 self.embed_dim,718 kernel_initializer=get_initializer(config.initializer_range),719 name="key",720 )721 self.value = keras.layers.Dense(722 self.embed_dim,723 kernel_initializer=get_initializer(config.initializer_range),724 name="value",725 )726 727 # separate projection layers for tokens with global attention728 self.query_global = keras.layers.Dense(729 self.embed_dim,730 kernel_initializer=get_initializer(config.initializer_range),731 name="query_global",732 )733 self.key_global = keras.layers.Dense(734 self.embed_dim,735 kernel_initializer=get_initializer(config.initializer_range),736 name="key_global",737 )738 self.value_global = keras.layers.Dense(739 self.embed_dim,740 kernel_initializer=get_initializer(config.initializer_range),741 name="value_global",742 )743 self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)744 self.global_dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)745 self.layer_id = layer_id746 attention_window = config.attention_window[self.layer_id]747 748 assert attention_window % 2 == 0, (749 f"`attention_window` for layer {self.layer_id} has to be an even value. Given {attention_window}"750 )751 assert attention_window > 0, (752 f"`attention_window` for layer {self.layer_id} has to be positive. Given {attention_window}"753 )754 755 self.one_sided_attn_window_size = attention_window // 2756 757 def build(self, input_shape=None):758 if not self.built:759 with tf.name_scope("query_global"):760 self.query_global.build((self.config.hidden_size,))761 with tf.name_scope("key_global"):762 self.key_global.build((self.config.hidden_size,))763 with tf.name_scope("value_global"):764 self.value_global.build((self.config.hidden_size,))765 766 if self.built:767 return768 self.built = True769 if getattr(self, "query", None) is not None:770 with tf.name_scope(self.query.name):771 self.query.build([None, None, self.config.hidden_size])772 if getattr(self, "key", None) is not None:773 with tf.name_scope(self.key.name):774 self.key.build([None, None, self.config.hidden_size])775 if getattr(self, "value", None) is not None:776 with tf.name_scope(self.value.name):777 self.value.build([None, None, self.config.hidden_size])778 if getattr(self, "query_global", None) is not None:779 with tf.name_scope(self.query_global.name):780 self.query_global.build([None, None, self.config.hidden_size])781 if getattr(self, "key_global", None) is not None:782 with tf.name_scope(self.key_global.name):783 self.key_global.build([None, None, self.config.hidden_size])784 if getattr(self, "value_global", None) is not None:785 with tf.name_scope(self.value_global.name):786 self.value_global.build([None, None, self.config.hidden_size])787 788 def call(789 self,790 inputs,791 training=False,792 ):793 """794 LongformerSelfAttention expects *len(hidden_states)* to be multiple of *attention_window*. Padding to795 *attention_window* happens in LongformerModel.forward to avoid redoing the padding on each layer.796 797 The *attention_mask* is changed in [`LongformerModel.forward`] from 0, 1, 2 to:798 799 - -10000: no attention800 - 0: local attention801 - +10000: global attention802 """803 # retrieve input args804 (805 hidden_states,806 attention_mask,807 layer_head_mask,808 is_index_masked,809 is_index_global_attn,810 is_global_attn,811 ) = inputs812 813 # project hidden states814 query_vectors = self.query(hidden_states)815 key_vectors = self.key(hidden_states)816 value_vectors = self.value(hidden_states)817 batch_size, seq_len, embed_dim = shape_list(hidden_states)818 819 tf.debugging.assert_equal(820 embed_dim,821 self.embed_dim,822 message=f"hidden_states should have embed_dim = {self.embed_dim}, but has {embed_dim}",823 )824 825 # normalize query826 query_vectors /= tf.math.sqrt(tf.cast(self.head_dim, dtype=query_vectors.dtype))827 query_vectors = tf.reshape(query_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))828 key_vectors = tf.reshape(key_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))829 830 # attn_probs = (batch_size, seq_len, num_heads, window*2+1)831 attn_scores = self._sliding_chunks_query_key_matmul(832 query_vectors, key_vectors, self.one_sided_attn_window_size833 )834 835 # values to pad for attention probs836 remove_from_windowed_attention_mask = attention_mask != 0837 # cast to fp32/fp16 then replace 1's with -inf838 float_mask = tf.cast(remove_from_windowed_attention_mask, dtype=query_vectors.dtype) * LARGE_NEGATIVE839 840 # diagonal mask with zeros everywhere and -inf inplace of padding841 diagonal_mask = self._sliding_chunks_query_key_matmul(842 tf.ones(shape_list(attention_mask)),843 float_mask,844 self.one_sided_attn_window_size,845 )846 847 # pad local attention probs848 attn_scores += diagonal_mask849 850 tf.debugging.assert_equal(851 shape_list(attn_scores),852 [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1],853 message=(854 f"attn_probs should be of size ({batch_size}, {seq_len}, {self.num_heads},"855 f" {self.one_sided_attn_window_size * 2 + 1}), but is of size {shape_list(attn_scores)}"856 ),857 )858 859 # compute global attn indices required through out forward fn860 (861 max_num_global_attn_indices,862 is_index_global_attn_nonzero,863 is_local_index_global_attn_nonzero,864 is_local_index_no_global_attn_nonzero,865 ) = self._get_global_attn_indices(is_index_global_attn)866 867 # this function is only relevant for global attention868 if is_global_attn:869 attn_scores = self._concat_with_global_key_attn_probs(870 attn_scores=attn_scores,871 query_vectors=query_vectors,872 key_vectors=key_vectors,873 max_num_global_attn_indices=max_num_global_attn_indices,874 is_index_global_attn_nonzero=is_index_global_attn_nonzero,875 is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,876 is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,877 )878 879 attn_probs = stable_softmax(attn_scores, axis=-1)880 881 # softmax sometimes inserts NaN if all positions are masked, replace them with 0882 # Make sure to create a mask with the proper shape:883 # if is_global_attn==True => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1]884 # if is_global_attn==False => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]885 if is_global_attn:886 masked_index = tf.tile(887 is_index_masked[:, :, None, None],888 (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1),889 )890 else:891 masked_index = tf.tile(892 is_index_masked[:, :, None, None],893 (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + 1),894 )895 attn_probs = tf.where(896 masked_index,897 tf.zeros(shape_list(masked_index), dtype=attn_probs.dtype),898 attn_probs,899 )900 901 if layer_head_mask is not None:902 tf.debugging.assert_equal(903 shape_list(layer_head_mask),904 [self.num_heads],905 message=(906 f"Head mask for a single layer should be of size {(self.num_heads)}, but is"907 f" {shape_list(layer_head_mask)}"908 ),909 )910 911 attn_probs = tf.reshape(layer_head_mask, (1, 1, -1, 1)) * attn_probs912 913 # apply dropout914 attn_probs = self.dropout(attn_probs, training=training)915 value_vectors = tf.reshape(value_vectors, (batch_size, seq_len, self.num_heads, self.head_dim))916 917 # if global attention, compute sum of global and local attn918 919 if is_global_attn:920 attn_output = self._compute_attn_output_with_global_indices(921 value_vectors=value_vectors,922 attn_probs=attn_probs,923 max_num_global_attn_indices=max_num_global_attn_indices,924 is_index_global_attn_nonzero=is_index_global_attn_nonzero,925 is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,926 )927 else:928 attn_output = self._sliding_chunks_matmul_attn_probs_value(929 attn_probs, value_vectors, self.one_sided_attn_window_size930 )931 932 tf.debugging.assert_equal(933 shape_list(attn_output), [batch_size, seq_len, self.num_heads, self.head_dim], message="Unexpected size"934 )935 936 attn_output = tf.reshape(attn_output, (batch_size, seq_len, embed_dim))937 938 # compute value for global attention and overwrite to attention output939 if is_global_attn:940 attn_output, global_attn_probs = self._compute_global_attn_output_from_hidden(941 attn_output=attn_output,942 hidden_states=hidden_states,943 max_num_global_attn_indices=max_num_global_attn_indices,944 layer_head_mask=layer_head_mask,945 is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,946 is_index_global_attn_nonzero=is_index_global_attn_nonzero,947 is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,948 is_index_masked=is_index_masked,949 training=training,950 )951 else:952 # Leave attn_output unchanged953 global_attn_probs = tf.zeros((batch_size, self.num_heads, max_num_global_attn_indices, seq_len))954 955 # make sure that local attention probabilities are set to 0 for indices of global attn956 # Make sure to create a mask with the proper shape:957 # if is_global_attn==True => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1]958 # if is_global_attn==False => [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]959 if is_global_attn:960 masked_global_attn_index = tf.tile(961 is_index_global_attn[:, :, None, None],962 (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1),963 )964 else:965 masked_global_attn_index = tf.tile(966 is_index_global_attn[:, :, None, None],967 (1, 1, self.num_heads, self.one_sided_attn_window_size * 2 + 1),968 )969 attn_probs = tf.where(970 masked_global_attn_index,971 tf.zeros(shape_list(masked_global_attn_index), dtype=attn_probs.dtype),972 attn_probs,973 )974 975 outputs = (attn_output, attn_probs, global_attn_probs)976 977 return outputs978 979 def _sliding_chunks_query_key_matmul(self, query, key, window_overlap):980 """981 Matrix multiplication of query and key tensors using with a sliding window attention pattern. This982 implementation splits the input into overlapping chunks of size 2w (e.g. 512 for pretrained Longformer) with an983 overlap of size window_overlap984 """985 batch_size, seq_len, num_heads, head_dim = shape_list(query)986 987 tf.debugging.assert_equal(988 seq_len % (window_overlap * 2),989 0,990 message=f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}",991 )992 tf.debugging.assert_equal(993 shape_list(query),994 shape_list(key),995 message=(996 f"Shape of query and key should be equal, but got query: {shape_list(query)} and key:"997 f" {shape_list(key)}"998 ),999 )1000 1001 chunks_count = seq_len // window_overlap - 11002 1003 # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size window_overlap * 21004 query = tf.reshape(1005 tf.transpose(query, (0, 2, 1, 3)),1006 (batch_size * num_heads, seq_len, head_dim),1007 )1008 key = tf.reshape(tf.transpose(key, (0, 2, 1, 3)), (batch_size * num_heads, seq_len, head_dim))1009 chunked_query = self._chunk(query, window_overlap)1010 chunked_key = self._chunk(key, window_overlap)1011 1012 # matrix multiplication1013 # bcxd: batch_size * num_heads x chunks x 2window_overlap x head_dim1014 # bcyd: batch_size * num_heads x chunks x 2window_overlap x head_dim1015 # bcxy: batch_size * num_heads x chunks x 2window_overlap x 2window_overlap1016 chunked_query = tf.cast(chunked_query, dtype=chunked_key.dtype)1017 chunked_attention_scores = tf.einsum("bcxd,bcyd->bcxy", chunked_query, chunked_key) # multiply1018 1019 # convert diagonals into columns1020 paddings = tf.convert_to_tensor([[0, 0], [0, 0], [0, 1], [0, 0]])1021 diagonal_chunked_attention_scores = self._pad_and_transpose_last_two_dims(chunked_attention_scores, paddings)1022 1023 # allocate space for the overall attention matrix where the chunks are combined. The last dimension1024 # has (window_overlap * 2 + 1) columns. The first (window_overlap) columns are the window_overlap lower triangles (attention from a word to1025 # window_overlap previous words). The following column is attention score from each word to itself, then1026 # followed by window_overlap columns for the upper triangle.1027 1028 # copy parts from diagonal_chunked_attention_scores into the combined matrix of attentions1029 # - copying the main diagonal and the upper triangle1030 # TODO: This code is most likely not very efficient and should be improved1031 diagonal_attn_scores_up_triang = tf.concat(1032 [1033 diagonal_chunked_attention_scores[:, :, :window_overlap, : window_overlap + 1],1034 diagonal_chunked_attention_scores[:, -1:, window_overlap:, : window_overlap + 1],1035 ],1036 axis=1,1037 )1038 1039 # - copying the lower triangle1040 diagonal_attn_scores_low_triang = tf.concat(1041 [1042 tf.zeros(1043 (batch_size * num_heads, 1, window_overlap, window_overlap),1044 dtype=diagonal_chunked_attention_scores.dtype,1045 ),1046 diagonal_chunked_attention_scores[:, :, -(window_overlap + 1) : -1, window_overlap + 1 :],1047 ],1048 axis=1,1049 )1050 diagonal_attn_scores_first_chunk = tf.concat(1051 [1052 tf.roll(1053 diagonal_chunked_attention_scores,1054 shift=[1, window_overlap],1055 axis=[2, 3],1056 )[:, :, :window_overlap, :window_overlap],1057 tf.zeros(1058 (batch_size * num_heads, 1, window_overlap, window_overlap),1059 dtype=diagonal_chunked_attention_scores.dtype,1060 ),1061 ],1062 axis=1,1063 )1064 first_chunk_mask = (1065 tf.tile(1066 tf.range(chunks_count + 1, dtype=tf.int64)[None, :, None, None],1067 (batch_size * num_heads, 1, window_overlap, window_overlap),1068 )1069 < 11070 )1071 diagonal_attn_scores_low_triang = tf.where(1072 first_chunk_mask,1073 diagonal_attn_scores_first_chunk,1074 diagonal_attn_scores_low_triang,1075 )1076 1077 # merging upper and lower triangle1078 diagonal_attention_scores = tf.concat(1079 [diagonal_attn_scores_low_triang, diagonal_attn_scores_up_triang], axis=-11080 )1081 1082 # separate batch_size and num_heads dimensions again1083 diagonal_attention_scores = tf.transpose(1084 tf.reshape(1085 diagonal_attention_scores,1086 (batch_size, num_heads, seq_len, 2 * window_overlap + 1),1087 ),1088 (0, 2, 1, 3),1089 )1090 1091 diagonal_attention_scores = self._mask_invalid_locations(diagonal_attention_scores, window_overlap)1092 1093 return diagonal_attention_scores1094 1095 @staticmethod1096 def _mask_invalid_locations(input_tensor, window_overlap):1097 # create correct upper triangle bool mask1098 mask_2d_upper = tf.reverse(1099 tf.linalg.band_part(tf.ones(shape=(window_overlap, window_overlap + 1)), -1, 0),1100 axis=[0],1101 )1102 1103 # pad to full matrix1104 padding = tf.convert_to_tensor(1105 [[0, shape_list(input_tensor)[1] - window_overlap], [0, shape_list(input_tensor)[3] - window_overlap - 1]]1106 )1107 1108 # create lower mask1109 mask_2d = tf.pad(mask_2d_upper, padding)1110 1111 # combine with upper mask1112 mask_2d = mask_2d + tf.reverse(mask_2d, axis=[0, 1])1113 1114 # broadcast to full matrix1115 mask_4d = tf.tile(mask_2d[None, :, None, :], (shape_list(input_tensor)[0], 1, 1, 1))1116 1117 # inf tensor used for masking1118 inf_tensor = -float("inf") * tf.ones_like(input_tensor)1119 1120 # mask1121 input_tensor = tf.where(tf.math.greater(mask_4d, 0), inf_tensor, input_tensor)1122 1123 return input_tensor1124 1125 def _sliding_chunks_matmul_attn_probs_value(self, attn_probs, value, window_overlap):1126 """1127 Same as _sliding_chunks_query_key_matmul but for attn_probs and value tensors. Returned tensor will be of the1128 same shape as `attn_probs`1129 """1130 1131 batch_size, seq_len, num_heads, head_dim = shape_list(value)1132 1133 tf.debugging.assert_equal(1134 seq_len % (window_overlap * 2), 0, message="Seq_len has to be multiple of 2 * window_overlap"1135 )1136 tf.debugging.assert_equal(1137 shape_list(attn_probs)[:3],1138 shape_list(value)[:3],1139 message="value and attn_probs must have same dims (except head_dim)",1140 )1141 tf.debugging.assert_equal(1142 shape_list(attn_probs)[3],1143 2 * window_overlap + 1,1144 message="attn_probs last dim has to be 2 * window_overlap + 1",1145 )1146 1147 chunks_count = seq_len // window_overlap - 11148 1149 # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size 2 window overlap1150 chunked_attn_probs = tf.reshape(1151 tf.transpose(attn_probs, (0, 2, 1, 3)),1152 (1153 batch_size * num_heads,1154 seq_len // window_overlap,1155 window_overlap,1156 2 * window_overlap + 1,1157 ),1158 )1159 1160 # group batch_size and num_heads dimensions into one1161 value = tf.reshape(1162 tf.transpose(value, (0, 2, 1, 3)),1163 (batch_size * num_heads, seq_len, head_dim),1164 )1165 1166 # pad seq_len with w at the beginning of the sequence and another window overlap at the end1167 paddings = tf.convert_to_tensor([[0, 0], [window_overlap, window_overlap], [0, 0]])1168 padded_value = tf.pad(value, paddings, constant_values=-1)1169 1170 # chunk padded_value into chunks of size 3 window overlap and an overlap of size window overlap1171 frame_size = 3 * window_overlap * head_dim1172 frame_hop_size = (shape_list(padded_value)[1] * head_dim - frame_size) // chunks_count1173 chunked_value = tf.signal.frame(1174 tf.reshape(padded_value, (batch_size * num_heads, -1)),1175 frame_size,1176 frame_hop_size,1177 )1178 chunked_value = tf.reshape(1179 chunked_value,1180 (batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim),1181 )1182 1183 tf.debugging.assert_equal(1184 shape_list(chunked_value),1185 [batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim],1186 message="Chunked value has the wrong shape",1187 )1188 1189 chunked_attn_probs = self._pad_and_diagonalize(chunked_attn_probs)1190 context = tf.einsum("bcwd,bcdh->bcwh", chunked_attn_probs, chunked_value)1191 context = tf.transpose(1192 tf.reshape(context, (batch_size, num_heads, seq_len, head_dim)),1193 (0, 2, 1, 3),1194 )1195 1196 return context1197 1198 @staticmethod1199 def _pad_and_transpose_last_two_dims(hidden_states_padded, paddings):1200 """pads rows and then flips rows and columns"""