Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team. All rights reserved.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"""TF 2.0 ConvBERT model."""16 17from __future__ import annotations18 19import numpy as np20import tensorflow as tf21 22from ...activations_tf import get_tf_activation23from ...modeling_tf_outputs import (24 TFBaseModelOutput,25 TFMaskedLMOutput,26 TFMultipleChoiceModelOutput,27 TFQuestionAnsweringModelOutput,28 TFSequenceClassifierOutput,29 TFTokenClassifierOutput,30)31from ...modeling_tf_utils import (32 TFMaskedLanguageModelingLoss,33 TFModelInputType,34 TFMultipleChoiceLoss,35 TFPreTrainedModel,36 TFQuestionAnsweringLoss,37 TFSequenceClassificationLoss,38 TFSequenceSummary,39 TFTokenClassificationLoss,40 get_initializer,41 keras,42 keras_serializable,43 unpack_inputs,44)45from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax46from ...utils import (47 add_code_sample_docstrings,48 add_start_docstrings,49 add_start_docstrings_to_model_forward,50 logging,51)52from .configuration_convbert import ConvBertConfig53 54 55logger = logging.get_logger(__name__)56 57_CHECKPOINT_FOR_DOC = "YituTech/conv-bert-base"58_CONFIG_FOR_DOC = "ConvBertConfig"59 60 61# Copied from transformers.models.albert.modeling_tf_albert.TFAlbertEmbeddings with Albert->ConvBert62class TFConvBertEmbeddings(keras.layers.Layer):63 """Construct the embeddings from word, position and token_type embeddings."""64 65 def __init__(self, config: ConvBertConfig, **kwargs):66 super().__init__(**kwargs)67 68 self.config = config69 self.embedding_size = config.embedding_size70 self.max_position_embeddings = config.max_position_embeddings71 self.initializer_range = config.initializer_range72 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")73 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)74 75 def build(self, input_shape=None):76 with tf.name_scope("word_embeddings"):77 self.weight = self.add_weight(78 name="weight",79 shape=[self.config.vocab_size, self.embedding_size],80 initializer=get_initializer(self.initializer_range),81 )82 83 with tf.name_scope("token_type_embeddings"):84 self.token_type_embeddings = self.add_weight(85 name="embeddings",86 shape=[self.config.type_vocab_size, self.embedding_size],87 initializer=get_initializer(self.initializer_range),88 )89 90 with tf.name_scope("position_embeddings"):91 self.position_embeddings = self.add_weight(92 name="embeddings",93 shape=[self.max_position_embeddings, self.embedding_size],94 initializer=get_initializer(self.initializer_range),95 )96 97 if self.built:98 return99 self.built = True100 if getattr(self, "LayerNorm", None) is not None:101 with tf.name_scope(self.LayerNorm.name):102 self.LayerNorm.build([None, None, self.config.embedding_size])103 104 # Copied from transformers.models.bert.modeling_tf_bert.TFBertEmbeddings.call105 def call(106 self,107 input_ids: tf.Tensor | None = None,108 position_ids: tf.Tensor | None = None,109 token_type_ids: tf.Tensor | None = None,110 inputs_embeds: tf.Tensor | None = None,111 past_key_values_length=0,112 training: bool = False,113 ) -> tf.Tensor:114 """115 Applies embedding based on inputs tensor.116 117 Returns:118 final_embeddings (`tf.Tensor`): output embedding tensor.119 """120 if input_ids is None and inputs_embeds is None:121 raise ValueError("Need to provide either `input_ids` or `input_embeds`.")122 123 if input_ids is not None:124 check_embeddings_within_bounds(input_ids, self.config.vocab_size)125 inputs_embeds = tf.gather(params=self.weight, indices=input_ids)126 127 input_shape = shape_list(inputs_embeds)[:-1]128 129 if token_type_ids is None:130 token_type_ids = tf.fill(dims=input_shape, value=0)131 132 if position_ids is None:133 position_ids = tf.expand_dims(134 tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0135 )136 137 position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)138 token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)139 final_embeddings = inputs_embeds + position_embeds + token_type_embeds140 final_embeddings = self.LayerNorm(inputs=final_embeddings)141 final_embeddings = self.dropout(inputs=final_embeddings, training=training)142 143 return final_embeddings144 145 146class TFConvBertSelfAttention(keras.layers.Layer):147 def __init__(self, config, **kwargs):148 super().__init__(**kwargs)149 150 if config.hidden_size % config.num_attention_heads != 0:151 raise ValueError(152 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "153 f"heads ({config.num_attention_heads})"154 )155 156 new_num_attention_heads = int(config.num_attention_heads / config.head_ratio)157 if new_num_attention_heads < 1:158 self.head_ratio = config.num_attention_heads159 num_attention_heads = 1160 else:161 num_attention_heads = new_num_attention_heads162 self.head_ratio = config.head_ratio163 164 self.num_attention_heads = num_attention_heads165 self.conv_kernel_size = config.conv_kernel_size166 167 if config.hidden_size % self.num_attention_heads != 0:168 raise ValueError("hidden_size should be divisible by num_attention_heads")169 170 self.attention_head_size = config.hidden_size // config.num_attention_heads171 self.all_head_size = self.num_attention_heads * self.attention_head_size172 self.query = keras.layers.Dense(173 self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"174 )175 self.key = keras.layers.Dense(176 self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"177 )178 self.value = keras.layers.Dense(179 self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"180 )181 182 self.key_conv_attn_layer = keras.layers.SeparableConv1D(183 self.all_head_size,184 self.conv_kernel_size,185 padding="same",186 activation=None,187 depthwise_initializer=get_initializer(1 / self.conv_kernel_size),188 pointwise_initializer=get_initializer(config.initializer_range),189 name="key_conv_attn_layer",190 )191 192 self.conv_kernel_layer = keras.layers.Dense(193 self.num_attention_heads * self.conv_kernel_size,194 activation=None,195 name="conv_kernel_layer",196 kernel_initializer=get_initializer(config.initializer_range),197 )198 199 self.conv_out_layer = keras.layers.Dense(200 self.all_head_size,201 activation=None,202 name="conv_out_layer",203 kernel_initializer=get_initializer(config.initializer_range),204 )205 206 self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)207 self.config = config208 209 def transpose_for_scores(self, x, batch_size):210 # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]211 x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))212 return tf.transpose(x, perm=[0, 2, 1, 3])213 214 def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):215 batch_size = shape_list(hidden_states)[0]216 mixed_query_layer = self.query(hidden_states)217 mixed_key_layer = self.key(hidden_states)218 mixed_value_layer = self.value(hidden_states)219 220 mixed_key_conv_attn_layer = self.key_conv_attn_layer(hidden_states)221 222 query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)223 key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)224 conv_attn_layer = tf.multiply(mixed_key_conv_attn_layer, mixed_query_layer)225 226 conv_kernel_layer = self.conv_kernel_layer(conv_attn_layer)227 conv_kernel_layer = tf.reshape(conv_kernel_layer, [-1, self.conv_kernel_size, 1])228 conv_kernel_layer = stable_softmax(conv_kernel_layer, axis=1)229 230 paddings = tf.constant(231 [232 [233 0,234 0,235 ],236 [int((self.conv_kernel_size - 1) / 2), int((self.conv_kernel_size - 1) / 2)],237 [0, 0],238 ]239 )240 241 conv_out_layer = self.conv_out_layer(hidden_states)242 conv_out_layer = tf.reshape(conv_out_layer, [batch_size, -1, self.all_head_size])243 conv_out_layer = tf.pad(conv_out_layer, paddings, "CONSTANT")244 245 unfold_conv_out_layer = tf.stack(246 [247 tf.slice(conv_out_layer, [0, i, 0], [batch_size, shape_list(mixed_query_layer)[1], self.all_head_size])248 for i in range(self.conv_kernel_size)249 ],250 axis=-1,251 )252 253 conv_out_layer = tf.reshape(unfold_conv_out_layer, [-1, self.attention_head_size, self.conv_kernel_size])254 255 conv_out_layer = tf.matmul(conv_out_layer, conv_kernel_layer)256 conv_out_layer = tf.reshape(conv_out_layer, [-1, self.all_head_size])257 258 # Take the dot product between "query" and "key" to get the raw attention scores.259 attention_scores = tf.matmul(260 query_layer, key_layer, transpose_b=True261 ) # (batch size, num_heads, seq_len_q, seq_len_k)262 dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores263 attention_scores = attention_scores / tf.math.sqrt(dk)264 265 if attention_mask is not None:266 # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)267 attention_scores = attention_scores + attention_mask268 269 # Normalize the attention scores to probabilities.270 attention_probs = stable_softmax(attention_scores, axis=-1)271 272 # This is actually dropping out entire tokens to attend to, which might273 # seem a bit unusual, but is taken from the original Transformer paper.274 attention_probs = self.dropout(attention_probs, training=training)275 276 # Mask heads if we want to277 if head_mask is not None:278 attention_probs = attention_probs * head_mask279 280 value_layer = tf.reshape(281 mixed_value_layer, [batch_size, -1, self.num_attention_heads, self.attention_head_size]282 )283 value_layer = tf.transpose(value_layer, [0, 2, 1, 3])284 285 context_layer = tf.matmul(attention_probs, value_layer)286 context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])287 288 conv_out = tf.reshape(conv_out_layer, [batch_size, -1, self.num_attention_heads, self.attention_head_size])289 context_layer = tf.concat([context_layer, conv_out], 2)290 context_layer = tf.reshape(291 context_layer, (batch_size, -1, self.head_ratio * self.all_head_size)292 ) # (batch_size, seq_len_q, all_head_size)293 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)294 295 return outputs296 297 def build(self, input_shape=None):298 if self.built:299 return300 self.built = True301 if getattr(self, "query", None) is not None:302 with tf.name_scope(self.query.name):303 self.query.build([None, None, self.config.hidden_size])304 if getattr(self, "key", None) is not None:305 with tf.name_scope(self.key.name):306 self.key.build([None, None, self.config.hidden_size])307 if getattr(self, "value", None) is not None:308 with tf.name_scope(self.value.name):309 self.value.build([None, None, self.config.hidden_size])310 if getattr(self, "key_conv_attn_layer", None) is not None:311 with tf.name_scope(self.key_conv_attn_layer.name):312 self.key_conv_attn_layer.build([None, None, self.config.hidden_size])313 if getattr(self, "conv_kernel_layer", None) is not None:314 with tf.name_scope(self.conv_kernel_layer.name):315 self.conv_kernel_layer.build([None, None, self.all_head_size])316 if getattr(self, "conv_out_layer", None) is not None:317 with tf.name_scope(self.conv_out_layer.name):318 self.conv_out_layer.build([None, None, self.config.hidden_size])319 320 321class TFConvBertSelfOutput(keras.layers.Layer):322 def __init__(self, config, **kwargs):323 super().__init__(**kwargs)324 325 self.dense = keras.layers.Dense(326 config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"327 )328 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")329 self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)330 self.config = config331 332 def call(self, hidden_states, input_tensor, training=False):333 hidden_states = self.dense(hidden_states)334 hidden_states = self.dropout(hidden_states, training=training)335 hidden_states = self.LayerNorm(hidden_states + input_tensor)336 337 return hidden_states338 339 def build(self, input_shape=None):340 if self.built:341 return342 self.built = True343 if getattr(self, "dense", None) is not None:344 with tf.name_scope(self.dense.name):345 self.dense.build([None, None, self.config.hidden_size])346 if getattr(self, "LayerNorm", None) is not None:347 with tf.name_scope(self.LayerNorm.name):348 self.LayerNorm.build([None, None, self.config.hidden_size])349 350 351class TFConvBertAttention(keras.layers.Layer):352 def __init__(self, config, **kwargs):353 super().__init__(**kwargs)354 355 self.self_attention = TFConvBertSelfAttention(config, name="self")356 self.dense_output = TFConvBertSelfOutput(config, name="output")357 358 def prune_heads(self, heads):359 raise NotImplementedError360 361 def call(self, input_tensor, attention_mask, head_mask, output_attentions, training=False):362 self_outputs = self.self_attention(363 input_tensor, attention_mask, head_mask, output_attentions, training=training364 )365 attention_output = self.dense_output(self_outputs[0], input_tensor, training=training)366 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them367 368 return outputs369 370 def build(self, input_shape=None):371 if self.built:372 return373 self.built = True374 if getattr(self, "self_attention", None) is not None:375 with tf.name_scope(self.self_attention.name):376 self.self_attention.build(None)377 if getattr(self, "dense_output", None) is not None:378 with tf.name_scope(self.dense_output.name):379 self.dense_output.build(None)380 381 382class GroupedLinearLayer(keras.layers.Layer):383 def __init__(self, input_size, output_size, num_groups, kernel_initializer, **kwargs):384 super().__init__(**kwargs)385 self.input_size = input_size386 self.output_size = output_size387 self.num_groups = num_groups388 self.kernel_initializer = kernel_initializer389 self.group_in_dim = self.input_size // self.num_groups390 self.group_out_dim = self.output_size // self.num_groups391 392 def build(self, input_shape=None):393 self.kernel = self.add_weight(394 "kernel",395 shape=[self.group_out_dim, self.group_in_dim, self.num_groups],396 initializer=self.kernel_initializer,397 trainable=True,398 )399 400 self.bias = self.add_weight(401 "bias", shape=[self.output_size], initializer=self.kernel_initializer, dtype=self.dtype, trainable=True402 )403 super().build(input_shape)404 405 def call(self, hidden_states):406 batch_size = shape_list(hidden_states)[0]407 x = tf.transpose(tf.reshape(hidden_states, [-1, self.num_groups, self.group_in_dim]), [1, 0, 2])408 x = tf.matmul(x, tf.transpose(self.kernel, [2, 1, 0]))409 x = tf.transpose(x, [1, 0, 2])410 x = tf.reshape(x, [batch_size, -1, self.output_size])411 x = tf.nn.bias_add(value=x, bias=self.bias)412 return x413 414 415class TFConvBertIntermediate(keras.layers.Layer):416 def __init__(self, config, **kwargs):417 super().__init__(**kwargs)418 if config.num_groups == 1:419 self.dense = keras.layers.Dense(420 config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"421 )422 else:423 self.dense = GroupedLinearLayer(424 config.hidden_size,425 config.intermediate_size,426 num_groups=config.num_groups,427 kernel_initializer=get_initializer(config.initializer_range),428 name="dense",429 )430 431 if isinstance(config.hidden_act, str):432 self.intermediate_act_fn = get_tf_activation(config.hidden_act)433 else:434 self.intermediate_act_fn = config.hidden_act435 self.config = config436 437 def call(self, hidden_states):438 hidden_states = self.dense(hidden_states)439 hidden_states = self.intermediate_act_fn(hidden_states)440 441 return hidden_states442 443 def build(self, input_shape=None):444 if self.built:445 return446 self.built = True447 if getattr(self, "dense", None) is not None:448 with tf.name_scope(self.dense.name):449 self.dense.build([None, None, self.config.hidden_size])450 451 452class TFConvBertOutput(keras.layers.Layer):453 def __init__(self, config, **kwargs):454 super().__init__(**kwargs)455 456 if config.num_groups == 1:457 self.dense = keras.layers.Dense(458 config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"459 )460 else:461 self.dense = GroupedLinearLayer(462 config.intermediate_size,463 config.hidden_size,464 num_groups=config.num_groups,465 kernel_initializer=get_initializer(config.initializer_range),466 name="dense",467 )468 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")469 self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)470 self.config = config471 472 def call(self, hidden_states, input_tensor, training=False):473 hidden_states = self.dense(hidden_states)474 hidden_states = self.dropout(hidden_states, training=training)475 hidden_states = self.LayerNorm(hidden_states + input_tensor)476 477 return hidden_states478 479 def build(self, input_shape=None):480 if self.built:481 return482 self.built = True483 if getattr(self, "LayerNorm", None) is not None:484 with tf.name_scope(self.LayerNorm.name):485 self.LayerNorm.build([None, None, self.config.hidden_size])486 if getattr(self, "dense", None) is not None:487 with tf.name_scope(self.dense.name):488 self.dense.build([None, None, self.config.intermediate_size])489 490 491class TFConvBertLayer(keras.layers.Layer):492 def __init__(self, config, **kwargs):493 super().__init__(**kwargs)494 495 self.attention = TFConvBertAttention(config, name="attention")496 self.intermediate = TFConvBertIntermediate(config, name="intermediate")497 self.bert_output = TFConvBertOutput(config, name="output")498 499 def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):500 attention_outputs = self.attention(501 hidden_states, attention_mask, head_mask, output_attentions, training=training502 )503 attention_output = attention_outputs[0]504 intermediate_output = self.intermediate(attention_output)505 layer_output = self.bert_output(intermediate_output, attention_output, training=training)506 outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them507 508 return outputs509 510 def build(self, input_shape=None):511 if self.built:512 return513 self.built = True514 if getattr(self, "attention", None) is not None:515 with tf.name_scope(self.attention.name):516 self.attention.build(None)517 if getattr(self, "intermediate", None) is not None:518 with tf.name_scope(self.intermediate.name):519 self.intermediate.build(None)520 if getattr(self, "bert_output", None) is not None:521 with tf.name_scope(self.bert_output.name):522 self.bert_output.build(None)523 524 525class TFConvBertEncoder(keras.layers.Layer):526 def __init__(self, config, **kwargs):527 super().__init__(**kwargs)528 529 self.layer = [TFConvBertLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]530 531 def call(532 self,533 hidden_states,534 attention_mask,535 head_mask,536 output_attentions,537 output_hidden_states,538 return_dict,539 training=False,540 ):541 all_hidden_states = () if output_hidden_states else None542 all_attentions = () if output_attentions else None543 544 for i, layer_module in enumerate(self.layer):545 if output_hidden_states:546 all_hidden_states = all_hidden_states + (hidden_states,)547 548 layer_outputs = layer_module(549 hidden_states, attention_mask, head_mask[i], output_attentions, training=training550 )551 hidden_states = layer_outputs[0]552 553 if output_attentions:554 all_attentions = all_attentions + (layer_outputs[1],)555 556 # Add last layer557 if output_hidden_states:558 all_hidden_states = all_hidden_states + (hidden_states,)559 560 if not return_dict:561 return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)562 563 return TFBaseModelOutput(564 last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions565 )566 567 def build(self, input_shape=None):568 if self.built:569 return570 self.built = True571 if getattr(self, "layer", None) is not None:572 for layer in self.layer:573 with tf.name_scope(layer.name):574 layer.build(None)575 576 577class TFConvBertPredictionHeadTransform(keras.layers.Layer):578 def __init__(self, config, **kwargs):579 super().__init__(**kwargs)580 581 self.dense = keras.layers.Dense(582 config.embedding_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"583 )584 585 if isinstance(config.hidden_act, str):586 self.transform_act_fn = get_tf_activation(config.hidden_act)587 else:588 self.transform_act_fn = config.hidden_act589 590 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")591 self.config = config592 593 def call(self, hidden_states):594 hidden_states = self.dense(hidden_states)595 hidden_states = self.transform_act_fn(hidden_states)596 hidden_states = self.LayerNorm(hidden_states)597 598 return hidden_states599 600 def build(self, input_shape=None):601 if self.built:602 return603 self.built = True604 if getattr(self, "dense", None) is not None:605 with tf.name_scope(self.dense.name):606 self.dense.build([None, None, self.config.hidden_size])607 if getattr(self, "LayerNorm", None) is not None:608 with tf.name_scope(self.LayerNorm.name):609 self.LayerNorm.build([None, None, self.config.hidden_size])610 611 612@keras_serializable613class TFConvBertMainLayer(keras.layers.Layer):614 config_class = ConvBertConfig615 616 def __init__(self, config, **kwargs):617 super().__init__(**kwargs)618 619 self.embeddings = TFConvBertEmbeddings(config, name="embeddings")620 621 if config.embedding_size != config.hidden_size:622 self.embeddings_project = keras.layers.Dense(config.hidden_size, name="embeddings_project")623 624 self.encoder = TFConvBertEncoder(config, name="encoder")625 self.config = config626 627 def get_input_embeddings(self):628 return self.embeddings629 630 def set_input_embeddings(self, value):631 self.embeddings.weight = value632 self.embeddings.vocab_size = value.shape[0]633 634 def _prune_heads(self, heads_to_prune):635 """636 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base637 class PreTrainedModel638 """639 raise NotImplementedError640 641 def get_extended_attention_mask(self, attention_mask, input_shape, dtype):642 if attention_mask is None:643 attention_mask = tf.fill(input_shape, 1)644 645 # We create a 3D attention mask from a 2D tensor mask.646 # Sizes are [batch_size, 1, 1, to_seq_length]647 # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]648 # this attention mask is more simple than the triangular masking of causal attention649 # used in OpenAI GPT, we just need to prepare the broadcast dimension here.650 extended_attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1]))651 652 # Since attention_mask is 1.0 for positions we want to attend and 0.0 for653 # masked positions, this operation will create a tensor which is 0.0 for654 # positions we want to attend and -10000.0 for masked positions.655 # Since we are adding it to the raw scores before the softmax, this is656 # effectively the same as removing these entirely.657 extended_attention_mask = tf.cast(extended_attention_mask, dtype)658 extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0659 660 return extended_attention_mask661 662 def get_head_mask(self, head_mask):663 if head_mask is not None:664 raise NotImplementedError665 else:666 head_mask = [None] * self.config.num_hidden_layers667 668 return head_mask669 670 @unpack_inputs671 def call(672 self,673 input_ids=None,674 attention_mask=None,675 token_type_ids=None,676 position_ids=None,677 head_mask=None,678 inputs_embeds=None,679 output_attentions=None,680 output_hidden_states=None,681 return_dict=None,682 training=False,683 ):684 if input_ids is not None and inputs_embeds is not None:685 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")686 elif input_ids is not None:687 input_shape = shape_list(input_ids)688 elif inputs_embeds is not None:689 input_shape = shape_list(inputs_embeds)[:-1]690 else:691 raise ValueError("You have to specify either input_ids or inputs_embeds")692 693 if attention_mask is None:694 attention_mask = tf.fill(input_shape, 1)695 696 if token_type_ids is None:697 token_type_ids = tf.fill(input_shape, 0)698 699 hidden_states = self.embeddings(input_ids, position_ids, token_type_ids, inputs_embeds, training=training)700 extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, hidden_states.dtype)701 head_mask = self.get_head_mask(head_mask)702 703 if hasattr(self, "embeddings_project"):704 hidden_states = self.embeddings_project(hidden_states, training=training)705 706 hidden_states = self.encoder(707 hidden_states,708 extended_attention_mask,709 head_mask,710 output_attentions,711 output_hidden_states,712 return_dict,713 training=training,714 )715 716 return hidden_states717 718 def build(self, input_shape=None):719 if self.built:720 return721 self.built = True722 if getattr(self, "embeddings", None) is not None:723 with tf.name_scope(self.embeddings.name):724 self.embeddings.build(None)725 if getattr(self, "encoder", None) is not None:726 with tf.name_scope(self.encoder.name):727 self.encoder.build(None)728 if getattr(self, "embeddings_project", None) is not None:729 with tf.name_scope(self.embeddings_project.name):730 self.embeddings_project.build([None, None, self.config.embedding_size])731 732 733class TFConvBertPreTrainedModel(TFPreTrainedModel):734 """735 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained736 models.737 """738 739 config_class = ConvBertConfig740 base_model_prefix = "convbert"741 742 743CONVBERT_START_DOCSTRING = r"""744 745 This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the746 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads747 etc.)748 749 This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it750 as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and751 behavior.752 753 <Tip>754 755 TensorFlow models and layers in `transformers` accept two formats as input:756 757 - having all inputs as keyword arguments (like PyTorch models), or758 - having all inputs as a list, tuple or dict in the first positional argument.759 760 The reason the second format is supported is that Keras methods prefer this format when passing inputs to models761 and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just762 pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second763 format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with764 the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first765 positional argument:766 767 - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`768 - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:769 `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`770 - a dictionary with one or several input Tensors associated to the input names given in the docstring:771 `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`772 773 Note that when creating models and layers with774 [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry775 about any of this, as you can just pass inputs like you would to any other Python function!776 777 </Tip>778 779 Args:780 config ([`ConvBertConfig`]): Model configuration class with all the parameters of the model.781 Initializing with a config file does not load the weights associated with the model, only the782 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.783"""784 785CONVBERT_INPUTS_DOCSTRING = r"""786 Args:787 input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):788 Indices of input sequence tokens in the vocabulary.789 790 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and791 [`PreTrainedTokenizer.encode`] for details.792 793 [What are input IDs?](../glossary#input-ids)794 attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):795 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:796 797 - 1 for tokens that are **not masked**,798 - 0 for tokens that are **masked**.799 800 [What are attention masks?](../glossary#attention-mask)801 token_type_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):802 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,803 1]`:804 805 - 0 corresponds to a *sentence A* token,806 - 1 corresponds to a *sentence B* token.807 808 [What are token type IDs?](../glossary#token-type-ids)809 position_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):810 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,811 config.max_position_embeddings - 1]`.812 813 [What are position IDs?](../glossary#position-ids)814 head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):815 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:816 817 - 1 indicates the head is **not masked**,818 - 0 indicates the head is **masked**.819 820 inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):821 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This822 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the823 model's internal embedding lookup matrix.824 output_attentions (`bool`, *optional*):825 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned826 tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the827 config will be used instead.828 output_hidden_states (`bool`, *optional*):829 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for830 more detail. This argument can be used only in eager mode, in graph mode the value in the config will be831 used instead.832 return_dict (`bool`, *optional*):833 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in834 eager mode, in graph mode the value will always be set to True.835 training (`bool`, *optional*, defaults to `False`):836 Whether or not to use the model in training mode (some modules like dropout modules have different837 behaviors between training and evaluation).838"""839 840 841@add_start_docstrings(842 "The bare ConvBERT Model transformer outputting raw hidden-states without any specific head on top.",843 CONVBERT_START_DOCSTRING,844)845class TFConvBertModel(TFConvBertPreTrainedModel):846 def __init__(self, config, *inputs, **kwargs):847 super().__init__(config, *inputs, **kwargs)848 849 self.convbert = TFConvBertMainLayer(config, name="convbert")850 851 @unpack_inputs852 @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))853 @add_code_sample_docstrings(854 checkpoint=_CHECKPOINT_FOR_DOC,855 output_type=TFBaseModelOutput,856 config_class=_CONFIG_FOR_DOC,857 )858 def call(859 self,860 input_ids: TFModelInputType | None = None,861 attention_mask: np.array | tf.Tensor | None = None,862 token_type_ids: np.array | tf.Tensor | None = None,863 position_ids: np.array | tf.Tensor | None = None,864 head_mask: np.array | tf.Tensor | None = None,865 inputs_embeds: tf.Tensor | None = None,866 output_attentions: bool | None = None,867 output_hidden_states: bool | None = None,868 return_dict: bool | None = None,869 training: bool = False,870 ) -> TFBaseModelOutput | tuple[tf.Tensor]:871 outputs = self.convbert(872 input_ids=input_ids,873 attention_mask=attention_mask,874 token_type_ids=token_type_ids,875 position_ids=position_ids,876 head_mask=head_mask,877 inputs_embeds=inputs_embeds,878 output_attentions=output_attentions,879 output_hidden_states=output_hidden_states,880 return_dict=return_dict,881 training=training,882 )883 884 return outputs885 886 def build(self, input_shape=None):887 if self.built:888 return889 self.built = True890 if getattr(self, "convbert", None) is not None:891 with tf.name_scope(self.convbert.name):892 self.convbert.build(None)893 894 895class TFConvBertMaskedLMHead(keras.layers.Layer):896 def __init__(self, config, input_embeddings, **kwargs):897 super().__init__(**kwargs)898 899 self.config = config900 self.embedding_size = config.embedding_size901 self.input_embeddings = input_embeddings902 903 def build(self, input_shape):904 self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")905 906 super().build(input_shape)907 908 def get_output_embeddings(self):909 return self.input_embeddings910 911 def set_output_embeddings(self, value):912 self.input_embeddings.weight = value913 self.input_embeddings.vocab_size = shape_list(value)[0]914 915 def get_bias(self):916 return {"bias": self.bias}917 918 def set_bias(self, value):919 self.bias = value["bias"]920 self.config.vocab_size = shape_list(value["bias"])[0]921 922 def call(self, hidden_states):923 seq_length = shape_list(tensor=hidden_states)[1]924 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.embedding_size])925 hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)926 hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])927 hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)928 929 return hidden_states930 931 932class TFConvBertGeneratorPredictions(keras.layers.Layer):933 def __init__(self, config, **kwargs):934 super().__init__(**kwargs)935 936 self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")937 self.dense = keras.layers.Dense(config.embedding_size, name="dense")938 self.config = config939 940 def call(self, generator_hidden_states, training=False):941 hidden_states = self.dense(generator_hidden_states)942 hidden_states = get_tf_activation("gelu")(hidden_states)943 hidden_states = self.LayerNorm(hidden_states)944 945 return hidden_states946 947 def build(self, input_shape=None):948 if self.built:949 return950 self.built = True951 if getattr(self, "LayerNorm", None) is not None:952 with tf.name_scope(self.LayerNorm.name):953 self.LayerNorm.build([None, None, self.config.embedding_size])954 if getattr(self, "dense", None) is not None:955 with tf.name_scope(self.dense.name):956 self.dense.build([None, None, self.config.hidden_size])957 958 959@add_start_docstrings("""ConvBERT Model with a `language modeling` head on top.""", CONVBERT_START_DOCSTRING)960class TFConvBertForMaskedLM(TFConvBertPreTrainedModel, TFMaskedLanguageModelingLoss):961 def __init__(self, config, *inputs, **kwargs):962 super().__init__(config, **kwargs)963 964 self.config = config965 self.convbert = TFConvBertMainLayer(config, name="convbert")966 self.generator_predictions = TFConvBertGeneratorPredictions(config, name="generator_predictions")967 968 if isinstance(config.hidden_act, str):969 self.activation = get_tf_activation(config.hidden_act)970 else:971 self.activation = config.hidden_act972 973 self.generator_lm_head = TFConvBertMaskedLMHead(config, self.convbert.embeddings, name="generator_lm_head")974 975 def get_lm_head(self):976 return self.generator_lm_head977 978 def get_prefix_bias_name(self):979 return self.name + "/" + self.generator_lm_head.name980 981 @unpack_inputs982 @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))983 @add_code_sample_docstrings(984 checkpoint=_CHECKPOINT_FOR_DOC,985 output_type=TFMaskedLMOutput,986 config_class=_CONFIG_FOR_DOC,987 )988 def call(989 self,990 input_ids: TFModelInputType | None = None,991 attention_mask: np.ndarray | tf.Tensor | None = None,992 token_type_ids: np.ndarray | tf.Tensor | None = None,993 position_ids: np.ndarray | tf.Tensor | None = None,994 head_mask: np.ndarray | tf.Tensor | None = None,995 inputs_embeds: tf.Tensor | None = None,996 output_attentions: bool | None = None,997 output_hidden_states: bool | None = None,998 return_dict: bool | None = None,999 labels: tf.Tensor | None = None,1000 training: bool | None = False,1001 ) -> tuple | TFMaskedLMOutput:1002 r"""1003 labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1004 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1005 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the1006 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1007 """1008 generator_hidden_states = self.convbert(1009 input_ids=input_ids,1010 attention_mask=attention_mask,1011 token_type_ids=token_type_ids,1012 position_ids=position_ids,1013 head_mask=head_mask,1014 inputs_embeds=inputs_embeds,1015 output_attentions=output_attentions,1016 output_hidden_states=output_hidden_states,1017 return_dict=return_dict,1018 training=training,1019 )1020 generator_sequence_output = generator_hidden_states[0]1021 prediction_scores = self.generator_predictions(generator_sequence_output, training=training)1022 prediction_scores = self.generator_lm_head(prediction_scores, training=training)1023 loss = None if labels is None else self.hf_compute_loss(labels, prediction_scores)1024 1025 if not return_dict:1026 output = (prediction_scores,) + generator_hidden_states[1:]1027 1028 return ((loss,) + output) if loss is not None else output1029 1030 return TFMaskedLMOutput(1031 loss=loss,1032 logits=prediction_scores,1033 hidden_states=generator_hidden_states.hidden_states,1034 attentions=generator_hidden_states.attentions,1035 )1036 1037 def build(self, input_shape=None):1038 if self.built:1039 return1040 self.built = True1041 if getattr(self, "convbert", None) is not None:1042 with tf.name_scope(self.convbert.name):1043 self.convbert.build(None)1044 if getattr(self, "generator_predictions", None) is not None:1045 with tf.name_scope(self.generator_predictions.name):1046 self.generator_predictions.build(None)1047 if getattr(self, "generator_lm_head", None) is not None:1048 with tf.name_scope(self.generator_lm_head.name):1049 self.generator_lm_head.build(None)1050 1051 1052class TFConvBertClassificationHead(keras.layers.Layer):1053 """Head for sentence-level classification tasks."""1054 1055 def __init__(self, config, **kwargs):1056 super().__init__(**kwargs)1057 1058 self.dense = keras.layers.Dense(1059 config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"1060 )1061 classifier_dropout = (1062 config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob1063 )1064 self.dropout = keras.layers.Dropout(classifier_dropout)1065 self.out_proj = keras.layers.Dense(1066 config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="out_proj"1067 )1068 1069 self.config = config1070 1071 def call(self, hidden_states, **kwargs):1072 x = hidden_states[:, 0, :] # take <s> token (equiv. to [CLS])1073 x = self.dropout(x)1074 x = self.dense(x)1075 x = get_tf_activation(self.config.hidden_act)(x)1076 x = self.dropout(x)1077 x = self.out_proj(x)1078 1079 return x1080 1081 def build(self, input_shape=None):1082 if self.built:1083 return1084 self.built = True1085 if getattr(self, "dense", None) is not None:1086 with tf.name_scope(self.dense.name):1087 self.dense.build([None, None, self.config.hidden_size])1088 if getattr(self, "out_proj", None) is not None:1089 with tf.name_scope(self.out_proj.name):1090 self.out_proj.build([None, None, self.config.hidden_size])1091 1092 1093@add_start_docstrings(1094 """1095 ConvBERT Model transformer with a sequence classification/regression head on top e.g., for GLUE tasks.1096 """,1097 CONVBERT_START_DOCSTRING,1098)1099class TFConvBertForSequenceClassification(TFConvBertPreTrainedModel, TFSequenceClassificationLoss):1100 def __init__(self, config, *inputs, **kwargs):1101 super().__init__(config, *inputs, **kwargs)1102 self.num_labels = config.num_labels1103 self.convbert = TFConvBertMainLayer(config, name="convbert")1104 self.classifier = TFConvBertClassificationHead(config, name="classifier")1105 1106 @unpack_inputs1107 @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))1108 @add_code_sample_docstrings(1109 checkpoint=_CHECKPOINT_FOR_DOC,1110 output_type=TFSequenceClassifierOutput,1111 config_class=_CONFIG_FOR_DOC,1112 )1113 def call(1114 self,1115 input_ids: TFModelInputType | None = None,1116 attention_mask: np.ndarray | tf.Tensor | None = None,1117 token_type_ids: np.ndarray | tf.Tensor | None = None,1118 position_ids: np.ndarray | tf.Tensor | None = None,1119 head_mask: np.ndarray | tf.Tensor | None = None,1120 inputs_embeds: tf.Tensor | None = None,1121 output_attentions: bool | None = None,1122 output_hidden_states: bool | None = None,1123 return_dict: bool | None = None,1124 labels: tf.Tensor | None = None,1125 training: bool | None = False,1126 ) -> tuple | TFSequenceClassifierOutput:1127 r"""1128 labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):1129 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1130 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1131 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1132 """1133 outputs = self.convbert(1134 input_ids,1135 attention_mask=attention_mask,1136 token_type_ids=token_type_ids,1137 position_ids=position_ids,1138 head_mask=head_mask,1139 inputs_embeds=inputs_embeds,1140 output_attentions=output_attentions,1141 output_hidden_states=output_hidden_states,1142 return_dict=return_dict,1143 training=training,1144 )1145 logits = self.classifier(outputs[0], training=training)1146 loss = None if labels is None else self.hf_compute_loss(labels, logits)1147 1148 if not return_dict:1149 output = (logits,) + outputs[1:]1150 1151 return ((loss,) + output) if loss is not None else output1152 1153 return TFSequenceClassifierOutput(1154 loss=loss,1155 logits=logits,1156 hidden_states=outputs.hidden_states,1157 attentions=outputs.attentions,1158 )1159 1160 def build(self, input_shape=None):1161 if self.built:1162 return1163 self.built = True1164 if getattr(self, "convbert", None) is not None:1165 with tf.name_scope(self.convbert.name):1166 self.convbert.build(None)1167 if getattr(self, "classifier", None) is not None:1168 with tf.name_scope(self.classifier.name):1169 self.classifier.build(None)1170 1171 1172@add_start_docstrings(1173 """1174 ConvBERT Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a1175 softmax) e.g. for RocStories/SWAG tasks.1176 """,1177 CONVBERT_START_DOCSTRING,1178)1179class TFConvBertForMultipleChoice(TFConvBertPreTrainedModel, TFMultipleChoiceLoss):1180 def __init__(self, config, *inputs, **kwargs):1181 super().__init__(config, *inputs, **kwargs)1182 1183 self.convbert = TFConvBertMainLayer(config, name="convbert")1184 self.sequence_summary = TFSequenceSummary(1185 config, initializer_range=config.initializer_range, name="sequence_summary"1186 )1187 self.classifier = keras.layers.Dense(1188 1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"1189 )1190 self.config = config1191 1192 @unpack_inputs1193 @add_start_docstrings_to_model_forward(1194 CONVBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")1195 )1196 @add_code_sample_docstrings(1197 checkpoint=_CHECKPOINT_FOR_DOC,1198 output_type=TFMultipleChoiceModelOutput,1199 config_class=_CONFIG_FOR_DOC,1200 )