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"""PyTorch RoFormer model."""16 17import math18import os19from typing import Callable, Optional, Union20 21import numpy as np22import torch23from torch import nn24from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss25 26from ...activations import ACT2FN, get_activation27from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache28from ...generation import GenerationMixin29from ...modeling_layers import GradientCheckpointingLayer30from ...modeling_outputs import (31 BaseModelOutputWithPastAndCrossAttentions,32 CausalLMOutputWithCrossAttentions,33 MaskedLMOutput,34 MultipleChoiceModelOutput,35 QuestionAnsweringModelOutput,36 SequenceClassifierOutput,37 TokenClassifierOutput,38)39from ...modeling_utils import PreTrainedModel40from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer41from ...utils import auto_docstring, logging42from ...utils.deprecation import deprecate_kwarg43from .configuration_roformer import RoFormerConfig44 45 46logger = logging.get_logger(__name__)47 48 49# Copied from transformers.models.marian.modeling_marian.MarianSinusoidalPositionalEmbedding with Marian->RoFormer50class RoFormerSinusoidalPositionalEmbedding(nn.Embedding):51 """This module produces sinusoidal positional embeddings of any length."""52 53 def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None) -> None:54 super().__init__(num_positions, embedding_dim)55 56 def _init_weight(self):57 """58 Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in59 the 2nd half of the vector. [dim // 2:]60 """61 n_pos, dim = self.weight.shape62 position_enc = np.array(63 [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]64 )65 out = torch.empty(n_pos, dim, dtype=self.weight.dtype, requires_grad=False)66 sentinel = dim // 2 if dim % 2 == 0 else (dim // 2) + 167 out[:, 0:sentinel] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))68 out[:, sentinel:] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))69 self.weight = nn.Parameter(out, requires_grad=False)70 71 @torch.no_grad()72 def forward(73 self, input_ids_shape: torch.Size, past_key_values_length: int = 0, position_ids: Optional[torch.Tensor] = None74 ) -> torch.Tensor:75 """`input_ids_shape` is expected to be [bsz x seqlen]."""76 if position_ids is None:77 bsz, seq_len = input_ids_shape[:2]78 position_ids = torch.arange(79 past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device80 )81 return super().forward(position_ids)82 83 84def load_tf_weights_in_roformer(model, config, tf_checkpoint_path):85 """Load tf checkpoints in a pytorch model."""86 try:87 import re88 89 import numpy as np90 import tensorflow as tf91 except ImportError:92 logger.error(93 "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "94 "https://www.tensorflow.org/install/ for installation instructions."95 )96 raise97 tf_path = os.path.abspath(tf_checkpoint_path)98 logger.info(f"Converting TensorFlow checkpoint from {tf_path}")99 # Load weights from TF model100 init_vars = tf.train.list_variables(tf_path)101 names = []102 arrays = []103 for name, shape in init_vars:104 logger.info(f"Loading TF weight {name} with shape {shape}")105 array = tf.train.load_variable(tf_path, name)106 names.append(name.replace("bert", "roformer"))107 arrays.append(array)108 109 for name, array in zip(names, arrays):110 name = name.split("/")111 # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v112 # which are not required for using pretrained model113 if any(114 n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]115 for n in name116 ):117 logger.info(f"Skipping {'/'.join(name)}")118 continue119 pointer = model120 for m_name in name:121 if re.fullmatch(r"[A-Za-z]+_\d+", m_name):122 scope_names = re.split(r"_(\d+)", m_name)123 else:124 scope_names = [m_name]125 if scope_names[0] == "kernel" or scope_names[0] == "gamma":126 pointer = getattr(pointer, "weight")127 elif scope_names[0] == "output_bias" or scope_names[0] == "beta":128 pointer = getattr(pointer, "bias")129 elif scope_names[0] == "output_weights":130 pointer = getattr(pointer, "weight")131 elif scope_names[0] == "squad":132 pointer = getattr(pointer, "classifier")133 else:134 try:135 pointer = getattr(pointer, scope_names[0])136 except AttributeError:137 logger.info(f"Skipping {'/'.join(name)}")138 continue139 if len(scope_names) >= 2:140 num = int(scope_names[1])141 pointer = pointer[num]142 if m_name[-11:] == "_embeddings":143 pointer = getattr(pointer, "weight")144 elif m_name == "kernel":145 array = np.transpose(array)146 try:147 if not pointer.shape == array.shape:148 raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")149 except AssertionError as e:150 e.args += (pointer.shape, array.shape)151 raise152 logger.info(f"Initialize PyTorch weight {name}")153 pointer.data = torch.from_numpy(array)154 return model155 156 157class RoFormerEmbeddings(nn.Module):158 """Construct the embeddings from word and token_type embeddings."""159 160 def __init__(self, config):161 super().__init__()162 self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)163 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)164 165 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load166 # any TensorFlow checkpoint file167 self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)168 self.dropout = nn.Dropout(config.hidden_dropout_prob)169 170 def forward(self, input_ids=None, token_type_ids=None, inputs_embeds=None):171 if input_ids is not None:172 input_shape = input_ids.size()173 else:174 input_shape = inputs_embeds.size()[:-1]175 176 if inputs_embeds is None:177 inputs_embeds = self.word_embeddings(input_ids)178 179 if token_type_ids is None:180 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=inputs_embeds.device)181 182 token_type_embeddings = self.token_type_embeddings(token_type_ids)183 184 embeddings = inputs_embeds + token_type_embeddings185 186 embeddings = self.LayerNorm(embeddings)187 embeddings = self.dropout(embeddings)188 return embeddings189 190 191class RoFormerSelfAttention(nn.Module):192 def __init__(self, config, layer_idx=None):193 super().__init__()194 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):195 raise ValueError(196 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "197 f"heads ({config.num_attention_heads})"198 )199 200 self.num_attention_heads = config.num_attention_heads201 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)202 self.all_head_size = self.num_attention_heads * self.attention_head_size203 204 self.query = nn.Linear(config.hidden_size, self.all_head_size)205 self.key = nn.Linear(config.hidden_size, self.all_head_size)206 self.value = nn.Linear(config.hidden_size, self.all_head_size)207 208 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)209 210 self.is_decoder = config.is_decoder211 self.rotary_value = config.rotary_value212 self.layer_idx = layer_idx213 214 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")215 def forward(216 self,217 hidden_states,218 attention_mask=None,219 sinusoidal_pos=None,220 head_mask=None,221 encoder_hidden_states=None,222 past_key_values=None,223 output_attentions=False,224 cache_position=None,225 ):226 batch_size, seq_length, _ = hidden_states.shape227 query_layer = (228 self.query(hidden_states)229 .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)230 .transpose(1, 2)231 )232 # If this is instantiated as a cross-attention module, the keys233 # and values come from an encoder; the attention mask needs to be234 # such that the encoder's padding tokens are not attended to.235 is_cross_attention = encoder_hidden_states is not None236 237 is_updated = False238 if past_key_values is not None:239 if isinstance(past_key_values, EncoderDecoderCache):240 is_updated = past_key_values.is_updated.get(self.layer_idx)241 if is_cross_attention:242 # after the first generated id, we can subsequently re-use all key/value_layer from cache243 curr_past_key_value = past_key_values.cross_attention_cache244 else:245 curr_past_key_value = past_key_values.self_attention_cache246 else:247 curr_past_key_value = past_key_values248 249 current_states = encoder_hidden_states if is_cross_attention else hidden_states250 if is_cross_attention and past_key_values is not None and is_updated:251 # reuse k,v, cross_attentions252 key_layer = curr_past_key_value.layers[self.layer_idx].keys253 value_layer = curr_past_key_value.layers[self.layer_idx].values254 else:255 key_layer = (256 self.key(current_states)257 .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)258 .transpose(1, 2)259 )260 value_layer = (261 self.value(current_states)262 .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)263 .transpose(1, 2)264 )265 266 # Apply RoPE if self attention267 if not is_cross_attention and sinusoidal_pos is not None:268 if self.rotary_value:269 query_layer, key_layer, value_layer = self.apply_rotary_position_embeddings(270 sinusoidal_pos, query_layer, key_layer, value_layer271 )272 else:273 query_layer, key_layer = self.apply_rotary_position_embeddings(274 sinusoidal_pos, query_layer, key_layer275 )276 277 if past_key_values is not None:278 # save all key/value_layer to cache to be re-used for fast auto-regressive generation279 cache_position = cache_position if not is_cross_attention else None280 key_layer, value_layer = curr_past_key_value.update(281 key_layer, value_layer, self.layer_idx, {"cache_position": cache_position}282 )283 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls284 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):285 past_key_values.is_updated[self.layer_idx] = True286 287 # Take the dot product between "query" and "key" to get the raw attention scores.288 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))289 290 attention_scores = attention_scores / math.sqrt(self.attention_head_size)291 if attention_mask is not None:292 # Apply the attention mask is (precomputed for all layers in RoFormerModel forward() function)293 attention_scores = attention_scores + attention_mask294 295 # Normalize the attention scores to probabilities.296 attention_probs = nn.functional.softmax(attention_scores, dim=-1)297 298 # This is actually dropping out entire tokens to attend to, which might299 # seem a bit unusual, but is taken from the original Transformer paper.300 attention_probs = self.dropout(attention_probs)301 302 # Mask heads if we want to303 if head_mask is not None:304 attention_probs = attention_probs * head_mask305 306 context_layer = torch.matmul(attention_probs, value_layer)307 308 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()309 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)310 context_layer = context_layer.view(*new_context_layer_shape)311 312 return context_layer, attention_probs313 314 @staticmethod315 def apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer, value_layer=None):316 # https://kexue.fm/archives/8265317 # sin [batch_size, num_heads, sequence_length, embed_size_per_head//2]318 # cos [batch_size, num_heads, sequence_length, embed_size_per_head//2]319 sin, cos = sinusoidal_pos.chunk(2, dim=-1)320 # sin [θ0,θ1,θ2......θd/2-1] -> sin_pos [θ0,θ0,θ1,θ1,θ2,θ2......θd/2-1,θd/2-1]321 sin_pos = torch.stack([sin, sin], dim=-1).reshape_as(sinusoidal_pos)322 # cos [θ0,θ1,θ2......θd/2-1] -> cos_pos [θ0,θ0,θ1,θ1,θ2,θ2......θd/2-1,θd/2-1]323 cos_pos = torch.stack([cos, cos], dim=-1).reshape_as(sinusoidal_pos)324 # rotate_half_query_layer [-q1,q0,-q3,q2......,-qd-1,qd-2]325 rotate_half_query_layer = torch.stack([-query_layer[..., 1::2], query_layer[..., ::2]], dim=-1).reshape_as(326 query_layer327 )328 query_layer = query_layer * cos_pos + rotate_half_query_layer * sin_pos329 # rotate_half_key_layer [-k1,k0,-k3,k2......,-kd-1,kd-2]330 rotate_half_key_layer = torch.stack([-key_layer[..., 1::2], key_layer[..., ::2]], dim=-1).reshape_as(key_layer)331 key_layer = key_layer * cos_pos + rotate_half_key_layer * sin_pos332 if value_layer is not None:333 # rotate_half_value_layer [-v1,v0,-v3,v2......,-vd-1,vd-2]334 rotate_half_value_layer = torch.stack([-value_layer[..., 1::2], value_layer[..., ::2]], dim=-1).reshape_as(335 value_layer336 )337 value_layer = value_layer * cos_pos + rotate_half_value_layer * sin_pos338 return query_layer, key_layer, value_layer339 return query_layer, key_layer340 341 342# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->RoFormer343class RoFormerSelfOutput(nn.Module):344 def __init__(self, config):345 super().__init__()346 self.dense = nn.Linear(config.hidden_size, config.hidden_size)347 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)348 self.dropout = nn.Dropout(config.hidden_dropout_prob)349 350 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:351 hidden_states = self.dense(hidden_states)352 hidden_states = self.dropout(hidden_states)353 hidden_states = self.LayerNorm(hidden_states + input_tensor)354 return hidden_states355 356 357class RoFormerAttention(nn.Module):358 def __init__(self, config, layer_idx=None):359 super().__init__()360 self.self = RoFormerSelfAttention(config, layer_idx=layer_idx)361 self.output = RoFormerSelfOutput(config)362 self.pruned_heads = set()363 364 # Copied from transformers.models.bert.modeling_bert.BertAttention.prune_heads365 def prune_heads(self, heads):366 if len(heads) == 0:367 return368 heads, index = find_pruneable_heads_and_indices(369 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads370 )371 372 # Prune linear layers373 self.self.query = prune_linear_layer(self.self.query, index)374 self.self.key = prune_linear_layer(self.self.key, index)375 self.self.value = prune_linear_layer(self.self.value, index)376 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)377 378 # Update hyper params and store pruned heads379 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)380 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads381 self.pruned_heads = self.pruned_heads.union(heads)382 383 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")384 def forward(385 self,386 hidden_states,387 attention_mask=None,388 sinusoidal_pos=None,389 head_mask=None,390 encoder_hidden_states=None,391 past_key_values=None,392 output_attentions=False,393 cache_position=None,394 ):395 self_outputs = self.self(396 hidden_states,397 attention_mask=attention_mask,398 sinusoidal_pos=sinusoidal_pos,399 head_mask=head_mask,400 encoder_hidden_states=encoder_hidden_states,401 past_key_values=past_key_values,402 output_attentions=output_attentions,403 cache_position=cache_position,404 )405 attention_output = self.output(self_outputs[0], hidden_states)406 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them407 return outputs408 409 410# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->RoFormer411class RoFormerIntermediate(nn.Module):412 def __init__(self, config):413 super().__init__()414 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)415 if isinstance(config.hidden_act, str):416 self.intermediate_act_fn = ACT2FN[config.hidden_act]417 else:418 self.intermediate_act_fn = config.hidden_act419 420 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:421 hidden_states = self.dense(hidden_states)422 hidden_states = self.intermediate_act_fn(hidden_states)423 return hidden_states424 425 426# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->RoFormer427class RoFormerOutput(nn.Module):428 def __init__(self, config):429 super().__init__()430 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)431 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)432 self.dropout = nn.Dropout(config.hidden_dropout_prob)433 434 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:435 hidden_states = self.dense(hidden_states)436 hidden_states = self.dropout(hidden_states)437 hidden_states = self.LayerNorm(hidden_states + input_tensor)438 return hidden_states439 440 441class RoFormerLayer(GradientCheckpointingLayer):442 def __init__(self, config, layer_idx=None):443 super().__init__()444 self.chunk_size_feed_forward = config.chunk_size_feed_forward445 self.seq_len_dim = 1446 self.attention = RoFormerAttention(config, layer_idx)447 self.is_decoder = config.is_decoder448 self.add_cross_attention = config.add_cross_attention449 if self.add_cross_attention:450 if not self.is_decoder:451 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")452 self.crossattention = RoFormerAttention(config, layer_idx)453 self.intermediate = RoFormerIntermediate(config)454 self.output = RoFormerOutput(config)455 456 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")457 def forward(458 self,459 hidden_states,460 attention_mask=None,461 sinusoidal_pos=None,462 head_mask=None,463 encoder_hidden_states=None,464 encoder_attention_mask=None,465 past_key_values=None,466 output_attentions=False,467 cache_position=None,468 ):469 self_attention_outputs = self.attention(470 hidden_states,471 attention_mask=attention_mask,472 sinusoidal_pos=sinusoidal_pos,473 head_mask=head_mask,474 output_attentions=output_attentions,475 past_key_values=past_key_values,476 cache_position=cache_position,477 )478 attention_output = self_attention_outputs[0]479 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights480 481 if self.is_decoder and encoder_hidden_states is not None:482 if not hasattr(self, "crossattention"):483 raise ValueError(484 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention "485 "layers by setting `config.add_cross_attention=True`"486 )487 488 cross_attention_outputs = self.crossattention(489 attention_output,490 attention_mask=encoder_attention_mask,491 sinusoidal_pos=sinusoidal_pos,492 head_mask=head_mask,493 encoder_hidden_states=encoder_hidden_states,494 past_key_values=past_key_values,495 output_attentions=output_attentions,496 cache_position=cache_position,497 )498 attention_output = cross_attention_outputs[0]499 outputs = outputs + cross_attention_outputs[1:] # add cross attentions if we output attention weights500 501 layer_output = apply_chunking_to_forward(502 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output503 )504 return (layer_output,) + outputs505 506 def feed_forward_chunk(self, attention_output):507 intermediate_output = self.intermediate(attention_output)508 layer_output = self.output(intermediate_output, attention_output)509 return layer_output510 511 512class RoFormerEncoder(nn.Module):513 def __init__(self, config):514 super().__init__()515 self.config = config516 self.embed_positions = RoFormerSinusoidalPositionalEmbedding(517 config.max_position_embeddings, config.hidden_size // config.num_attention_heads518 )519 self.layer = nn.ModuleList([RoFormerLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])520 self.gradient_checkpointing = False521 522 def forward(523 self,524 hidden_states,525 attention_mask=None,526 head_mask=None,527 encoder_hidden_states=None,528 encoder_attention_mask=None,529 past_key_values=None,530 use_cache=None,531 output_attentions=False,532 output_hidden_states=False,533 return_dict=True,534 cache_position=None,535 ):536 if self.gradient_checkpointing and self.training:537 if use_cache:538 logger.warning_once(539 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."540 )541 use_cache = False542 543 if use_cache and past_key_values is None:544 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))545 if use_cache and isinstance(past_key_values, tuple):546 logger.warning_once(547 "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "548 "You should pass an instance of `EncoderDecoderCache` instead, e.g. "549 "`past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`."550 )551 past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)552 553 all_hidden_states = () if output_hidden_states else None554 all_self_attentions = () if output_attentions else None555 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None556 557 past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0558 559 # [sequence_length, embed_size_per_head] -> [batch_size, num_heads, sequence_length, embed_size_per_head]560 sinusoidal_pos = self.embed_positions(hidden_states.shape[:-1], past_key_values_length)[None, None, :, :]561 562 for i, layer_module in enumerate(self.layer):563 if output_hidden_states:564 all_hidden_states = all_hidden_states + (hidden_states,)565 566 layer_head_mask = head_mask[i] if head_mask is not None else None567 568 layer_outputs = layer_module(569 hidden_states,570 attention_mask,571 sinusoidal_pos,572 layer_head_mask,573 encoder_hidden_states,574 encoder_attention_mask,575 past_key_values,576 output_attentions,577 cache_position,578 )579 580 hidden_states = layer_outputs[0]581 if output_attentions:582 all_self_attentions = all_self_attentions + (layer_outputs[1],)583 if self.config.add_cross_attention:584 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)585 586 if output_hidden_states:587 all_hidden_states = all_hidden_states + (hidden_states,)588 589 if not return_dict:590 return tuple(591 v592 for v in [593 hidden_states,594 past_key_values,595 all_hidden_states,596 all_self_attentions,597 all_cross_attentions,598 ]599 if v is not None600 )601 return BaseModelOutputWithPastAndCrossAttentions(602 last_hidden_state=hidden_states,603 past_key_values=past_key_values,604 hidden_states=all_hidden_states,605 attentions=all_self_attentions,606 cross_attentions=all_cross_attentions,607 )608 609 610# Copied from transformers.models.xlm.modeling_xlm.XLMSequenceSummary with XLM->RoFormer611class RoFormerSequenceSummary(nn.Module):612 r"""613 Compute a single vector summary of a sequence hidden states.614 615 Args:616 config ([`RoFormerConfig`]):617 The config used by the model. Relevant arguments in the config class of the model are (refer to the actual618 config class of your model for the default values it uses):619 620 - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:621 622 - `"last"` -- Take the last token hidden state (like XLNet)623 - `"first"` -- Take the first token hidden state (like Bert)624 - `"mean"` -- Take the mean of all tokens hidden states625 - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)626 - `"attn"` -- Not implemented now, use multi-head attention627 628 - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.629 - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes630 (otherwise to `config.hidden_size`).631 - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,632 another string or `None` will add no activation.633 - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.634 - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.635 """636 637 def __init__(self, config: RoFormerConfig):638 super().__init__()639 640 self.summary_type = getattr(config, "summary_type", "last")641 if self.summary_type == "attn":642 # We should use a standard multi-head attention module with absolute positional embedding for that.643 # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276644 # We can probably just use the multi-head attention module of PyTorch >=1.1.0645 raise NotImplementedError646 647 self.summary = nn.Identity()648 if hasattr(config, "summary_use_proj") and config.summary_use_proj:649 if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0:650 num_classes = config.num_labels651 else:652 num_classes = config.hidden_size653 self.summary = nn.Linear(config.hidden_size, num_classes)654 655 activation_string = getattr(config, "summary_activation", None)656 self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity()657 658 self.first_dropout = nn.Identity()659 if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:660 self.first_dropout = nn.Dropout(config.summary_first_dropout)661 662 self.last_dropout = nn.Identity()663 if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:664 self.last_dropout = nn.Dropout(config.summary_last_dropout)665 666 def forward(667 self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None668 ) -> torch.FloatTensor:669 """670 Compute a single vector summary of a sequence hidden states.671 672 Args:673 hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):674 The hidden states of the last layer.675 cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):676 Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.677 678 Returns:679 `torch.FloatTensor`: The summary of the sequence hidden states.680 """681 if self.summary_type == "last":682 output = hidden_states[:, -1]683 elif self.summary_type == "first":684 output = hidden_states[:, 0]685 elif self.summary_type == "mean":686 output = hidden_states.mean(dim=1)687 elif self.summary_type == "cls_index":688 if cls_index is None:689 cls_index = torch.full_like(690 hidden_states[..., :1, :],691 hidden_states.shape[-2] - 1,692 dtype=torch.long,693 )694 else:695 cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)696 cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))697 # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states698 output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size)699 elif self.summary_type == "attn":700 raise NotImplementedError701 702 output = self.first_dropout(output)703 output = self.summary(output)704 output = self.activation(output)705 output = self.last_dropout(output)706 707 return output708 709 710class RoFormerPredictionHeadTransform(nn.Module):711 def __init__(self, config):712 super().__init__()713 self.dense = nn.Linear(config.hidden_size, config.embedding_size)714 if isinstance(config.hidden_act, str):715 self.transform_act_fn = ACT2FN[config.hidden_act]716 else:717 self.transform_act_fn = config.hidden_act718 self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)719 720 def forward(self, hidden_states):721 hidden_states = self.dense(hidden_states)722 hidden_states = self.transform_act_fn(hidden_states)723 hidden_states = self.LayerNorm(hidden_states)724 return hidden_states725 726 727class RoFormerLMPredictionHead(nn.Module):728 def __init__(self, config):729 super().__init__()730 self.transform = RoFormerPredictionHeadTransform(config)731 732 # The output weights are the same as the input embeddings, but there is733 # an output-only bias for each token.734 self.decoder = nn.Linear(config.embedding_size, config.vocab_size, bias=False)735 736 self.bias = nn.Parameter(torch.zeros(config.vocab_size))737 738 # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`739 self.decoder.bias = self.bias740 741 def _tie_weights(self) -> None:742 self.decoder.bias = self.bias743 744 def forward(self, hidden_states):745 hidden_states = self.transform(hidden_states)746 hidden_states = self.decoder(hidden_states)747 return hidden_states748 749 750# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->RoFormer751class RoFormerOnlyMLMHead(nn.Module):752 def __init__(self, config):753 super().__init__()754 self.predictions = RoFormerLMPredictionHead(config)755 756 def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:757 prediction_scores = self.predictions(sequence_output)758 return prediction_scores759 760 761@auto_docstring762class RoFormerPreTrainedModel(PreTrainedModel):763 config: RoFormerConfig764 load_tf_weights = load_tf_weights_in_roformer765 base_model_prefix = "roformer"766 supports_gradient_checkpointing = True767 768 def _init_weights(self, module):769 """Initialize the weights"""770 if isinstance(module, nn.Linear):771 # Slightly different from the TF version which uses truncated_normal for initialization772 # cf https://github.com/pytorch/pytorch/pull/5617773 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)774 if module.bias is not None:775 module.bias.data.zero_()776 elif isinstance(module, RoFormerSinusoidalPositionalEmbedding):777 module._init_weight()778 elif isinstance(module, nn.Embedding):779 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)780 if module.padding_idx is not None:781 module.weight.data[module.padding_idx].zero_()782 elif isinstance(module, nn.LayerNorm):783 module.bias.data.zero_()784 module.weight.data.fill_(1.0)785 elif isinstance(module, RoFormerLMPredictionHead):786 module.bias.data.zero_()787 788 789@auto_docstring(790 custom_intro="""791 792 The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of793 cross-attention is added between the self-attention layers, following the architecture described in [Attention is794 all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,795 Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.796 797 To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set798 to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and799 `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.800 """801)802class RoFormerModel(RoFormerPreTrainedModel):803 def __init__(self, config):804 super().__init__(config)805 self.config = config806 self.embeddings = RoFormerEmbeddings(config)807 808 if config.embedding_size != config.hidden_size:809 self.embeddings_project = nn.Linear(config.embedding_size, config.hidden_size)810 811 self.encoder = RoFormerEncoder(config)812 813 # Initialize weights and apply final processing814 self.post_init()815 816 def get_input_embeddings(self):817 return self.embeddings.word_embeddings818 819 def set_input_embeddings(self, value):820 self.embeddings.word_embeddings = value821 822 def _prune_heads(self, heads_to_prune):823 """824 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base825 class PreTrainedModel826 """827 for layer, heads in heads_to_prune.items():828 self.encoder.layer[layer].attention.prune_heads(heads)829 830 @auto_docstring831 def forward(832 self,833 input_ids: Optional[torch.LongTensor] = None,834 attention_mask: Optional[torch.FloatTensor] = None,835 token_type_ids: Optional[torch.LongTensor] = None,836 head_mask: Optional[torch.FloatTensor] = None,837 inputs_embeds: Optional[torch.FloatTensor] = None,838 encoder_hidden_states: Optional[torch.FloatTensor] = None,839 encoder_attention_mask: Optional[torch.FloatTensor] = None,840 past_key_values: Optional[Cache] = None,841 use_cache: Optional[bool] = None,842 output_attentions: Optional[bool] = None,843 output_hidden_states: Optional[bool] = None,844 return_dict: Optional[bool] = None,845 cache_position: Optional[torch.Tensor] = None,846 ) -> Union[BaseModelOutputWithPastAndCrossAttentions, tuple[torch.Tensor]]:847 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions848 output_hidden_states = (849 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states850 )851 return_dict = return_dict if return_dict is not None else self.config.use_return_dict852 853 if self.config.is_decoder:854 use_cache = use_cache if use_cache is not None else self.config.use_cache855 else:856 use_cache = False857 858 if input_ids is not None and inputs_embeds is not None:859 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")860 elif input_ids is not None:861 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)862 input_shape = input_ids.size()863 elif inputs_embeds is not None:864 input_shape = inputs_embeds.size()[:-1]865 else:866 raise ValueError("You have to specify either input_ids or inputs_embeds")867 868 batch_size, seq_length = input_shape869 device = input_ids.device if input_ids is not None else inputs_embeds.device870 871 past_key_values_length = 0872 if past_key_values is not None:873 past_key_values_length = (874 past_key_values[0][0].shape[-2]875 if not isinstance(past_key_values, Cache)876 else past_key_values.get_seq_length()877 )878 879 if attention_mask is None:880 attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)881 if token_type_ids is None:882 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)883 884 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]885 # ourselves in which case we just need to make it broadcastable to all heads.886 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)887 888 # If a 2D or 3D attention mask is provided for the cross-attention889 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]890 if self.config.is_decoder and encoder_hidden_states is not None:891 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()892 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)893 if encoder_attention_mask is None:894 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)895 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)896 else:897 encoder_extended_attention_mask = None898 899 # Prepare head mask if needed900 # 1.0 in head_mask indicate we keep the head901 # attention_probs has shape bsz x n_heads x N x N902 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]903 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]904 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)905 906 embedding_output = self.embeddings(907 input_ids=input_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds908 )909 if hasattr(self, "embeddings_project"):910 embedding_output = self.embeddings_project(embedding_output)911 912 encoder_outputs = self.encoder(913 embedding_output,914 attention_mask=extended_attention_mask,915 head_mask=head_mask,916 encoder_hidden_states=encoder_hidden_states,917 encoder_attention_mask=encoder_extended_attention_mask,918 past_key_values=past_key_values,919 use_cache=use_cache,920 output_attentions=output_attentions,921 output_hidden_states=output_hidden_states,922 return_dict=return_dict,923 cache_position=cache_position,924 )925 sequence_output = encoder_outputs[0]926 927 if not return_dict:928 return (sequence_output,) + encoder_outputs[1:]929 930 return BaseModelOutputWithPastAndCrossAttentions(931 last_hidden_state=sequence_output,932 past_key_values=encoder_outputs.past_key_values,933 hidden_states=encoder_outputs.hidden_states,934 attentions=encoder_outputs.attentions,935 cross_attentions=encoder_outputs.cross_attentions,936 )937 938 939@auto_docstring940class RoFormerForMaskedLM(RoFormerPreTrainedModel):941 _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]942 943 def __init__(self, config):944 super().__init__(config)945 946 if config.is_decoder:947 logger.warning(948 "If you want to use `RoFormerForMaskedLM` make sure `config.is_decoder=False` for "949 "bi-directional self-attention."950 )951 952 self.roformer = RoFormerModel(config)953 self.cls = RoFormerOnlyMLMHead(config)954 955 # Initialize weights and apply final processing956 self.post_init()957 958 def get_output_embeddings(self):959 return self.cls.predictions.decoder960 961 def set_output_embeddings(self, new_embeddings):962 self.cls.predictions.decoder = new_embeddings963 self.cls.predictions.bias = new_embeddings.bias964 965 @auto_docstring966 def forward(967 self,968 input_ids: Optional[torch.LongTensor] = None,969 attention_mask: Optional[torch.FloatTensor] = None,970 token_type_ids: Optional[torch.LongTensor] = None,971 head_mask: Optional[torch.FloatTensor] = None,972 inputs_embeds: Optional[torch.FloatTensor] = None,973 encoder_hidden_states: Optional[torch.FloatTensor] = None,974 encoder_attention_mask: Optional[torch.FloatTensor] = None,975 labels: Optional[torch.LongTensor] = None,976 output_attentions: Optional[bool] = None,977 output_hidden_states: Optional[bool] = None,978 return_dict: Optional[bool] = None,979 ) -> Union[MaskedLMOutput, tuple[torch.Tensor]]:980 r"""981 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):982 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,983 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the984 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.985 """986 return_dict = return_dict if return_dict is not None else self.config.use_return_dict987 988 outputs = self.roformer(989 input_ids,990 attention_mask=attention_mask,991 token_type_ids=token_type_ids,992 head_mask=head_mask,993 inputs_embeds=inputs_embeds,994 encoder_hidden_states=encoder_hidden_states,995 encoder_attention_mask=encoder_attention_mask,996 output_attentions=output_attentions,997 output_hidden_states=output_hidden_states,998 return_dict=return_dict,999 )1000 1001 sequence_output = outputs[0]1002 prediction_scores = self.cls(sequence_output)1003 1004 masked_lm_loss = None1005 if labels is not None:1006 loss_fct = CrossEntropyLoss() # -100 index = padding token1007 masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))1008 1009 if not return_dict:1010 output = (prediction_scores,) + outputs[1:]1011 return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output1012 1013 return MaskedLMOutput(1014 loss=masked_lm_loss,1015 logits=prediction_scores,1016 hidden_states=outputs.hidden_states,1017 attentions=outputs.attentions,1018 )1019 1020 def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):1021 input_shape = input_ids.shape1022 effective_batch_size = input_shape[0]1023 1024 # add a dummy token1025 assert self.config.pad_token_id is not None, "The PAD token should be defined for generation"1026 attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1)1027 dummy_token = torch.full(1028 (effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device1029 )1030 input_ids = torch.cat([input_ids, dummy_token], dim=1)1031 1032 return {"input_ids": input_ids, "attention_mask": attention_mask}1033 1034 1035@auto_docstring(1036 custom_intro="""1037 RoFormer Model with a `language modeling` head on top for CLM fine-tuning.1038 """1039)1040class RoFormerForCausalLM(RoFormerPreTrainedModel, GenerationMixin):1041 _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]1042 1043 def __init__(self, config):1044 super().__init__(config)1045 1046 if not config.is_decoder:1047 logger.warning("If you want to use `RoFormerForCausalLM` as a standalone, add `is_decoder=True.`")1048 1049 self.roformer = RoFormerModel(config)1050 self.cls = RoFormerOnlyMLMHead(config)1051 1052 # Initialize weights and apply final processing1053 self.post_init()1054 1055 def get_output_embeddings(self):1056 return self.cls.predictions.decoder1057 1058 def set_output_embeddings(self, new_embeddings):1059 self.cls.predictions.decoder = new_embeddings1060 self.cls.predictions.bias = new_embeddings.bias1061 1062 @auto_docstring1063 def forward(1064 self,1065 input_ids: Optional[torch.LongTensor] = None,1066 attention_mask: Optional[torch.FloatTensor] = None,1067 token_type_ids: Optional[torch.LongTensor] = None,1068 inputs_embeds: Optional[torch.FloatTensor] = None,1069 encoder_hidden_states: Optional[torch.FloatTensor] = None,1070 encoder_attention_mask: Optional[torch.FloatTensor] = None,1071 head_mask: Optional[torch.FloatTensor] = None,1072 cross_attn_head_mask: Optional[torch.Tensor] = None,1073 past_key_values: Optional[Cache] = None,1074 labels: Optional[torch.LongTensor] = None,1075 use_cache: Optional[bool] = None,1076 output_attentions: Optional[bool] = None,1077 output_hidden_states: Optional[bool] = None,1078 return_dict: Optional[bool] = None,1079 cache_position: Optional[torch.Tensor] = None,1080 **kwargs,1081 ) -> Union[CausalLMOutputWithCrossAttentions, tuple[torch.Tensor]]:1082 r"""1083 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1084 Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in1085 `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are1086 ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.1087 1088 Example:1089 1090 ```python1091 >>> from transformers import AutoTokenizer, RoFormerForCausalLM, RoFormerConfig1092 >>> import torch1093 1094 >>> tokenizer = AutoTokenizer.from_pretrained("junnyu/roformer_chinese_base")1095 >>> config = RoFormerConfig.from_pretrained("junnyu/roformer_chinese_base")1096 >>> config.is_decoder = True1097 >>> model = RoFormerForCausalLM.from_pretrained("junnyu/roformer_chinese_base", config=config)1098 1099 >>> inputs = tokenizer("今天天气非常好。", return_tensors="pt")1100 >>> outputs = model(**inputs)1101 1102 >>> prediction_logits = outputs.logits1103 ```"""1104 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1105 1106 outputs = self.roformer(1107 input_ids,1108 attention_mask=attention_mask,1109 token_type_ids=token_type_ids,1110 head_mask=head_mask,1111 inputs_embeds=inputs_embeds,1112 encoder_hidden_states=encoder_hidden_states,1113 encoder_attention_mask=encoder_attention_mask,1114 past_key_values=past_key_values,1115 use_cache=use_cache,1116 output_attentions=output_attentions,1117 output_hidden_states=output_hidden_states,1118 return_dict=return_dict,1119 cache_position=cache_position,1120 )1121 1122 sequence_output = outputs[0]1123 prediction_scores = self.cls(sequence_output)1124 1125 lm_loss = None1126 if labels is not None:1127 lm_loss = self.loss_function(1128 prediction_scores,1129 labels,1130 vocab_size=self.config.vocab_size,1131 **kwargs,1132 )1133 1134 if not return_dict:1135 output = (prediction_scores,) + outputs[1:]1136 return ((lm_loss,) + output) if lm_loss is not None else output1137 1138 return CausalLMOutputWithCrossAttentions(1139 loss=lm_loss,1140 logits=prediction_scores,1141 past_key_values=outputs.past_key_values,1142 hidden_states=outputs.hidden_states,1143 attentions=outputs.attentions,1144 cross_attentions=outputs.cross_attentions,1145 )1146 1147 1148class RoFormerClassificationHead(nn.Module):1149 """Head for sentence-level classification tasks."""1150 1151 def __init__(self, config):1152 super().__init__()1153 self.dense = nn.Linear(config.hidden_size, config.hidden_size)1154 self.dropout = nn.Dropout(config.hidden_dropout_prob)1155 self.out_proj = nn.Linear(config.hidden_size, config.num_labels)1156 1157 self.config = config1158 1159 def forward(self, features, **kwargs):1160 x = features[:, 0, :] # take <s> token (equiv. to [CLS])1161 x = self.dropout(x)1162 x = self.dense(x)1163 x = ACT2FN[self.config.hidden_act](x)1164 x = self.dropout(x)1165 x = self.out_proj(x)1166 return x1167 1168 1169@auto_docstring(1170 custom_intro="""1171 RoFormer Model transformer with a sequence classification/regression head on top (a linear layer on top of the1172 pooled output) e.g. for GLUE tasks.1173 """1174)1175class RoFormerForSequenceClassification(RoFormerPreTrainedModel):1176 def __init__(self, config):1177 super().__init__(config)1178 self.num_labels = config.num_labels1179 self.roformer = RoFormerModel(config)1180 self.classifier = RoFormerClassificationHead(config)1181 1182 # Initialize weights and apply final processing1183 self.post_init()1184 1185 @auto_docstring1186 def forward(1187 self,1188 input_ids: Optional[torch.LongTensor] = None,1189 attention_mask: Optional[torch.FloatTensor] = None,1190 token_type_ids: Optional[torch.LongTensor] = None,1191 head_mask: Optional[torch.FloatTensor] = None,1192 inputs_embeds: Optional[torch.FloatTensor] = None,1193 labels: Optional[torch.LongTensor] = None,1194 output_attentions: Optional[bool] = None,1195 output_hidden_states: Optional[bool] = None,1196 return_dict: Optional[bool] = None,1197 ) -> Union[SequenceClassifierOutput, tuple[torch.Tensor]]:1198 r"""1199 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1200 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,