Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2018 Salesforce and HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""TF 2.0 CTRL model."""17 18from __future__ import annotations19 20import numpy as np21import tensorflow as tf22 23from ...modeling_tf_outputs import TFBaseModelOutputWithPast, TFCausalLMOutputWithPast, TFSequenceClassifierOutput24from ...modeling_tf_utils import (25 TFCausalLanguageModelingLoss,26 TFModelInputType,27 TFPreTrainedModel,28 TFSequenceClassificationLoss,29 get_initializer,30 keras,31 keras_serializable,32 unpack_inputs,33)34from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax35from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging36from .configuration_ctrl import CTRLConfig37 38 39logger = logging.get_logger(__name__)40 41_CHECKPOINT_FOR_DOC = "Salesforce/ctrl"42_CONFIG_FOR_DOC = "CTRLConfig"43 44 45def angle_defn(pos, i, d_model_size):46 angle_rates = 1 / np.power(10000, (2 * (i // 2)) / d_model_size)47 return pos * angle_rates48 49 50def positional_encoding(position, d_model_size):51 # create the sinusoidal pattern for the positional encoding52 angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)53 54 sines = np.sin(angle_rads[:, 0::2])55 cosines = np.cos(angle_rads[:, 1::2])56 pos_encoding = tf.convert_to_tensor(np.concatenate([sines, cosines], axis=-1))57 58 return pos_encoding59 60 61def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):62 # calculate attention63 matmul_qk = tf.matmul(q, k, transpose_b=True)64 65 dk = tf.cast(shape_list(k)[-1], dtype=matmul_qk.dtype)66 scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)67 68 if mask is not None:69 scaled_attention_logits += tf.cast(mask * -1e4, dtype=scaled_attention_logits.dtype)70 71 if attention_mask is not None:72 # Apply the attention mask73 attention_mask = tf.cast(attention_mask, dtype=scaled_attention_logits.dtype)74 scaled_attention_logits = scaled_attention_logits + attention_mask75 76 attention_weights = stable_softmax(scaled_attention_logits, axis=-1)77 78 # Mask heads if we want to79 if head_mask is not None:80 attention_weights = attention_weights * head_mask81 82 output = tf.matmul(attention_weights, v)83 84 return output, attention_weights85 86 87class TFMultiHeadAttention(keras.layers.Layer):88 def __init__(self, d_model_size, num_heads, output_attentions=False, **kwargs):89 super().__init__(**kwargs)90 self.num_heads = num_heads91 self.d_model_size = d_model_size92 self.output_attentions = output_attentions93 94 self.depth = int(d_model_size / self.num_heads)95 96 self.Wq = keras.layers.Dense(d_model_size, name="Wq")97 self.Wk = keras.layers.Dense(d_model_size, name="Wk")98 self.Wv = keras.layers.Dense(d_model_size, name="Wv")99 100 self.dense = keras.layers.Dense(d_model_size, name="dense")101 102 def split_into_heads(self, x, batch_size):103 x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))104 return tf.transpose(x, perm=[0, 2, 1, 3])105 106 def call(self, v, k, q, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions, training=False):107 batch_size = shape_list(q)[0]108 109 q = self.Wq(q)110 k = self.Wk(k)111 v = self.Wv(v)112 113 q = self.split_into_heads(q, batch_size)114 k = self.split_into_heads(k, batch_size)115 v = self.split_into_heads(v, batch_size)116 117 if layer_past is not None:118 past_key, past_value = tf.unstack(layer_past, axis=0)119 k = tf.concat((past_key, k), axis=-2)120 v = tf.concat((past_value, v), axis=-2)121 122 if use_cache:123 present = tf.stack((k, v), axis=0)124 else:125 present = (None,)126 127 output = scaled_dot_product_attention(q, k, v, mask, attention_mask, head_mask)128 scaled_attention = tf.transpose(output[0], perm=[0, 2, 1, 3])129 attn = output[1]130 original_size_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model_size))131 output = self.dense(original_size_attention)132 outputs = (output, present)133 134 if output_attentions:135 outputs = outputs + (attn,)136 137 return outputs138 139 def build(self, input_shape=None):140 if self.built:141 return142 self.built = True143 if getattr(self, "Wq", None) is not None:144 with tf.name_scope(self.Wq.name):145 self.Wq.build([None, None, self.d_model_size])146 if getattr(self, "Wk", None) is not None:147 with tf.name_scope(self.Wk.name):148 self.Wk.build([None, None, self.d_model_size])149 if getattr(self, "Wv", None) is not None:150 with tf.name_scope(self.Wv.name):151 self.Wv.build([None, None, self.d_model_size])152 if getattr(self, "dense", None) is not None:153 with tf.name_scope(self.dense.name):154 self.dense.build([None, None, self.d_model_size])155 156 157class TFPointWiseFeedForwardLayer(keras.layers.Layer):158 def __init__(self, d_model_size, dff, **kwargs):159 super().__init__(**kwargs)160 161 self.dense_0 = keras.layers.Dense(dff, activation="relu", name="0")162 self.dense_2 = keras.layers.Dense(d_model_size, name="2")163 self.d_model_size = d_model_size164 self.dff = dff165 166 def call(self, inputs, trainable=False):167 dense_0_output = self.dense_0(inputs)168 dense_2_output = self.dense_2(dense_0_output)169 170 return dense_2_output171 172 def build(self, input_shape=None):173 if self.built:174 return175 self.built = True176 if getattr(self, "dense_0", None) is not None:177 with tf.name_scope(self.dense_0.name):178 self.dense_0.build([None, None, self.d_model_size])179 if getattr(self, "dense_2", None) is not None:180 with tf.name_scope(self.dense_2.name):181 self.dense_2.build([None, None, self.dff])182 183 184class TFEncoderLayer(keras.layers.Layer):185 def __init__(186 self, d_model_size, num_heads, dff, rate=0.1, layer_norm_epsilon=1e-6, output_attentions=False, **kwargs187 ):188 super().__init__(**kwargs)189 190 self.output_attentions = output_attentions191 192 self.multi_head_attention = TFMultiHeadAttention(193 d_model_size, num_heads, output_attentions=self.output_attentions, name="multi_head_attention"194 )195 self.ffn = TFPointWiseFeedForwardLayer(d_model_size, dff, name="ffn")196 197 self.layernorm1 = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm1")198 self.layernorm2 = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm2")199 200 self.dropout1 = keras.layers.Dropout(rate)201 self.dropout2 = keras.layers.Dropout(rate)202 self.d_model_size = d_model_size203 204 def call(self, x, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions, training=False):205 normed = self.layernorm1(x)206 attn_outputs = self.multi_head_attention(207 normed,208 normed,209 normed,210 mask,211 layer_past,212 attention_mask,213 head_mask,214 use_cache,215 output_attentions,216 training=training,217 )218 attn_output = attn_outputs[0]219 attn_output = self.dropout1(attn_output, training=training)220 out1 = x + attn_output221 222 out2 = self.layernorm2(out1)223 ffn_output = self.ffn(out2)224 ffn_output = self.dropout2(ffn_output, training=training)225 out2 = out1 + ffn_output226 227 outputs = (out2,) + attn_outputs[1:]228 return outputs229 230 def build(self, input_shape=None):231 if self.built:232 return233 self.built = True234 if getattr(self, "multi_head_attention", None) is not None:235 with tf.name_scope(self.multi_head_attention.name):236 self.multi_head_attention.build(None)237 if getattr(self, "ffn", None) is not None:238 with tf.name_scope(self.ffn.name):239 self.ffn.build(None)240 if getattr(self, "layernorm1", None) is not None:241 with tf.name_scope(self.layernorm1.name):242 self.layernorm1.build([None, None, self.d_model_size])243 if getattr(self, "layernorm2", None) is not None:244 with tf.name_scope(self.layernorm2.name):245 self.layernorm2.build([None, None, self.d_model_size])246 247 248@keras_serializable249class TFCTRLMainLayer(keras.layers.Layer):250 config_class = CTRLConfig251 252 def __init__(self, config, **kwargs):253 super().__init__(**kwargs)254 255 self.config = config256 self.output_hidden_states = config.output_hidden_states257 self.output_attentions = config.output_attentions258 self.use_cache = config.use_cache259 self.return_dict = config.use_return_dict260 261 self.d_model_size = config.n_embd262 self.num_layers = config.n_layer263 264 self.pos_encoding = positional_encoding(config.n_positions, self.d_model_size)265 266 self.w = keras.layers.Embedding(267 input_dim=config.vocab_size,268 output_dim=config.n_embd,269 embeddings_initializer=get_initializer(config.initializer_range),270 name="w",271 )272 273 self.dropout = keras.layers.Dropout(config.embd_pdrop)274 self.h = [275 TFEncoderLayer(276 config.n_embd,277 config.n_head,278 config.dff,279 config.resid_pdrop,280 config.layer_norm_epsilon,281 self.output_attentions,282 name=f"h_._{i}",283 )284 for i in range(config.n_layer)285 ]286 self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="layernorm")287 288 def get_input_embeddings(self):289 return self.w290 291 def set_input_embeddings(self, new_embeddings):292 self.w = new_embeddings293 294 def _prune_heads(self, heads_to_prune):295 """296 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}297 """298 raise NotImplementedError299 300 @unpack_inputs301 def call(302 self,303 input_ids: TFModelInputType | None = None,304 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,305 attention_mask: np.ndarray | tf.Tensor | None = None,306 token_type_ids: np.ndarray | tf.Tensor | None = None,307 position_ids: np.ndarray | tf.Tensor | None = None,308 head_mask: np.ndarray | tf.Tensor | None = None,309 inputs_embeds: np.ndarray | tf.Tensor | None = None,310 use_cache: bool | None = None,311 output_attentions: bool | None = None,312 output_hidden_states: bool | None = None,313 return_dict: bool | None = None,314 training: bool | None = False,315 ) -> tuple | TFBaseModelOutputWithPast:316 # If using past key value states, only the last tokens317 # should be given as an input318 if past_key_values is not None:319 if input_ids is not None:320 input_ids = input_ids[:, -1:]321 if inputs_embeds is not None:322 inputs_embeds = inputs_embeds[:, -1:]323 if token_type_ids is not None:324 token_type_ids = token_type_ids[:, -1:]325 326 if input_ids is not None and inputs_embeds is not None:327 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")328 elif input_ids is not None:329 input_shape = shape_list(input_ids)330 input_ids = tf.reshape(input_ids, [-1, input_shape[-1]])331 elif inputs_embeds is not None:332 input_shape = shape_list(inputs_embeds)[:-1]333 else:334 raise ValueError("You have to specify either input_ids or inputs_embeds")335 336 if past_key_values is None:337 past_length = 0338 past_key_values = [None] * len(self.h)339 else:340 past_length = shape_list(past_key_values[0][0])[-2]341 if position_ids is None:342 position_ids = tf.expand_dims(tf.range(past_length, input_shape[-1] + past_length, dtype=tf.int32), axis=0)343 position_ids = tf.tile(position_ids, [input_shape[0], 1])344 345 # Attention mask.346 if attention_mask is not None:347 # We create a 3D attention mask from a 2D tensor mask.348 # Sizes are [batch_size, 1, 1, to_seq_length]349 # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]350 # this attention mask is more simple than the triangular masking of causal attention351 # used in OpenAI GPT, we just need to prepare the broadcast dimension here.352 attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1] + past_length))353 354 # Since attention_mask is 1.0 for positions we want to attend and 0.0 for355 # masked positions, this operation will create a tensor which is 0.0 for356 # positions we want to attend and -10000.0 for masked positions.357 # Since we are adding it to the raw scores before the softmax, this is358 # effectively the same as removing these entirely.359 360 one_cst = tf.constant(1.0)361 ten_thousand_cst = tf.constant(-10000.0)362 attention_mask = tf.cast(attention_mask, dtype=one_cst.dtype)363 attention_mask = tf.multiply(tf.subtract(one_cst, attention_mask), ten_thousand_cst)364 365 # Prepare head mask if needed366 # 1.0 in head_mask indicate we keep the head367 # attention_probs has shape bsz x n_heads x N x N368 # head_mask has shape n_layer x batch x n_heads x N x N369 if head_mask is not None:370 raise NotImplementedError371 else:372 head_mask = [None] * self.num_layers373 374 if token_type_ids is not None:375 token_type_ids = tf.reshape(token_type_ids, [-1, shape_list(token_type_ids)[-1]])376 token_type_embeds = self.w(token_type_ids)377 token_type_embeds *= tf.math.sqrt(tf.cast(self.d_model_size, dtype=token_type_embeds.dtype))378 else:379 token_type_embeds = tf.constant(0.0)380 position_ids = tf.reshape(position_ids, [-1, shape_list(position_ids)[-1]])381 382 if inputs_embeds is None:383 check_embeddings_within_bounds(input_ids, self.w.input_dim)384 inputs_embeds = self.w(input_ids)385 seq_len = input_shape[-1]386 mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)387 388 inputs_embeds *= tf.math.sqrt(tf.cast(self.d_model_size, inputs_embeds.dtype))389 390 pos_embeds = tf.gather(self.pos_encoding, position_ids)391 pos_embeds = tf.cast(pos_embeds, dtype=token_type_embeds.dtype)392 hidden_states = inputs_embeds + pos_embeds + token_type_embeds393 394 hidden_states = self.dropout(hidden_states, training=training)395 396 output_shape = input_shape + [shape_list(hidden_states)[-1]]397 presents = () if use_cache else None398 all_hidden_states = () if output_hidden_states else None399 all_attentions = () if output_attentions else None400 for i, (h, layer_past) in enumerate(zip(self.h, past_key_values)):401 if output_hidden_states:402 all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)403 outputs = h(404 hidden_states,405 mask,406 layer_past,407 attention_mask,408 head_mask[i],409 use_cache,410 output_attentions,411 training=training,412 )413 hidden_states, present = outputs[:2]414 415 if use_cache:416 presents = presents + (present,)417 418 if output_attentions:419 all_attentions = all_attentions + (outputs[2],)420 421 hidden_states = self.layernorm(hidden_states)422 hidden_states = tf.reshape(hidden_states, output_shape)423 if output_hidden_states:424 all_hidden_states = all_hidden_states + (hidden_states,)425 426 if output_attentions:427 # let the number of heads free (-1) so we can extract attention even after head pruning428 attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]429 all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)430 431 if not return_dict:432 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)433 434 return TFBaseModelOutputWithPast(435 last_hidden_state=hidden_states,436 past_key_values=presents,437 hidden_states=all_hidden_states,438 attentions=all_attentions,439 )440 441 def build(self, input_shape=None):442 if self.built:443 return444 self.built = True445 if getattr(self, "w", None) is not None:446 with tf.name_scope(self.w.name):447 self.w.build(None)448 if getattr(self, "layernorm", None) is not None:449 with tf.name_scope(self.layernorm.name):450 self.layernorm.build([None, None, self.config.n_embd])451 if getattr(self, "h", None) is not None:452 for layer in self.h:453 with tf.name_scope(layer.name):454 layer.build(None)455 456 457class TFCTRLPreTrainedModel(TFPreTrainedModel):458 """459 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained460 models.461 """462 463 config_class = CTRLConfig464 base_model_prefix = "transformer"465 466 467CTRL_START_DOCSTRING = r"""468 469 This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the470 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads471 etc.)472 473 This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it474 as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and475 behavior.476 477 <Tip>478 479 TensorFlow models and layers in `transformers` accept two formats as input:480 481 - having all inputs as keyword arguments (like PyTorch models), or482 - having all inputs as a list, tuple or dict in the first positional argument.483 484 The reason the second format is supported is that Keras methods prefer this format when passing inputs to models485 and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just486 pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second487 format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with488 the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first489 positional argument:490 491 - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`492 - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:493 `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`494 - a dictionary with one or several input Tensors associated to the input names given in the docstring:495 `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`496 497 Note that when creating models and layers with498 [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry499 about any of this, as you can just pass inputs like you would to any other Python function!500 501 </Tip>502 503 Parameters:504 config ([`CTRLConfig`]): Model configuration class with all the parameters of the model.505 Initializing with a config file does not load the weights associated with the model, only the506 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.507"""508 509CTRL_INPUTS_DOCSTRING = r"""510 Args:511 input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, input_ids_length)`):512 `input_ids_length` = `sequence_length` if `past` is `None` else `past[0].shape[-2]` (`sequence_length` of513 input past key value states).514 515 Indices of input sequence tokens in the vocabulary.516 517 If `past` is used, only input IDs that do not have their past calculated should be passed as `input_ids`.518 519 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and520 [`PreTrainedTokenizer.encode`] for details.521 522 [What are input IDs?](../glossary#input-ids)523 past (`list[tf.Tensor]` of length `config.n_layers`):524 Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see525 `past` output below). Can be used to speed up sequential decoding. The token ids which have their past526 given to this model should not be passed as input ids as they have already been computed.527 attention_mask (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):528 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:529 530 - 1 for tokens that are **not masked**,531 - 0 for tokens that are **masked**.532 533 [What are attention masks?](../glossary#attention-mask)534 token_type_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, input_ids_length)`, *optional*):535 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,536 1]`:537 538 - 0 corresponds to a *sentence A* token,539 - 1 corresponds to a *sentence B* token.540 541 [What are token type IDs?](../glossary#token-type-ids)542 position_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, input_ids_length)`, *optional*):543 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,544 config.max_position_embeddings - 1]`.545 546 [What are position IDs?](../glossary#position-ids)547 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):548 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:549 550 - 1 indicates the head is **not masked**,551 - 0 indicates the head is **masked**.552 553 inputs_embeds (`tf.Tensor` or `Numpy array` of shape `(batch_size, input_ids_length, hidden_size)`, *optional*):554 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This555 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the556 model's internal embedding lookup matrix.557 use_cache (`bool`, *optional*):558 If set to `True`, `past` key value states are returned and can be used to speed up decoding (see `past`).559 output_attentions (`bool`, *optional*):560 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned561 tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the562 config will be used instead.563 output_hidden_states (`bool`, *optional*):564 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for565 more detail. This argument can be used only in eager mode, in graph mode the value in the config will be566 used instead.567 return_dict (`bool`, *optional*):568 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in569 eager mode, in graph mode the value will always be set to True.570 training (`bool`, *optional*, defaults to `False`):571 Whether or not to use the model in training mode (some modules like dropout modules have different572 behaviors between training and evaluation).573"""574 575 576@add_start_docstrings(577 "The bare CTRL Model transformer outputting raw hidden-states without any specific head on top.",578 CTRL_START_DOCSTRING,579)580class TFCTRLModel(TFCTRLPreTrainedModel):581 def __init__(self, config, *inputs, **kwargs):582 super().__init__(config, *inputs, **kwargs)583 self.transformer = TFCTRLMainLayer(config, name="transformer")584 585 @unpack_inputs586 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)587 @add_code_sample_docstrings(588 checkpoint=_CHECKPOINT_FOR_DOC,589 output_type=TFBaseModelOutputWithPast,590 config_class=_CONFIG_FOR_DOC,591 )592 def call(593 self,594 input_ids: TFModelInputType | None = None,595 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,596 attention_mask: np.ndarray | tf.Tensor | None = None,597 token_type_ids: np.ndarray | tf.Tensor | None = None,598 position_ids: np.ndarray | tf.Tensor | None = None,599 head_mask: np.ndarray | tf.Tensor | None = None,600 inputs_embeds: np.ndarray | tf.Tensor | None = None,601 use_cache: bool | None = None,602 output_attentions: bool | None = None,603 output_hidden_states: bool | None = None,604 return_dict: bool | None = None,605 training: bool | None = False,606 ) -> tuple | TFBaseModelOutputWithPast:607 outputs = self.transformer(608 input_ids=input_ids,609 past_key_values=past_key_values,610 attention_mask=attention_mask,611 token_type_ids=token_type_ids,612 position_ids=position_ids,613 head_mask=head_mask,614 inputs_embeds=inputs_embeds,615 use_cache=use_cache,616 output_attentions=output_attentions,617 output_hidden_states=output_hidden_states,618 return_dict=return_dict,619 training=training,620 )621 return outputs622 623 def build(self, input_shape=None):624 if self.built:625 return626 self.built = True627 if getattr(self, "transformer", None) is not None:628 with tf.name_scope(self.transformer.name):629 self.transformer.build(None)630 631 632class TFCTRLBiasLayer(keras.layers.Layer):633 """634 Bias as a layer. It is used for serialization purposes: `keras.Model.save_weights` stores on a per-layer basis,635 so all weights have to be registered in a layer.636 """637 638 def __init__(self, shape, initializer, trainable, name, **kwargs):639 super().__init__(name=name, **kwargs)640 self.shape = shape641 self.initializer = initializer642 self.trainable = trainable643 644 def build(self, input_shape):645 self.bias = self.add_weight(646 name="bias", shape=self.shape, initializer=self.initializer, trainable=self.trainable647 )648 super().build(input_shape)649 650 def call(self, x):651 return x + self.bias652 653 654@add_start_docstrings(655 """656 The CTRL Model transformer with a language modeling head on top (linear layer with weights tied to the input657 embeddings).658 """,659 CTRL_START_DOCSTRING,660)661class TFCTRLLMHeadModel(TFCTRLPreTrainedModel, TFCausalLanguageModelingLoss):662 def __init__(self, config, *inputs, **kwargs):663 super().__init__(config, *inputs, **kwargs)664 self.transformer = TFCTRLMainLayer(config, name="transformer")665 self.bias_layer = TFCTRLBiasLayer(666 name="lm_head", shape=[1, config.vocab_size], initializer="zeros", trainable=True667 )668 669 def get_output_embeddings(self):670 return self.get_input_embeddings()671 672 def set_output_embeddings(self, value):673 self.set_input_embeddings(value)674 675 def get_bias(self):676 return {"lm_head.bias": self.bias_layer.bias}677 678 def set_bias(self, value):679 # Replaces the existing layers containing bias for correct (de)serialization.680 vocab_size = value["lm_head.bias"].shape[-1]681 self.bias_layer = TFCTRLBiasLayer(682 name="final_logits_bias", shape=[1, vocab_size], initializer="zeros", trainable=True683 )684 self.bias_layer.build(None)685 self.bias_layer.bias.assign(value["lm_head.bias"])686 687 # Copied from transformers.models.gpt2.modeling_tf_gpt2.TFGPT2LMHeadModel.prepare_inputs_for_generation688 def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):689 token_type_ids = kwargs.get("token_type_ids")690 # only last token for inputs_ids if past is defined in kwargs691 if past_key_values:692 inputs = tf.expand_dims(inputs[:, -1], -1)693 if token_type_ids is not None:694 token_type_ids = tf.expand_dims(token_type_ids[:, -1], -1)695 696 position_ids = kwargs.get("position_ids")697 attention_mask = kwargs.get("attention_mask")698 699 if attention_mask is not None and position_ids is None:700 position_ids = tf.math.cumsum(attention_mask, axis=-1, exclusive=True)701 if past_key_values:702 position_ids = tf.expand_dims(position_ids[:, -1], -1)703 704 return {705 "input_ids": inputs,706 "attention_mask": attention_mask,707 "position_ids": position_ids,708 "past_key_values": past_key_values,709 "use_cache": use_cache,710 "token_type_ids": token_type_ids,711 }712 713 @unpack_inputs714 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)715 @add_code_sample_docstrings(716 checkpoint=_CHECKPOINT_FOR_DOC,717 output_type=TFCausalLMOutputWithPast,718 config_class=_CONFIG_FOR_DOC,719 )720 def call(721 self,722 input_ids: TFModelInputType | None = None,723 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,724 attention_mask: np.ndarray | tf.Tensor | None = None,725 token_type_ids: np.ndarray | tf.Tensor | None = None,726 position_ids: np.ndarray | tf.Tensor | None = None,727 head_mask: np.ndarray | tf.Tensor | None = None,728 inputs_embeds: np.ndarray | tf.Tensor | None = None,729 use_cache: bool | None = None,730 output_attentions: bool | None = None,731 output_hidden_states: bool | None = None,732 return_dict: bool | None = None,733 labels: np.ndarray | tf.Tensor | None = None,734 training: bool | None = False,735 ) -> tuple | TFCausalLMOutputWithPast:736 r"""737 labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):738 Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,739 config.vocab_size - 1]`.740 """741 transformer_outputs = self.transformer(742 input_ids=input_ids,743 past_key_values=past_key_values,744 attention_mask=attention_mask,745 token_type_ids=token_type_ids,746 position_ids=position_ids,747 head_mask=head_mask,748 inputs_embeds=inputs_embeds,749 use_cache=use_cache,750 output_attentions=output_attentions,751 output_hidden_states=output_hidden_states,752 return_dict=return_dict,753 training=training,754 )755 hidden_states = transformer_outputs[0]756 logits = tf.matmul(hidden_states, self.transformer.w.weights, transpose_b=True)757 logits = self.bias_layer(logits)758 759 loss = None760 if labels is not None:761 # shift labels to the left and cut last logit token762 shifted_logits = logits[:, :-1]763 labels = labels[:, 1:]764 loss = self.hf_compute_loss(labels, shifted_logits)765 766 if not return_dict:767 output = (logits,) + transformer_outputs[1:]768 return ((loss,) + output) if loss is not None else output769 770 return TFCausalLMOutputWithPast(771 loss=loss,772 logits=logits,773 past_key_values=transformer_outputs.past_key_values,774 hidden_states=transformer_outputs.hidden_states,775 attentions=transformer_outputs.attentions,776 )777 778 def build(self, input_shape=None):779 if self.built:780 return781 self.built = True782 if getattr(self, "transformer", None) is not None:783 with tf.name_scope(self.transformer.name):784 self.transformer.build(None)785 if getattr(self, "bias_layer", None) is not None:786 with tf.name_scope(self.bias_layer.name):787 self.bias_layer.build(None)788 789 790@add_start_docstrings(791 """792 The CTRL Model transformer with a sequence classification head on top (linear layer).793 794 [`TFCTRLForSequenceClassification`] uses the last token in order to do the classification, as other causal models795 (e.g. GPT-1, GPT-2) do.796 797 Since it does classification on the last token, it requires to know the position of the last token. If a798 `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If799 no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the800 padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in801 each row of the batch).802 """,803 CTRL_START_DOCSTRING,804)805class TFCTRLForSequenceClassification(TFCTRLPreTrainedModel, TFSequenceClassificationLoss):806 def __init__(self, config, *inputs, **kwargs):807 super().__init__(config, *inputs, **kwargs)808 self.num_labels = config.num_labels809 self.classifier = keras.layers.Dense(810 config.num_labels,811 kernel_initializer=get_initializer(config.initializer_range),812 name="classifier",813 use_bias=False,814 )815 self.transformer = TFCTRLMainLayer(config, name="transformer")816 self.config = config817 818 def get_output_embeddings(self):819 # Remove after transformers v4.32. Fix this model's `test_model_common_attributes` test too.820 logger.warning(821 "Sequence classification models do not have output embeddings. `.get_output_embeddings` will be removed "822 "in transformers v4.32."823 )824 return self.transformer.w825 826 @unpack_inputs827 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)828 @add_code_sample_docstrings(829 checkpoint=_CHECKPOINT_FOR_DOC,830 output_type=TFSequenceClassifierOutput,831 config_class=_CONFIG_FOR_DOC,832 )833 def call(834 self,835 input_ids: TFModelInputType | None = None,836 past_key_values: tuple[tuple[np.ndarray | tf.Tensor]] | None = None,837 attention_mask: np.ndarray | tf.Tensor | None = None,838 token_type_ids: np.ndarray | tf.Tensor | None = None,839 position_ids: np.ndarray | tf.Tensor | None = None,840 head_mask: np.ndarray | tf.Tensor | None = None,841 inputs_embeds: np.ndarray | tf.Tensor | None = None,842 use_cache: bool | None = None,843 output_attentions: bool | None = None,844 output_hidden_states: bool | None = None,845 return_dict: bool | None = None,846 labels: np.ndarray | tf.Tensor | None = None,847 training: bool | None = False,848 ) -> tuple | TFSequenceClassifierOutput:849 r"""850 labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):851 Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,852 config.vocab_size - 1]`.853 """854 855 transformer_outputs = self.transformer(856 input_ids=input_ids,857 past_key_values=past_key_values,858 attention_mask=attention_mask,859 token_type_ids=token_type_ids,860 position_ids=position_ids,861 head_mask=head_mask,862 inputs_embeds=inputs_embeds,863 use_cache=use_cache,864 output_attentions=output_attentions,865 output_hidden_states=output_hidden_states,866 return_dict=return_dict,867 training=training,868 )869 hidden_states = transformer_outputs[0]870 logits = self.classifier(hidden_states)871 logits_shape = shape_list(logits)872 batch_size = logits_shape[0]873 874 if self.config.pad_token_id is None:875 last_non_pad_token = tf.fill((batch_size,), value=logits_shape[1] - 1)876 else:877 if input_ids is not None:878 token_indices = tf.range(shape_list(input_ids)[-1])879 non_pad_mask = tf.cast(input_ids != self.config.pad_token_id, token_indices.dtype)880 last_non_pad_token = tf.reduce_max(token_indices * non_pad_mask, axis=-1)881 else:882 last_non_pad_token = tf.fill((batch_size,), value=logits_shape[1] - 1)883 logger.warning_once(884 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "885 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"886 )887 loss = None888 889 pooled_logits = tf.gather(logits, last_non_pad_token, batch_dims=1, axis=1)890 891 if labels is not None:892 if self.config.pad_token_id is None and logits_shape[0] != 1:893 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")894 895 loss = self.hf_compute_loss(tf.reshape(labels, [-1]), tf.reshape(pooled_logits, [-1, self.num_labels]))896 897 if not return_dict:898 output = (pooled_logits,) + transformer_outputs[1:]899 return ((loss,) + output) if loss is not None else output900 901 return TFSequenceClassifierOutput(902 loss=loss,903 logits=pooled_logits,904 hidden_states=transformer_outputs.hidden_states,905 attentions=transformer_outputs.attentions,906 )907 908 def build(self, input_shape=None):909 if self.built:910 return911 self.built = True912 if getattr(self, "classifier", None) is not None:913 with tf.name_scope(self.classifier.name):914 self.classifier.build([None, None, self.config.n_embd])915 if getattr(self, "transformer", None) is not None:916 with tf.name_scope(self.transformer.name):917 self.transformer.build(None)918 919 920__all__ = ["TFCTRLForSequenceClassification", "TFCTRLLMHeadModel", "TFCTRLModel", "TFCTRLPreTrainedModel"]921