Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch Data2VecText model."""16 17import math18from typing import Optional, Union19 20import torch21from torch import nn22from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss23 24from ...activations import ACT2FN, gelu25from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache26from ...generation import GenerationMixin27from ...modeling_layers import GradientCheckpointingLayer28from ...modeling_outputs import (29 BaseModelOutputWithPastAndCrossAttentions,30 BaseModelOutputWithPoolingAndCrossAttentions,31 CausalLMOutputWithCrossAttentions,32 MaskedLMOutput,33 MultipleChoiceModelOutput,34 QuestionAnsweringModelOutput,35 SequenceClassifierOutput,36 TokenClassifierOutput,37)38from ...modeling_utils import PreTrainedModel39from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer40from ...utils import auto_docstring, logging41from ...utils.deprecation import deprecate_kwarg42from .configuration_data2vec_text import Data2VecTextConfig43 44 45logger = logging.get_logger(__name__)46 47 48_HIDDEN_STATES_START_POSITION = 249 50 51# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->Data2VecText52class Data2VecTextForTextEmbeddings(nn.Module):53 """54 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.55 """56 57 # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__58 def __init__(self, config):59 super().__init__()60 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)61 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)62 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)63 64 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load65 # any TensorFlow checkpoint file66 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)67 self.dropout = nn.Dropout(config.hidden_dropout_prob)68 # position_ids (1, len position emb) is contiguous in memory and exported when serialized69 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")70 self.register_buffer(71 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False72 )73 self.register_buffer(74 "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False75 )76 77 # End copy78 self.padding_idx = config.pad_token_id79 self.position_embeddings = nn.Embedding(80 config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx81 )82 83 def forward(84 self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=085 ):86 if position_ids is None:87 if input_ids is not None:88 # Create the position ids from the input token ids. Any padded tokens remain padded.89 position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)90 else:91 position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)92 93 if input_ids is not None:94 input_shape = input_ids.size()95 else:96 input_shape = inputs_embeds.size()[:-1]97 98 seq_length = input_shape[1]99 100 # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs101 # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves102 # issue #5664103 if token_type_ids is None:104 if hasattr(self, "token_type_ids"):105 buffered_token_type_ids = self.token_type_ids[:, :seq_length]106 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)107 token_type_ids = buffered_token_type_ids_expanded108 else:109 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)110 111 if inputs_embeds is None:112 inputs_embeds = self.word_embeddings(input_ids)113 token_type_embeddings = self.token_type_embeddings(token_type_ids)114 115 embeddings = inputs_embeds + token_type_embeddings116 if self.position_embedding_type == "absolute":117 position_embeddings = self.position_embeddings(position_ids)118 embeddings += position_embeddings119 embeddings = self.LayerNorm(embeddings)120 embeddings = self.dropout(embeddings)121 return embeddings122 123 def create_position_ids_from_inputs_embeds(self, inputs_embeds):124 """125 We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.126 127 Args:128 inputs_embeds: torch.Tensor129 130 Returns: torch.Tensor131 """132 input_shape = inputs_embeds.size()[:-1]133 sequence_length = input_shape[1]134 135 position_ids = torch.arange(136 self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device137 )138 return position_ids.unsqueeze(0).expand(input_shape)139 140 141# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfAttention with Roberta->Data2VecText142class Data2VecTextSelfAttention(nn.Module):143 def __init__(self, config, position_embedding_type=None, layer_idx=None):144 super().__init__()145 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):146 raise ValueError(147 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "148 f"heads ({config.num_attention_heads})"149 )150 151 self.num_attention_heads = config.num_attention_heads152 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)153 self.all_head_size = self.num_attention_heads * self.attention_head_size154 155 self.query = nn.Linear(config.hidden_size, self.all_head_size)156 self.key = nn.Linear(config.hidden_size, self.all_head_size)157 self.value = nn.Linear(config.hidden_size, self.all_head_size)158 159 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)160 self.position_embedding_type = position_embedding_type or getattr(161 config, "position_embedding_type", "absolute"162 )163 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":164 self.max_position_embeddings = config.max_position_embeddings165 self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)166 167 self.is_decoder = config.is_decoder168 self.layer_idx = layer_idx169 170 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")171 def forward(172 self,173 hidden_states: torch.Tensor,174 attention_mask: Optional[torch.FloatTensor] = None,175 head_mask: Optional[torch.FloatTensor] = None,176 encoder_hidden_states: Optional[torch.FloatTensor] = None,177 past_key_values: Optional[Cache] = None,178 output_attentions: Optional[bool] = False,179 cache_position: Optional[torch.Tensor] = None,180 ) -> tuple[torch.Tensor]:181 batch_size, seq_length, _ = hidden_states.shape182 query_layer = self.query(hidden_states)183 query_layer = query_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(184 1, 2185 )186 187 is_updated = False188 is_cross_attention = encoder_hidden_states is not None189 if past_key_values is not None:190 if isinstance(past_key_values, EncoderDecoderCache):191 is_updated = past_key_values.is_updated.get(self.layer_idx)192 if is_cross_attention:193 # after the first generated id, we can subsequently re-use all key/value_layer from cache194 curr_past_key_value = past_key_values.cross_attention_cache195 else:196 curr_past_key_value = past_key_values.self_attention_cache197 else:198 curr_past_key_value = past_key_values199 200 current_states = encoder_hidden_states if is_cross_attention else hidden_states201 if is_cross_attention and past_key_values is not None and is_updated:202 # reuse k,v, cross_attentions203 key_layer = curr_past_key_value.layers[self.layer_idx].keys204 value_layer = curr_past_key_value.layers[self.layer_idx].values205 else:206 key_layer = self.key(current_states)207 key_layer = key_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(208 1, 2209 )210 value_layer = self.value(current_states)211 value_layer = value_layer.view(212 batch_size, -1, self.num_attention_heads, self.attention_head_size213 ).transpose(1, 2)214 215 if past_key_values is not None:216 # save all key/value_layer to cache to be re-used for fast auto-regressive generation217 cache_position = cache_position if not is_cross_attention else None218 key_layer, value_layer = curr_past_key_value.update(219 key_layer, value_layer, self.layer_idx, {"cache_position": cache_position}220 )221 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls222 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):223 past_key_values.is_updated[self.layer_idx] = True224 225 # Take the dot product between "query" and "key" to get the raw attention scores.226 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))227 228 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":229 query_length, key_length = query_layer.shape[2], key_layer.shape[2]230 if past_key_values is not None:231 position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(232 -1, 1233 )234 else:235 position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)236 position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)237 distance = position_ids_l - position_ids_r238 239 positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)240 positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility241 242 if self.position_embedding_type == "relative_key":243 relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)244 attention_scores = attention_scores + relative_position_scores245 elif self.position_embedding_type == "relative_key_query":246 relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)247 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)248 attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key249 250 attention_scores = attention_scores / math.sqrt(self.attention_head_size)251 if attention_mask is not None:252 # Apply the attention mask is (precomputed for all layers in Data2VecTextModel forward() function)253 attention_scores = attention_scores + attention_mask254 255 # Normalize the attention scores to probabilities.256 attention_probs = nn.functional.softmax(attention_scores, dim=-1)257 258 # This is actually dropping out entire tokens to attend to, which might259 # seem a bit unusual, but is taken from the original Transformer paper.260 attention_probs = self.dropout(attention_probs)261 262 # Mask heads if we want to263 if head_mask is not None:264 attention_probs = attention_probs * head_mask265 266 context_layer = torch.matmul(attention_probs, value_layer)267 268 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()269 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)270 context_layer = context_layer.view(new_context_layer_shape)271 272 return context_layer, attention_probs273 274 275# Copied from transformers.models.bert.modeling_bert.BertSelfOutput276class Data2VecTextSelfOutput(nn.Module):277 def __init__(self, config):278 super().__init__()279 self.dense = nn.Linear(config.hidden_size, config.hidden_size)280 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)281 self.dropout = nn.Dropout(config.hidden_dropout_prob)282 283 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:284 hidden_states = self.dense(hidden_states)285 hidden_states = self.dropout(hidden_states)286 hidden_states = self.LayerNorm(hidden_states + input_tensor)287 return hidden_states288 289 290DATA2VEC_TEXT_SELF_ATTENTION_CLASSES = {291 "eager": Data2VecTextSelfAttention,292}293 294 295# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Data2VecText,BERT->DATA2VEC_TEXT296class Data2VecTextAttention(nn.Module):297 def __init__(self, config, position_embedding_type=None, layer_idx=None):298 super().__init__()299 self.self = DATA2VEC_TEXT_SELF_ATTENTION_CLASSES[config._attn_implementation](300 config,301 position_embedding_type=position_embedding_type,302 layer_idx=layer_idx,303 )304 self.output = Data2VecTextSelfOutput(config)305 self.pruned_heads = set()306 307 def prune_heads(self, heads):308 if len(heads) == 0:309 return310 heads, index = find_pruneable_heads_and_indices(311 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads312 )313 314 # Prune linear layers315 self.self.query = prune_linear_layer(self.self.query, index)316 self.self.key = prune_linear_layer(self.self.key, index)317 self.self.value = prune_linear_layer(self.self.value, index)318 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)319 320 # Update hyper params and store pruned heads321 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)322 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads323 self.pruned_heads = self.pruned_heads.union(heads)324 325 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")326 def forward(327 self,328 hidden_states: torch.Tensor,329 attention_mask: Optional[torch.FloatTensor] = None,330 head_mask: Optional[torch.FloatTensor] = None,331 encoder_hidden_states: Optional[torch.FloatTensor] = None,332 past_key_values: Optional[Cache] = None,333 output_attentions: Optional[bool] = False,334 cache_position: Optional[torch.Tensor] = None,335 ) -> tuple[torch.Tensor]:336 self_outputs = self.self(337 hidden_states,338 attention_mask=attention_mask,339 head_mask=head_mask,340 encoder_hidden_states=encoder_hidden_states,341 past_key_values=past_key_values,342 output_attentions=output_attentions,343 cache_position=cache_position,344 )345 attention_output = self.output(self_outputs[0], hidden_states)346 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them347 return outputs348 349 350# Copied from transformers.models.bert.modeling_bert.BertIntermediate351class Data2VecTextIntermediate(nn.Module):352 def __init__(self, config):353 super().__init__()354 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)355 if isinstance(config.hidden_act, str):356 self.intermediate_act_fn = ACT2FN[config.hidden_act]357 else:358 self.intermediate_act_fn = config.hidden_act359 360 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:361 hidden_states = self.dense(hidden_states)362 hidden_states = self.intermediate_act_fn(hidden_states)363 return hidden_states364 365 366# Copied from transformers.models.bert.modeling_bert.BertOutput367class Data2VecTextOutput(nn.Module):368 def __init__(self, config):369 super().__init__()370 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)371 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)372 self.dropout = nn.Dropout(config.hidden_dropout_prob)373 374 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:375 hidden_states = self.dense(hidden_states)376 hidden_states = self.dropout(hidden_states)377 hidden_states = self.LayerNorm(hidden_states + input_tensor)378 return hidden_states379 380 381# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Data2VecText382class Data2VecTextLayer(GradientCheckpointingLayer):383 def __init__(self, config, layer_idx=None):384 super().__init__()385 self.chunk_size_feed_forward = config.chunk_size_feed_forward386 self.seq_len_dim = 1387 self.attention = Data2VecTextAttention(config, layer_idx=layer_idx)388 self.is_decoder = config.is_decoder389 self.add_cross_attention = config.add_cross_attention390 if self.add_cross_attention:391 if not self.is_decoder:392 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")393 self.crossattention = Data2VecTextAttention(394 config, position_embedding_type="absolute", layer_idx=layer_idx395 )396 self.intermediate = Data2VecTextIntermediate(config)397 self.output = Data2VecTextOutput(config)398 399 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")400 def forward(401 self,402 hidden_states: torch.Tensor,403 attention_mask: Optional[torch.FloatTensor] = None,404 head_mask: Optional[torch.FloatTensor] = None,405 encoder_hidden_states: Optional[torch.FloatTensor] = None,406 encoder_attention_mask: Optional[torch.FloatTensor] = None,407 past_key_values: Optional[Cache] = None,408 output_attentions: Optional[bool] = False,409 cache_position: Optional[torch.Tensor] = None,410 ) -> tuple[torch.Tensor]:411 self_attention_outputs = self.attention(412 hidden_states,413 attention_mask=attention_mask,414 head_mask=head_mask,415 output_attentions=output_attentions,416 past_key_values=past_key_values,417 cache_position=cache_position,418 )419 attention_output = self_attention_outputs[0]420 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights421 422 if self.is_decoder and encoder_hidden_states is not None:423 if not hasattr(self, "crossattention"):424 raise ValueError(425 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"426 " by setting `config.add_cross_attention=True`"427 )428 429 cross_attention_outputs = self.crossattention(430 attention_output,431 attention_mask=encoder_attention_mask,432 head_mask=head_mask,433 encoder_hidden_states=encoder_hidden_states,434 past_key_values=past_key_values,435 output_attentions=output_attentions,436 cache_position=cache_position,437 )438 attention_output = cross_attention_outputs[0]439 outputs = outputs + cross_attention_outputs[1:] # add cross attentions if we output attention weights440 441 layer_output = apply_chunking_to_forward(442 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output443 )444 outputs = (layer_output,) + outputs445 446 return outputs447 448 def feed_forward_chunk(self, attention_output):449 intermediate_output = self.intermediate(attention_output)450 layer_output = self.output(intermediate_output, attention_output)451 return layer_output452 453 454# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Data2VecText455class Data2VecTextEncoder(nn.Module):456 def __init__(self, config, layer_idx=None):457 super().__init__()458 self.config = config459 self.layer = nn.ModuleList([Data2VecTextLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])460 self.gradient_checkpointing = False461 462 def forward(463 self,464 hidden_states: torch.Tensor,465 attention_mask: Optional[torch.FloatTensor] = None,466 head_mask: Optional[torch.FloatTensor] = None,467 encoder_hidden_states: Optional[torch.FloatTensor] = None,468 encoder_attention_mask: Optional[torch.FloatTensor] = None,469 past_key_values: Optional[Cache] = None,470 use_cache: Optional[bool] = None,471 output_attentions: Optional[bool] = False,472 output_hidden_states: Optional[bool] = False,473 return_dict: Optional[bool] = True,474 cache_position: Optional[torch.Tensor] = None,475 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:476 all_hidden_states = () if output_hidden_states else None477 all_self_attentions = () if output_attentions else None478 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None479 480 if self.gradient_checkpointing and self.training:481 if use_cache:482 logger.warning_once(483 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."484 )485 use_cache = False486 487 if use_cache and self.config.is_decoder and past_key_values is None:488 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))489 490 if use_cache and self.config.is_decoder and isinstance(past_key_values, tuple):491 logger.warning_once(492 "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "493 "You should pass an instance of `EncoderDecoderCache` instead, e.g. "494 "`past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`."495 )496 past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)497 498 for i, layer_module in enumerate(self.layer):499 if output_hidden_states:500 all_hidden_states = all_hidden_states + (hidden_states,)501 502 layer_head_mask = head_mask[i] if head_mask is not None else None503 504 layer_outputs = layer_module(505 hidden_states,506 attention_mask,507 layer_head_mask,508 encoder_hidden_states, # as a positional argument for gradient checkpointing509 encoder_attention_mask=encoder_attention_mask,510 past_key_values=past_key_values,511 output_attentions=output_attentions,512 cache_position=cache_position,513 )514 515 hidden_states = layer_outputs[0]516 if output_attentions:517 all_self_attentions = all_self_attentions + (layer_outputs[1],)518 if self.config.add_cross_attention:519 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)520 521 if output_hidden_states:522 all_hidden_states = all_hidden_states + (hidden_states,)523 524 if not return_dict:525 return tuple(526 v527 for v in [528 hidden_states,529 past_key_values,530 all_hidden_states,531 all_self_attentions,532 all_cross_attentions,533 ]534 if v is not None535 )536 return BaseModelOutputWithPastAndCrossAttentions(537 last_hidden_state=hidden_states,538 past_key_values=past_key_values,539 hidden_states=all_hidden_states,540 attentions=all_self_attentions,541 cross_attentions=all_cross_attentions,542 )543 544 545# Copied from transformers.models.bert.modeling_bert.BertPooler546class Data2VecTextPooler(nn.Module):547 def __init__(self, config):548 super().__init__()549 self.dense = nn.Linear(config.hidden_size, config.hidden_size)550 self.activation = nn.Tanh()551 552 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:553 # We "pool" the model by simply taking the hidden state corresponding554 # to the first token.555 first_token_tensor = hidden_states[:, 0]556 pooled_output = self.dense(first_token_tensor)557 pooled_output = self.activation(pooled_output)558 return pooled_output559 560 561@auto_docstring562class Data2VecTextPreTrainedModel(PreTrainedModel):563 config: Data2VecTextConfig564 base_model_prefix = "data2vec_text"565 supports_gradient_checkpointing = True566 _no_split_modules = ["Data2VecTextForTextEmbeddings", "Data2VecTextLayer"]567 568 def _init_weights(self, module):569 """Initialize the weights"""570 if isinstance(module, nn.Linear):571 # Slightly different from the TF version which uses truncated_normal for initialization572 # cf https://github.com/pytorch/pytorch/pull/5617573 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)574 if module.bias is not None:575 module.bias.data.zero_()576 elif isinstance(module, nn.Embedding):577 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)578 if module.padding_idx is not None:579 module.weight.data[module.padding_idx].zero_()580 elif isinstance(module, nn.LayerNorm):581 if hasattr(module, "bias") and module.bias is not None:582 module.bias.data.zero_()583 if hasattr(module, "weight") and module.weight is not None:584 module.weight.data.fill_(1.0)585 586 587@auto_docstring588class Data2VecTextModel(Data2VecTextPreTrainedModel):589 """590 591 The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of592 cross-attention is added between the self-attention layers, following the architecture described in *Attention is593 all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz594 Kaiser and Illia Polosukhin.595 596 To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set597 to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and598 `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.599 600 .. _*Attention is all you need*: https://huggingface.co/papers/1706.03762601 602 """603 604 def __init__(self, config, add_pooling_layer=True):605 r"""606 add_pooling_layer (bool, *optional*, defaults to `True`):607 Whether to add a pooling layer608 """609 super().__init__(config)610 self.config = config611 612 self.embeddings = Data2VecTextForTextEmbeddings(config)613 self.encoder = Data2VecTextEncoder(config)614 615 self.pooler = Data2VecTextPooler(config) if add_pooling_layer else None616 617 # Initialize weights and apply final processing618 self.post_init()619 620 def get_input_embeddings(self):621 return self.embeddings.word_embeddings622 623 def set_input_embeddings(self, value):624 self.embeddings.word_embeddings = value625 626 def _prune_heads(self, heads_to_prune):627 """628 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base629 class PreTrainedModel630 """631 for layer, heads in heads_to_prune.items():632 self.encoder.layer[layer].attention.prune_heads(heads)633 634 @auto_docstring635 def forward(636 self,637 input_ids: Optional[torch.Tensor] = None,638 attention_mask: Optional[torch.Tensor] = None,639 token_type_ids: Optional[torch.Tensor] = None,640 position_ids: Optional[torch.Tensor] = None,641 head_mask: Optional[torch.Tensor] = None,642 inputs_embeds: Optional[torch.Tensor] = None,643 encoder_hidden_states: Optional[torch.Tensor] = None,644 encoder_attention_mask: Optional[torch.Tensor] = None,645 past_key_values: Optional[Cache] = None,646 use_cache: Optional[bool] = None,647 output_attentions: Optional[bool] = None,648 output_hidden_states: Optional[bool] = None,649 return_dict: Optional[bool] = None,650 cache_position: Optional[torch.Tensor] = None,651 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:652 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions653 output_hidden_states = (654 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states655 )656 return_dict = return_dict if return_dict is not None else self.config.use_return_dict657 658 if self.config.is_decoder:659 use_cache = use_cache if use_cache is not None else self.config.use_cache660 else:661 use_cache = False662 663 if input_ids is not None and inputs_embeds is not None:664 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")665 elif input_ids is not None:666 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)667 input_shape = input_ids.size()668 elif inputs_embeds is not None:669 input_shape = inputs_embeds.size()[:-1]670 else:671 raise ValueError("You have to specify either input_ids or inputs_embeds")672 673 batch_size, seq_length = input_shape674 device = input_ids.device if input_ids is not None else inputs_embeds.device675 676 past_key_values_length = 0677 if past_key_values is not None:678 past_key_values_length = (679 past_key_values[0][0].shape[-2]680 if not isinstance(past_key_values, Cache)681 else past_key_values.get_seq_length()682 )683 684 if attention_mask is None:685 attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)686 687 if token_type_ids is None:688 if hasattr(self.embeddings, "token_type_ids"):689 buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]690 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)691 token_type_ids = buffered_token_type_ids_expanded692 else:693 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)694 695 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]696 # ourselves in which case we just need to make it broadcastable to all heads.697 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)698 699 # If a 2D or 3D attention mask is provided for the cross-attention700 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]701 if self.config.is_decoder and encoder_hidden_states is not None:702 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()703 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)704 if encoder_attention_mask is None:705 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)706 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)707 else:708 encoder_extended_attention_mask = None709 710 # Prepare head mask if needed711 # 1.0 in head_mask indicate we keep the head712 # attention_probs has shape bsz x n_heads x N x N713 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]714 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]715 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)716 717 embedding_output = self.embeddings(718 input_ids=input_ids,719 position_ids=position_ids,720 token_type_ids=token_type_ids,721 inputs_embeds=inputs_embeds,722 past_key_values_length=past_key_values_length,723 )724 encoder_outputs = self.encoder(725 embedding_output,726 attention_mask=extended_attention_mask,727 head_mask=head_mask,728 encoder_hidden_states=encoder_hidden_states,729 encoder_attention_mask=encoder_extended_attention_mask,730 past_key_values=past_key_values,731 use_cache=use_cache,732 output_attentions=output_attentions,733 output_hidden_states=output_hidden_states,734 return_dict=return_dict,735 cache_position=cache_position,736 )737 sequence_output = encoder_outputs[0]738 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None739 740 if not return_dict:741 return (sequence_output, pooled_output) + encoder_outputs[1:]742 743 return BaseModelOutputWithPoolingAndCrossAttentions(744 last_hidden_state=sequence_output,745 pooler_output=pooled_output,746 past_key_values=encoder_outputs.past_key_values,747 hidden_states=encoder_outputs.hidden_states,748 attentions=encoder_outputs.attentions,749 cross_attentions=encoder_outputs.cross_attentions,750 )751 752 753@auto_docstring(754 custom_intro="""755 Data2VecText Model with a `language modeling` head on top for CLM fine-tuning.756 """757)758class Data2VecTextForCausalLM(Data2VecTextPreTrainedModel, GenerationMixin):759 _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]760 761 def __init__(self, config):762 super().__init__(config)763 764 if not config.is_decoder:765 logger.warning("If you want to use `Data2VecTextLMHeadModel` as a standalone, add `is_decoder=True.`")766 767 self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)768 self.lm_head = Data2VecTextLMHead(config)769 770 # Initialize weights and apply final processing771 self.post_init()772 773 def get_output_embeddings(self):774 return self.lm_head.decoder775 776 def set_output_embeddings(self, new_embeddings):777 self.lm_head.decoder = new_embeddings778 779 @auto_docstring780 def forward(781 self,782 input_ids: Optional[torch.LongTensor] = None,783 attention_mask: Optional[torch.FloatTensor] = None,784 token_type_ids: Optional[torch.LongTensor] = None,785 position_ids: Optional[torch.LongTensor] = None,786 head_mask: Optional[torch.FloatTensor] = None,787 inputs_embeds: Optional[torch.FloatTensor] = None,788 encoder_hidden_states: Optional[torch.FloatTensor] = None,789 encoder_attention_mask: Optional[torch.FloatTensor] = None,790 labels: Optional[torch.LongTensor] = None,791 past_key_values: Optional[Cache] = None,792 use_cache: Optional[bool] = None,793 output_attentions: Optional[bool] = None,794 output_hidden_states: Optional[bool] = None,795 return_dict: Optional[bool] = None,796 cache_position: Optional[torch.Tensor] = None,797 **kwargs,798 ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:799 r"""800 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):801 Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in802 `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are803 ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`804 805 Example:806 807 ```python808 >>> from transformers import AutoTokenizer, Data2VecTextForCausalLM, Data2VecTextConfig809 >>> import torch810 811 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/data2vec-text-base")812 >>> config = Data2VecTextConfig.from_pretrained("facebook/data2vec-text-base")813 >>> config.is_decoder = True814 >>> model = Data2VecTextForCausalLM.from_pretrained("facebook/data2vec-text-base", config=config)815 816 >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")817 >>> outputs = model(**inputs)818 819 >>> prediction_logits = outputs.logits820 ```"""821 return_dict = return_dict if return_dict is not None else self.config.use_return_dict822 if labels is not None:823 use_cache = False824 825 outputs = self.data2vec_text(826 input_ids,827 attention_mask=attention_mask,828 token_type_ids=token_type_ids,829 position_ids=position_ids,830 head_mask=head_mask,831 inputs_embeds=inputs_embeds,832 encoder_hidden_states=encoder_hidden_states,833 encoder_attention_mask=encoder_attention_mask,834 past_key_values=past_key_values,835 use_cache=use_cache,836 output_attentions=output_attentions,837 output_hidden_states=output_hidden_states,838 return_dict=return_dict,839 cache_position=cache_position,840 )841 842 sequence_output = outputs[0]843 prediction_scores = self.lm_head(sequence_output)844 845 lm_loss = None846 if labels is not None:847 lm_loss = self.loss_function(848 prediction_scores,849 labels,850 vocab_size=self.config.vocab_size,851 **kwargs,852 )853 854 if not return_dict:855 output = (prediction_scores,) + outputs[2:]856 return ((lm_loss,) + output) if lm_loss is not None else output857 858 return CausalLMOutputWithCrossAttentions(859 loss=lm_loss,860 logits=prediction_scores,861 past_key_values=outputs.past_key_values,862 hidden_states=outputs.hidden_states,863 attentions=outputs.attentions,864 cross_attentions=outputs.cross_attentions,865 )866 867 868@auto_docstring869class Data2VecTextForMaskedLM(Data2VecTextPreTrainedModel):870 _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]871 872 def __init__(self, config):873 super().__init__(config)874 875 if config.is_decoder:876 logger.warning(877 "If you want to use `Data2VecTextForMaskedLM` make sure `config.is_decoder=False` for "878 "bi-directional self-attention."879 )880 881 self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)882 self.lm_head = Data2VecTextLMHead(config)883 884 # Initialize weights and apply final processing885 self.post_init()886 887 def get_output_embeddings(self):888 return self.lm_head.decoder889 890 def set_output_embeddings(self, new_embeddings):891 self.lm_head.decoder = new_embeddings892 893 @auto_docstring894 def forward(895 self,896 input_ids: Optional[torch.LongTensor] = None,897 attention_mask: Optional[torch.FloatTensor] = None,898 token_type_ids: Optional[torch.LongTensor] = None,899 position_ids: Optional[torch.LongTensor] = None,900 head_mask: Optional[torch.FloatTensor] = None,901 inputs_embeds: Optional[torch.FloatTensor] = None,902 encoder_hidden_states: Optional[torch.FloatTensor] = None,903 encoder_attention_mask: Optional[torch.FloatTensor] = None,904 labels: Optional[torch.LongTensor] = None,905 output_attentions: Optional[bool] = None,906 output_hidden_states: Optional[bool] = None,907 return_dict: Optional[bool] = None,908 ) -> Union[tuple, MaskedLMOutput]:909 r"""910 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):911 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,912 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the913 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`914 """915 return_dict = return_dict if return_dict is not None else self.config.use_return_dict916 917 outputs = self.data2vec_text(918 input_ids,919 attention_mask=attention_mask,920 token_type_ids=token_type_ids,921 position_ids=position_ids,922 head_mask=head_mask,923 inputs_embeds=inputs_embeds,924 encoder_hidden_states=encoder_hidden_states,925 encoder_attention_mask=encoder_attention_mask,926 output_attentions=output_attentions,927 output_hidden_states=output_hidden_states,928 return_dict=return_dict,929 )930 sequence_output = outputs[0]931 prediction_scores = self.lm_head(sequence_output)932 933 masked_lm_loss = None934 if labels is not None:935 loss_fct = CrossEntropyLoss()936 937 labels = labels.to(prediction_scores.device)938 masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))939 940 if not return_dict:941 output = (prediction_scores,) + outputs[2:]942 return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output943 944 return MaskedLMOutput(945 loss=masked_lm_loss,946 logits=prediction_scores,947 hidden_states=outputs.hidden_states,948 attentions=outputs.attentions,949 )950 951 952# Copied from transformers.models.roberta.modeling_roberta.RobertaLMHead with Roberta->Data2VecText953class Data2VecTextLMHead(nn.Module):954 """Data2VecText Head for masked language modeling."""955 956 def __init__(self, config):957 super().__init__()958 self.dense = nn.Linear(config.hidden_size, config.hidden_size)959 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)960 961 self.decoder = nn.Linear(config.hidden_size, config.vocab_size)962 self.bias = nn.Parameter(torch.zeros(config.vocab_size))963 self.decoder.bias = self.bias964 965 def forward(self, features, **kwargs):966 x = self.dense(features)967 x = gelu(x)968 x = self.layer_norm(x)969 970 # project back to size of vocabulary with bias971 x = self.decoder(x)972 973 return x974 975 def _tie_weights(self):976 # To tie those two weights if they get disconnected (on TPU or when the bias is resized)977 # For accelerate compatibility and to not break backward compatibility978 if self.decoder.bias.device.type == "meta":979 self.decoder.bias = self.bias980 else:981 self.bias = self.decoder.bias982 983 984@auto_docstring(985 custom_intro="""986 Data2VecText Model transformer with a sequence classification/regression head on top (a linear layer on top of the987 pooled output) e.g. for GLUE tasks.988 """989)990class Data2VecTextForSequenceClassification(Data2VecTextPreTrainedModel):991 def __init__(self, config):992 super().__init__(config)993 self.num_labels = config.num_labels994 self.config = config995 996 self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)997 self.classifier = Data2VecTextClassificationHead(config)998 999 # Initialize weights and apply final processing1000 self.post_init()1001 1002 @auto_docstring1003 def forward(1004 self,1005 input_ids: Optional[torch.LongTensor] = None,1006 attention_mask: Optional[torch.FloatTensor] = None,1007 token_type_ids: Optional[torch.LongTensor] = None,1008 position_ids: Optional[torch.LongTensor] = None,1009 head_mask: Optional[torch.FloatTensor] = None,1010 inputs_embeds: Optional[torch.FloatTensor] = None,1011 labels: Optional[torch.LongTensor] = None,1012 output_attentions: Optional[bool] = None,1013 output_hidden_states: Optional[bool] = None,1014 return_dict: Optional[bool] = None,1015 ) -> Union[tuple, SequenceClassifierOutput]:1016 r"""1017 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1018 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1019 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1020 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1021 """1022 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1023 1024 outputs = self.data2vec_text(1025 input_ids,1026 attention_mask=attention_mask,1027 token_type_ids=token_type_ids,1028 position_ids=position_ids,1029 head_mask=head_mask,1030 inputs_embeds=inputs_embeds,1031 output_attentions=output_attentions,1032 output_hidden_states=output_hidden_states,1033 return_dict=return_dict,1034 )1035 sequence_output = outputs[0]1036 logits = self.classifier(sequence_output)1037 1038 loss = None1039 if labels is not None:1040 labels = labels.to(logits.device)1041 1042 if self.config.problem_type is None:1043 if self.num_labels == 1:1044 self.config.problem_type = "regression"1045 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):1046 self.config.problem_type = "single_label_classification"1047 else:1048 self.config.problem_type = "multi_label_classification"1049 1050 if self.config.problem_type == "regression":1051 loss_fct = MSELoss()1052 if self.num_labels == 1:1053 loss = loss_fct(logits.squeeze(), labels.squeeze())1054 else:1055 loss = loss_fct(logits, labels)1056 elif self.config.problem_type == "single_label_classification":1057 loss_fct = CrossEntropyLoss()1058 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))1059 elif self.config.problem_type == "multi_label_classification":1060 loss_fct = BCEWithLogitsLoss()1061 loss = loss_fct(logits, labels)1062 1063 if not return_dict:1064 output = (logits,) + outputs[2:]1065 return ((loss,) + output) if loss is not None else output1066 1067 return SequenceClassifierOutput(1068 loss=loss,1069 logits=logits,1070 hidden_states=outputs.hidden_states,1071 attentions=outputs.attentions,1072 )1073 1074 1075@auto_docstring1076class Data2VecTextForMultipleChoice(Data2VecTextPreTrainedModel):1077 def __init__(self, config):1078 super().__init__(config)1079 1080 self.data2vec_text = Data2VecTextModel(config)1081 self.dropout = nn.Dropout(config.hidden_dropout_prob)1082 self.classifier = nn.Linear(config.hidden_size, 1)1083 1084 # Initialize weights and apply final processing1085 self.post_init()1086 1087 @auto_docstring1088 def forward(1089 self,1090 input_ids: Optional[torch.LongTensor] = None,1091 token_type_ids: Optional[torch.LongTensor] = None,1092 attention_mask: Optional[torch.FloatTensor] = None,1093 labels: Optional[torch.LongTensor] = None,1094 position_ids: Optional[torch.LongTensor] = None,1095 head_mask: Optional[torch.FloatTensor] = None,1096 inputs_embeds: Optional[torch.FloatTensor] = None,1097 output_attentions: Optional[bool] = None,1098 output_hidden_states: Optional[bool] = None,1099 return_dict: Optional[bool] = None,1100 ) -> Union[tuple, MultipleChoiceModelOutput]:1101 r"""1102 input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`):1103 Indices of input sequence tokens in the vocabulary.1104 1105 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and1106 [`PreTrainedTokenizer.__call__`] for details.1107 1108 [What are input IDs?](../glossary#input-ids)1109 token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):1110 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1111 1]`:1112 1113 - 0 corresponds to a *sentence A* token,1114 - 1 corresponds to a *sentence B* token.1115 1116 [What are token type IDs?](../glossary#token-type-ids)1117 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1118 Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,1119 num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See1120 `input_ids` above)1121 position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):1122 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,1123 config.max_position_embeddings - 1]`.1124 1125 [What are position IDs?](../glossary#position-ids)1126 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, hidden_size)`, *optional*):1127 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This1128 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the1129 model's internal embedding lookup matrix.1130 """1131 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1132 num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]1133 1134 flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None1135 flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None1136 flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None1137 flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None1138 flat_inputs_embeds = (1139 inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))1140 if inputs_embeds is not None1141 else None1142 )1143 1144 outputs = self.data2vec_text(1145 flat_input_ids,1146 position_ids=flat_position_ids,1147 token_type_ids=flat_token_type_ids,1148 attention_mask=flat_attention_mask,1149 head_mask=head_mask,1150 inputs_embeds=flat_inputs_embeds,1151 output_attentions=output_attentions,1152 output_hidden_states=output_hidden_states,1153 return_dict=return_dict,1154 )1155 pooled_output = outputs[1]1156 1157 pooled_output = self.dropout(pooled_output)1158 logits = self.classifier(pooled_output)1159 reshaped_logits = logits.view(-1, num_choices)1160 1161 loss = None1162 if labels is not None:1163 loss_fct = CrossEntropyLoss()1164 1165 labels = labels.to(reshaped_logits.device)1166 loss = loss_fct(reshaped_logits, labels)1167 1168 if not return_dict:1169 output = (reshaped_logits,) + outputs[2:]1170 return ((loss,) + output) if loss is not None else output1171 1172 return MultipleChoiceModelOutput(1173 loss=loss,1174 logits=reshaped_logits,1175 hidden_states=outputs.hidden_states,1176 attentions=outputs.attentions,1177 )1178 1179 1180@auto_docstring1181class Data2VecTextForTokenClassification(Data2VecTextPreTrainedModel):1182 def __init__(self, config):1183 super().__init__(config)1184 self.num_labels = config.num_labels1185 1186 self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)1187 classifier_dropout = (1188 config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob1189 )1190 self.dropout = nn.Dropout(classifier_dropout)1191 self.classifier = nn.Linear(config.hidden_size, config.num_labels)1192 1193 # Initialize weights and apply final processing1194 self.post_init()1195 1196 @auto_docstring1197 def forward(1198 self,1199 input_ids: Optional[torch.LongTensor] = None,1200 attention_mask: Optional[torch.FloatTensor] = None,