Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2019 Facebook AI Research and the HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""PyTorch XLM-RoBERTa model."""17 18import math19from typing import Optional, Union20 21import torch22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ...activations import ACT2FN, gelu26from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache27from ...generation import GenerationMixin28from ...modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa, _prepare_4d_causal_attention_mask_for_sdpa29from ...modeling_layers import GradientCheckpointingLayer30from ...modeling_outputs import (31 BaseModelOutputWithPastAndCrossAttentions,32 BaseModelOutputWithPoolingAndCrossAttentions,33 CausalLMOutputWithCrossAttentions,34 MaskedLMOutput,35 MultipleChoiceModelOutput,36 QuestionAnsweringModelOutput,37 SequenceClassifierOutput,38 TokenClassifierOutput,39)40from ...modeling_utils import PreTrainedModel41from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer42from ...utils import auto_docstring, logging43from ...utils.deprecation import deprecate_kwarg44from .configuration_xlm_roberta import XLMRobertaConfig45 46 47logger = logging.get_logger(__name__)48 49 50# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->XLMRoberta51class XLMRobertaEmbeddings(nn.Module):52 """53 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.54 """55 56 # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__57 def __init__(self, config):58 super().__init__()59 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)60 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)61 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)62 63 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load64 # any TensorFlow checkpoint file65 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)66 self.dropout = nn.Dropout(config.hidden_dropout_prob)67 # position_ids (1, len position emb) is contiguous in memory and exported when serialized68 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")69 self.register_buffer(70 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False71 )72 self.register_buffer(73 "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False74 )75 76 # End copy77 self.padding_idx = config.pad_token_id78 self.position_embeddings = nn.Embedding(79 config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx80 )81 82 def forward(83 self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=084 ):85 if position_ids is None:86 if input_ids is not None:87 # Create the position ids from the input token ids. Any padded tokens remain padded.88 position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)89 else:90 position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)91 92 if input_ids is not None:93 input_shape = input_ids.size()94 else:95 input_shape = inputs_embeds.size()[:-1]96 97 seq_length = input_shape[1]98 99 # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs100 # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves101 # issue #5664102 if token_type_ids is None:103 if hasattr(self, "token_type_ids"):104 buffered_token_type_ids = self.token_type_ids[:, :seq_length]105 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)106 token_type_ids = buffered_token_type_ids_expanded107 else:108 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)109 110 if inputs_embeds is None:111 inputs_embeds = self.word_embeddings(input_ids)112 token_type_embeddings = self.token_type_embeddings(token_type_ids)113 114 embeddings = inputs_embeds + token_type_embeddings115 if self.position_embedding_type == "absolute":116 position_embeddings = self.position_embeddings(position_ids)117 embeddings += position_embeddings118 embeddings = self.LayerNorm(embeddings)119 embeddings = self.dropout(embeddings)120 return embeddings121 122 def create_position_ids_from_inputs_embeds(self, inputs_embeds):123 """124 We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.125 126 Args:127 inputs_embeds: torch.Tensor128 129 Returns: torch.Tensor130 """131 input_shape = inputs_embeds.size()[:-1]132 sequence_length = input_shape[1]133 134 position_ids = torch.arange(135 self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device136 )137 return position_ids.unsqueeze(0).expand(input_shape)138 139 140# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfAttention with Roberta->XLMRoberta141class XLMRobertaSelfAttention(nn.Module):142 def __init__(self, config, position_embedding_type=None, layer_idx=None):143 super().__init__()144 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):145 raise ValueError(146 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "147 f"heads ({config.num_attention_heads})"148 )149 150 self.num_attention_heads = config.num_attention_heads151 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)152 self.all_head_size = self.num_attention_heads * self.attention_head_size153 154 self.query = nn.Linear(config.hidden_size, self.all_head_size)155 self.key = nn.Linear(config.hidden_size, self.all_head_size)156 self.value = nn.Linear(config.hidden_size, self.all_head_size)157 158 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)159 self.position_embedding_type = position_embedding_type or getattr(160 config, "position_embedding_type", "absolute"161 )162 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":163 self.max_position_embeddings = config.max_position_embeddings164 self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)165 166 self.is_decoder = config.is_decoder167 self.layer_idx = layer_idx168 169 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")170 def forward(171 self,172 hidden_states: torch.Tensor,173 attention_mask: Optional[torch.FloatTensor] = None,174 head_mask: Optional[torch.FloatTensor] = None,175 encoder_hidden_states: Optional[torch.FloatTensor] = None,176 past_key_values: Optional[Cache] = None,177 output_attentions: Optional[bool] = False,178 cache_position: Optional[torch.Tensor] = None,179 ) -> tuple[torch.Tensor]:180 batch_size, seq_length, _ = hidden_states.shape181 query_layer = self.query(hidden_states)182 query_layer = query_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(183 1, 2184 )185 186 is_updated = False187 is_cross_attention = encoder_hidden_states is not None188 if past_key_values is not None:189 if isinstance(past_key_values, EncoderDecoderCache):190 is_updated = past_key_values.is_updated.get(self.layer_idx)191 if is_cross_attention:192 # after the first generated id, we can subsequently re-use all key/value_layer from cache193 curr_past_key_value = past_key_values.cross_attention_cache194 else:195 curr_past_key_value = past_key_values.self_attention_cache196 else:197 curr_past_key_value = past_key_values198 199 current_states = encoder_hidden_states if is_cross_attention else hidden_states200 if is_cross_attention and past_key_values is not None and is_updated:201 # reuse k,v, cross_attentions202 key_layer = curr_past_key_value.layers[self.layer_idx].keys203 value_layer = curr_past_key_value.layers[self.layer_idx].values204 else:205 key_layer = self.key(current_states)206 key_layer = key_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(207 1, 2208 )209 value_layer = self.value(current_states)210 value_layer = value_layer.view(211 batch_size, -1, self.num_attention_heads, self.attention_head_size212 ).transpose(1, 2)213 214 if past_key_values is not None:215 # save all key/value_layer to cache to be re-used for fast auto-regressive generation216 cache_position = cache_position if not is_cross_attention else None217 key_layer, value_layer = curr_past_key_value.update(218 key_layer, value_layer, self.layer_idx, {"cache_position": cache_position}219 )220 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls221 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):222 past_key_values.is_updated[self.layer_idx] = True223 224 # Take the dot product between "query" and "key" to get the raw attention scores.225 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))226 227 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":228 query_length, key_length = query_layer.shape[2], key_layer.shape[2]229 if past_key_values is not None:230 position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(231 -1, 1232 )233 else:234 position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)235 position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)236 distance = position_ids_l - position_ids_r237 238 positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)239 positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility240 241 if self.position_embedding_type == "relative_key":242 relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)243 attention_scores = attention_scores + relative_position_scores244 elif self.position_embedding_type == "relative_key_query":245 relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)246 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)247 attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key248 249 attention_scores = attention_scores / math.sqrt(self.attention_head_size)250 if attention_mask is not None:251 # Apply the attention mask is (precomputed for all layers in XLMRobertaModel forward() function)252 attention_scores = attention_scores + attention_mask253 254 # Normalize the attention scores to probabilities.255 attention_probs = nn.functional.softmax(attention_scores, dim=-1)256 257 # This is actually dropping out entire tokens to attend to, which might258 # seem a bit unusual, but is taken from the original Transformer paper.259 attention_probs = self.dropout(attention_probs)260 261 # Mask heads if we want to262 if head_mask is not None:263 attention_probs = attention_probs * head_mask264 265 context_layer = torch.matmul(attention_probs, value_layer)266 267 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()268 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)269 context_layer = context_layer.view(new_context_layer_shape)270 271 return context_layer, attention_probs272 273 274# Copied from transformers.models.roberta.modeling_roberta.RobertaSdpaSelfAttention with Roberta->XLMRoberta275class XLMRobertaSdpaSelfAttention(XLMRobertaSelfAttention):276 def __init__(self, config, position_embedding_type=None, layer_idx=None):277 super().__init__(config, position_embedding_type=position_embedding_type, layer_idx=layer_idx)278 self.dropout_prob = config.attention_probs_dropout_prob279 280 # Adapted from XLMRobertaSelfAttention281 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")282 def forward(283 self,284 hidden_states: torch.Tensor,285 attention_mask: Optional[torch.Tensor] = None,286 head_mask: Optional[torch.FloatTensor] = None,287 encoder_hidden_states: Optional[torch.FloatTensor] = None,288 past_key_values: Optional[Cache] = None,289 output_attentions: Optional[bool] = False,290 cache_position: Optional[torch.Tensor] = None,291 ) -> tuple[torch.Tensor]:292 if self.position_embedding_type != "absolute" or output_attentions or head_mask is not None:293 # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once implemented.294 logger.warning_once(295 "XLMRobertaSdpaSelfAttention is used but `torch.nn.functional.scaled_dot_product_attention` does not support "296 "non-absolute `position_embedding_type` or `output_attentions=True` or `head_mask`. Falling back to "297 "the manual attention implementation, but specifying the manual implementation will be required from "298 "Transformers version v5.0.0 onwards. This warning can be removed using the argument "299 '`attn_implementation="eager"` when loading the model.'300 )301 return super().forward(302 hidden_states,303 attention_mask,304 head_mask,305 encoder_hidden_states,306 past_key_values,307 output_attentions,308 cache_position,309 )310 311 bsz, tgt_len, _ = hidden_states.size()312 313 query_layer = (314 self.query(hidden_states).view(bsz, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2)315 )316 317 is_updated = False318 is_cross_attention = encoder_hidden_states is not None319 current_states = encoder_hidden_states if is_cross_attention else hidden_states320 if past_key_values is not None:321 if isinstance(past_key_values, EncoderDecoderCache):322 is_updated = past_key_values.is_updated.get(self.layer_idx)323 if is_cross_attention:324 # after the first generated id, we can subsequently re-use all key/value_states from cache325 curr_past_key_value = past_key_values.cross_attention_cache326 else:327 curr_past_key_value = past_key_values.self_attention_cache328 else:329 curr_past_key_value = past_key_values330 331 current_states = encoder_hidden_states if is_cross_attention else hidden_states332 if is_cross_attention and past_key_values is not None and is_updated:333 # reuse k,v, cross_attentions334 key_layer = curr_past_key_value.layers[self.layer_idx].keys335 value_layer = curr_past_key_value.layers[self.layer_idx].values336 else:337 key_layer = (338 self.key(current_states)339 .view(bsz, -1, self.num_attention_heads, self.attention_head_size)340 .transpose(1, 2)341 )342 value_layer = (343 self.value(current_states)344 .view(bsz, -1, self.num_attention_heads, self.attention_head_size)345 .transpose(1, 2)346 )347 348 if past_key_values is not None:349 # save all key/value_layer to cache to be re-used for fast auto-regressive generation350 cache_position = cache_position if not is_cross_attention else None351 key_layer, value_layer = curr_past_key_value.update(352 key_layer, value_layer, self.layer_idx, {"cache_position": cache_position}353 )354 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls355 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):356 past_key_values.is_updated[self.layer_idx] = True357 358 # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment359 # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.360 # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create361 # a causal mask in case tgt_len == 1.362 is_causal = self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1363 364 attn_output = torch.nn.functional.scaled_dot_product_attention(365 query_layer,366 key_layer,367 value_layer,368 attn_mask=attention_mask,369 dropout_p=self.dropout_prob if self.training else 0.0,370 is_causal=is_causal,371 )372 373 attn_output = attn_output.transpose(1, 2)374 attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size)375 376 return attn_output, None377 378 379# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfOutput with Roberta->XLMRoberta380class XLMRobertaSelfOutput(nn.Module):381 def __init__(self, config):382 super().__init__()383 self.dense = nn.Linear(config.hidden_size, config.hidden_size)384 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)385 self.dropout = nn.Dropout(config.hidden_dropout_prob)386 387 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:388 hidden_states = self.dense(hidden_states)389 hidden_states = self.dropout(hidden_states)390 hidden_states = self.LayerNorm(hidden_states + input_tensor)391 return hidden_states392 393 394XLM_ROBERTA_SELF_ATTENTION_CLASSES = {395 "eager": XLMRobertaSelfAttention,396 "sdpa": XLMRobertaSdpaSelfAttention,397}398 399 400# Copied from transformers.models.roberta.modeling_roberta.RobertaAttention with Roberta->XLMRoberta,ROBERTA->XLM_ROBERTA401class XLMRobertaAttention(nn.Module):402 def __init__(self, config, position_embedding_type=None, layer_idx=None):403 super().__init__()404 self.self = XLM_ROBERTA_SELF_ATTENTION_CLASSES[config._attn_implementation](405 config,406 position_embedding_type=position_embedding_type,407 layer_idx=layer_idx,408 )409 self.output = XLMRobertaSelfOutput(config)410 self.pruned_heads = set()411 412 def prune_heads(self, heads):413 if len(heads) == 0:414 return415 heads, index = find_pruneable_heads_and_indices(416 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads417 )418 419 # Prune linear layers420 self.self.query = prune_linear_layer(self.self.query, index)421 self.self.key = prune_linear_layer(self.self.key, index)422 self.self.value = prune_linear_layer(self.self.value, index)423 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)424 425 # Update hyper params and store pruned heads426 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)427 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads428 self.pruned_heads = self.pruned_heads.union(heads)429 430 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")431 def forward(432 self,433 hidden_states: torch.Tensor,434 attention_mask: Optional[torch.FloatTensor] = None,435 head_mask: Optional[torch.FloatTensor] = None,436 encoder_hidden_states: Optional[torch.FloatTensor] = None,437 past_key_values: Optional[Cache] = None,438 output_attentions: Optional[bool] = False,439 cache_position: Optional[torch.Tensor] = None,440 ) -> tuple[torch.Tensor]:441 self_outputs = self.self(442 hidden_states,443 attention_mask=attention_mask,444 head_mask=head_mask,445 encoder_hidden_states=encoder_hidden_states,446 past_key_values=past_key_values,447 output_attentions=output_attentions,448 cache_position=cache_position,449 )450 attention_output = self.output(self_outputs[0], hidden_states)451 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them452 return outputs453 454 455# Copied from transformers.models.roberta.modeling_roberta.RobertaIntermediate with Roberta->XLMRoberta456class XLMRobertaIntermediate(nn.Module):457 def __init__(self, config):458 super().__init__()459 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)460 if isinstance(config.hidden_act, str):461 self.intermediate_act_fn = ACT2FN[config.hidden_act]462 else:463 self.intermediate_act_fn = config.hidden_act464 465 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:466 hidden_states = self.dense(hidden_states)467 hidden_states = self.intermediate_act_fn(hidden_states)468 return hidden_states469 470 471# Copied from transformers.models.roberta.modeling_roberta.RobertaOutput with Roberta->XLMRoberta472class XLMRobertaOutput(nn.Module):473 def __init__(self, config):474 super().__init__()475 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)476 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)477 self.dropout = nn.Dropout(config.hidden_dropout_prob)478 479 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:480 hidden_states = self.dense(hidden_states)481 hidden_states = self.dropout(hidden_states)482 hidden_states = self.LayerNorm(hidden_states + input_tensor)483 return hidden_states484 485 486# Copied from transformers.models.roberta.modeling_roberta.RobertaLayer with Roberta->XLMRoberta487class XLMRobertaLayer(GradientCheckpointingLayer):488 def __init__(self, config, layer_idx=None):489 super().__init__()490 self.chunk_size_feed_forward = config.chunk_size_feed_forward491 self.seq_len_dim = 1492 self.attention = XLMRobertaAttention(config, layer_idx=layer_idx)493 self.is_decoder = config.is_decoder494 self.add_cross_attention = config.add_cross_attention495 if self.add_cross_attention:496 if not self.is_decoder:497 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")498 self.crossattention = XLMRobertaAttention(config, position_embedding_type="absolute", layer_idx=layer_idx)499 self.intermediate = XLMRobertaIntermediate(config)500 self.output = XLMRobertaOutput(config)501 502 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")503 def forward(504 self,505 hidden_states: torch.Tensor,506 attention_mask: Optional[torch.FloatTensor] = None,507 head_mask: Optional[torch.FloatTensor] = None,508 encoder_hidden_states: Optional[torch.FloatTensor] = None,509 encoder_attention_mask: Optional[torch.FloatTensor] = None,510 past_key_values: Optional[Cache] = None,511 output_attentions: Optional[bool] = False,512 cache_position: Optional[torch.Tensor] = None,513 ) -> tuple[torch.Tensor]:514 self_attention_outputs = self.attention(515 hidden_states,516 attention_mask=attention_mask,517 head_mask=head_mask,518 output_attentions=output_attentions,519 past_key_values=past_key_values,520 cache_position=cache_position,521 )522 attention_output = self_attention_outputs[0]523 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights524 525 if self.is_decoder and encoder_hidden_states is not None:526 if not hasattr(self, "crossattention"):527 raise ValueError(528 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"529 " by setting `config.add_cross_attention=True`"530 )531 532 cross_attention_outputs = self.crossattention(533 attention_output,534 attention_mask=encoder_attention_mask,535 head_mask=head_mask,536 encoder_hidden_states=encoder_hidden_states,537 past_key_values=past_key_values,538 output_attentions=output_attentions,539 cache_position=cache_position,540 )541 attention_output = cross_attention_outputs[0]542 outputs = outputs + cross_attention_outputs[1:] # add cross attentions if we output attention weights543 544 layer_output = apply_chunking_to_forward(545 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output546 )547 outputs = (layer_output,) + outputs548 549 return outputs550 551 def feed_forward_chunk(self, attention_output):552 intermediate_output = self.intermediate(attention_output)553 layer_output = self.output(intermediate_output, attention_output)554 return layer_output555 556 557# Copied from transformers.models.roberta.modeling_roberta.RobertaEncoder with Roberta->XLMRoberta558class XLMRobertaEncoder(nn.Module):559 def __init__(self, config, layer_idx=None):560 super().__init__()561 self.config = config562 self.layer = nn.ModuleList([XLMRobertaLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])563 self.gradient_checkpointing = False564 565 def forward(566 self,567 hidden_states: torch.Tensor,568 attention_mask: Optional[torch.FloatTensor] = None,569 head_mask: Optional[torch.FloatTensor] = None,570 encoder_hidden_states: Optional[torch.FloatTensor] = None,571 encoder_attention_mask: Optional[torch.FloatTensor] = None,572 past_key_values: Optional[Cache] = None,573 use_cache: Optional[bool] = None,574 output_attentions: Optional[bool] = False,575 output_hidden_states: Optional[bool] = False,576 return_dict: Optional[bool] = True,577 cache_position: Optional[torch.Tensor] = None,578 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:579 all_hidden_states = () if output_hidden_states else None580 all_self_attentions = () if output_attentions else None581 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None582 583 if self.gradient_checkpointing and self.training:584 if use_cache:585 logger.warning_once(586 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."587 )588 use_cache = False589 590 if use_cache and self.config.is_decoder and past_key_values is None:591 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))592 593 if use_cache and self.config.is_decoder and isinstance(past_key_values, tuple):594 logger.warning_once(595 "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "596 "You should pass an instance of `EncoderDecoderCache` instead, e.g. "597 "`past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`."598 )599 past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)600 601 for i, layer_module in enumerate(self.layer):602 if output_hidden_states:603 all_hidden_states = all_hidden_states + (hidden_states,)604 605 layer_head_mask = head_mask[i] if head_mask is not None else None606 607 layer_outputs = layer_module(608 hidden_states,609 attention_mask,610 layer_head_mask,611 encoder_hidden_states, # as a positional argument for gradient checkpointing612 encoder_attention_mask=encoder_attention_mask,613 past_key_values=past_key_values,614 output_attentions=output_attentions,615 cache_position=cache_position,616 )617 618 hidden_states = layer_outputs[0]619 if output_attentions:620 all_self_attentions = all_self_attentions + (layer_outputs[1],)621 if self.config.add_cross_attention:622 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)623 624 if output_hidden_states:625 all_hidden_states = all_hidden_states + (hidden_states,)626 627 if not return_dict:628 return tuple(629 v630 for v in [631 hidden_states,632 past_key_values,633 all_hidden_states,634 all_self_attentions,635 all_cross_attentions,636 ]637 if v is not None638 )639 return BaseModelOutputWithPastAndCrossAttentions(640 last_hidden_state=hidden_states,641 past_key_values=past_key_values,642 hidden_states=all_hidden_states,643 attentions=all_self_attentions,644 cross_attentions=all_cross_attentions,645 )646 647 648# Copied from transformers.models.roberta.modeling_roberta.RobertaPooler with Roberta->XLMRoberta649class XLMRobertaPooler(nn.Module):650 def __init__(self, config):651 super().__init__()652 self.dense = nn.Linear(config.hidden_size, config.hidden_size)653 self.activation = nn.Tanh()654 655 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:656 # We "pool" the model by simply taking the hidden state corresponding657 # to the first token.658 first_token_tensor = hidden_states[:, 0]659 pooled_output = self.dense(first_token_tensor)660 pooled_output = self.activation(pooled_output)661 return pooled_output662 663 664@auto_docstring665# Copied from transformers.models.roberta.modeling_roberta.RobertaPreTrainedModel with Roberta->XLMRoberta666class XLMRobertaPreTrainedModel(PreTrainedModel):667 config: XLMRobertaConfig668 base_model_prefix = "roberta"669 supports_gradient_checkpointing = True670 _no_split_modules = ["XLMRobertaEmbeddings", "XLMRobertaSelfAttention", "XLMRobertaSdpaSelfAttention"]671 _supports_sdpa = True672 673 # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights with BertLMPredictionHead->XLMRobertaLMHead674 def _init_weights(self, module):675 """Initialize the weights"""676 if isinstance(module, nn.Linear):677 # Slightly different from the TF version which uses truncated_normal for initialization678 # cf https://github.com/pytorch/pytorch/pull/5617679 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)680 if module.bias is not None:681 module.bias.data.zero_()682 elif isinstance(module, nn.Embedding):683 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)684 if module.padding_idx is not None:685 module.weight.data[module.padding_idx].zero_()686 elif isinstance(module, nn.LayerNorm):687 module.bias.data.zero_()688 module.weight.data.fill_(1.0)689 elif isinstance(module, XLMRobertaLMHead):690 module.bias.data.zero_()691 692 693@auto_docstring694# Copied from transformers.models.roberta.modeling_roberta.RobertaModel with Roberta->XLMRoberta, ROBERTA->XLM_ROBERTA695class XLMRobertaModel(XLMRobertaPreTrainedModel):696 _no_split_modules = ["XLMRobertaEmbeddings", "XLMRobertaLayer"]697 698 def __init__(self, config, add_pooling_layer=True):699 r"""700 add_pooling_layer (bool, *optional*, defaults to `True`):701 Whether to add a pooling layer702 """703 super().__init__(config)704 self.config = config705 706 self.embeddings = XLMRobertaEmbeddings(config)707 self.encoder = XLMRobertaEncoder(config)708 709 self.pooler = XLMRobertaPooler(config) if add_pooling_layer else None710 711 self.attn_implementation = config._attn_implementation712 self.position_embedding_type = config.position_embedding_type713 714 # Initialize weights and apply final processing715 self.post_init()716 717 def get_input_embeddings(self):718 return self.embeddings.word_embeddings719 720 def set_input_embeddings(self, value):721 self.embeddings.word_embeddings = value722 723 def _prune_heads(self, heads_to_prune):724 """725 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base726 class PreTrainedModel727 """728 for layer, heads in heads_to_prune.items():729 self.encoder.layer[layer].attention.prune_heads(heads)730 731 @auto_docstring732 def forward(733 self,734 input_ids: Optional[torch.Tensor] = None,735 attention_mask: Optional[torch.Tensor] = None,736 token_type_ids: Optional[torch.Tensor] = None,737 position_ids: Optional[torch.Tensor] = None,738 head_mask: Optional[torch.Tensor] = None,739 inputs_embeds: Optional[torch.Tensor] = None,740 encoder_hidden_states: Optional[torch.Tensor] = None,741 encoder_attention_mask: Optional[torch.Tensor] = None,742 past_key_values: Optional[Cache] = None,743 use_cache: Optional[bool] = None,744 output_attentions: Optional[bool] = None,745 output_hidden_states: Optional[bool] = None,746 return_dict: Optional[bool] = None,747 cache_position: Optional[torch.Tensor] = None,748 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:749 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions750 output_hidden_states = (751 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states752 )753 return_dict = return_dict if return_dict is not None else self.config.use_return_dict754 755 if self.config.is_decoder:756 use_cache = use_cache if use_cache is not None else self.config.use_cache757 else:758 use_cache = False759 760 if input_ids is not None and inputs_embeds is not None:761 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")762 elif input_ids is not None:763 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)764 input_shape = input_ids.size()765 elif inputs_embeds is not None:766 input_shape = inputs_embeds.size()[:-1]767 else:768 raise ValueError("You have to specify either input_ids or inputs_embeds")769 770 batch_size, seq_length = input_shape771 device = input_ids.device if input_ids is not None else inputs_embeds.device772 773 past_key_values_length = 0774 if past_key_values is not None:775 past_key_values_length = (776 past_key_values[0][0].shape[-2]777 if not isinstance(past_key_values, Cache)778 else past_key_values.get_seq_length()779 )780 781 if token_type_ids is None:782 if hasattr(self.embeddings, "token_type_ids"):783 buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]784 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)785 token_type_ids = buffered_token_type_ids_expanded786 else:787 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)788 789 embedding_output = self.embeddings(790 input_ids=input_ids,791 position_ids=position_ids,792 token_type_ids=token_type_ids,793 inputs_embeds=inputs_embeds,794 past_key_values_length=past_key_values_length,795 )796 797 if attention_mask is None:798 attention_mask = torch.ones((batch_size, seq_length + past_key_values_length), device=device)799 800 use_sdpa_attention_masks = (801 self.attn_implementation == "sdpa"802 and self.position_embedding_type == "absolute"803 and head_mask is None804 and not output_attentions805 )806 807 # Expand the attention mask808 if use_sdpa_attention_masks and attention_mask.dim() == 2:809 # Expand the attention mask for SDPA.810 # [bsz, seq_len] -> [bsz, 1, seq_len, seq_len]811 if self.config.is_decoder:812 extended_attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(813 attention_mask,814 input_shape,815 embedding_output,816 past_key_values_length,817 )818 else:819 extended_attention_mask = _prepare_4d_attention_mask_for_sdpa(820 attention_mask, embedding_output.dtype, tgt_len=seq_length821 )822 else:823 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]824 # ourselves in which case we just need to make it broadcastable to all heads.825 extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)826 827 # If a 2D or 3D attention mask is provided for the cross-attention828 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]829 if self.config.is_decoder and encoder_hidden_states is not None:830 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()831 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)832 if encoder_attention_mask is None:833 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)834 835 if use_sdpa_attention_masks and encoder_attention_mask.dim() == 2:836 # Expand the attention mask for SDPA.837 # [bsz, seq_len] -> [bsz, 1, seq_len, seq_len]838 encoder_extended_attention_mask = _prepare_4d_attention_mask_for_sdpa(839 encoder_attention_mask, embedding_output.dtype, tgt_len=seq_length840 )841 else:842 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)843 else:844 encoder_extended_attention_mask = None845 846 # Prepare head mask if needed847 # 1.0 in head_mask indicate we keep the head848 # attention_probs has shape bsz x n_heads x N x N849 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]850 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]851 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)852 853 encoder_outputs = self.encoder(854 embedding_output,855 attention_mask=extended_attention_mask,856 head_mask=head_mask,857 encoder_hidden_states=encoder_hidden_states,858 encoder_attention_mask=encoder_extended_attention_mask,859 past_key_values=past_key_values,860 use_cache=use_cache,861 output_attentions=output_attentions,862 output_hidden_states=output_hidden_states,863 return_dict=return_dict,864 cache_position=cache_position,865 )866 sequence_output = encoder_outputs[0]867 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None868 869 if not return_dict:870 return (sequence_output, pooled_output) + encoder_outputs[1:]871 872 return BaseModelOutputWithPoolingAndCrossAttentions(873 last_hidden_state=sequence_output,874 pooler_output=pooled_output,875 past_key_values=encoder_outputs.past_key_values,876 hidden_states=encoder_outputs.hidden_states,877 attentions=encoder_outputs.attentions,878 cross_attentions=encoder_outputs.cross_attentions,879 )880 881 882@auto_docstring(883 custom_intro="""884 XLM-RoBERTa Model with a `language modeling` head on top for CLM fine-tuning.885 """886)887# Copied from transformers.models.roberta.modeling_roberta.RobertaForCausalLM with Roberta->XLMRoberta, ROBERTA->XLM_ROBERTA888class XLMRobertaForCausalLM(XLMRobertaPreTrainedModel, GenerationMixin):889 _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]890 891 def __init__(self, config):892 super().__init__(config)893 894 if not config.is_decoder:895 logger.warning("If you want to use `XLMRobertaLMHeadModel` as a standalone, add `is_decoder=True.`")896 897 self.roberta = XLMRobertaModel(config, add_pooling_layer=False)898 self.lm_head = XLMRobertaLMHead(config)899 900 # Initialize weights and apply final processing901 self.post_init()902 903 def get_output_embeddings(self):904 return self.lm_head.decoder905 906 def set_output_embeddings(self, new_embeddings):907 self.lm_head.decoder = new_embeddings908 909 @auto_docstring910 def forward(911 self,912 input_ids: Optional[torch.LongTensor] = None,913 attention_mask: Optional[torch.FloatTensor] = None,914 token_type_ids: Optional[torch.LongTensor] = None,915 position_ids: Optional[torch.LongTensor] = None,916 head_mask: Optional[torch.FloatTensor] = None,917 inputs_embeds: Optional[torch.FloatTensor] = None,918 encoder_hidden_states: Optional[torch.FloatTensor] = None,919 encoder_attention_mask: Optional[torch.FloatTensor] = None,920 labels: Optional[torch.LongTensor] = None,921 past_key_values: Optional[Cache] = None,922 use_cache: Optional[bool] = None,923 output_attentions: Optional[bool] = None,924 output_hidden_states: Optional[bool] = None,925 return_dict: Optional[bool] = None,926 **kwargs,927 ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:928 r"""929 token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):930 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1]`:931 932 - 0 corresponds to a *sentence A* token,933 - 1 corresponds to a *sentence B* token.934 This parameter can only be used when the model is initialized with `type_vocab_size` parameter with value935 >= 2. All the value in this tensor should be always < type_vocab_size.936 937 [What are token type IDs?](../glossary#token-type-ids)938 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):939 Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in940 `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are941 ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`942 943 Example:944 945 ```python946 >>> from transformers import AutoTokenizer, XLMRobertaForCausalLM, AutoConfig947 >>> import torch948 949 >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base")950 >>> config = AutoConfig.from_pretrained("FacebookAI/roberta-base")951 >>> config.is_decoder = True952 >>> model = XLMRobertaForCausalLM.from_pretrained("FacebookAI/roberta-base", config=config)953 954 >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")955 >>> outputs = model(**inputs)956 957 >>> prediction_logits = outputs.logits958 ```"""959 return_dict = return_dict if return_dict is not None else self.config.use_return_dict960 if labels is not None:961 use_cache = False962 963 outputs = self.roberta(964 input_ids,965 attention_mask=attention_mask,966 token_type_ids=token_type_ids,967 position_ids=position_ids,968 head_mask=head_mask,969 inputs_embeds=inputs_embeds,970 encoder_hidden_states=encoder_hidden_states,971 encoder_attention_mask=encoder_attention_mask,972 past_key_values=past_key_values,973 use_cache=use_cache,974 output_attentions=output_attentions,975 output_hidden_states=output_hidden_states,976 return_dict=return_dict,977 )978 979 sequence_output = outputs[0]980 prediction_scores = self.lm_head(sequence_output)981 982 lm_loss = None983 if labels is not None:984 # move labels to correct device to enable model parallelism985 labels = labels.to(prediction_scores.device)986 lm_loss = self.loss_function(987 prediction_scores,988 labels,989 vocab_size=self.config.vocab_size,990 **kwargs,991 )992 993 if not return_dict:994 output = (prediction_scores,) + outputs[2:]995 return ((lm_loss,) + output) if lm_loss is not None else output996 997 return CausalLMOutputWithCrossAttentions(998 loss=lm_loss,999 logits=prediction_scores,1000 past_key_values=outputs.past_key_values,1001 hidden_states=outputs.hidden_states,1002 attentions=outputs.attentions,1003 cross_attentions=outputs.cross_attentions,1004 )1005 1006 1007@auto_docstring1008# Copied from transformers.models.roberta.modeling_roberta.RobertaForMaskedLM with Roberta->XLMRoberta, ROBERTA->XLM_ROBERTA1009class XLMRobertaForMaskedLM(XLMRobertaPreTrainedModel):1010 _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]1011 1012 def __init__(self, config):1013 super().__init__(config)1014 1015 if config.is_decoder:1016 logger.warning(1017 "If you want to use `XLMRobertaForMaskedLM` make sure `config.is_decoder=False` for "1018 "bi-directional self-attention."1019 )1020 1021 self.roberta = XLMRobertaModel(config, add_pooling_layer=False)1022 self.lm_head = XLMRobertaLMHead(config)1023 1024 # Initialize weights and apply final processing1025 self.post_init()1026 1027 def get_output_embeddings(self):1028 return self.lm_head.decoder1029 1030 def set_output_embeddings(self, new_embeddings):1031 self.lm_head.decoder = new_embeddings1032 1033 @auto_docstring1034 def forward(1035 self,1036 input_ids: Optional[torch.LongTensor] = None,1037 attention_mask: Optional[torch.FloatTensor] = None,1038 token_type_ids: Optional[torch.LongTensor] = None,1039 position_ids: Optional[torch.LongTensor] = None,1040 head_mask: Optional[torch.FloatTensor] = None,1041 inputs_embeds: Optional[torch.FloatTensor] = None,1042 encoder_hidden_states: Optional[torch.FloatTensor] = None,1043 encoder_attention_mask: Optional[torch.FloatTensor] = None,1044 labels: Optional[torch.LongTensor] = None,1045 output_attentions: Optional[bool] = None,1046 output_hidden_states: Optional[bool] = None,1047 return_dict: Optional[bool] = None,1048 ) -> Union[tuple[torch.Tensor], MaskedLMOutput]:1049 r"""1050 token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1051 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1]`:1052 1053 - 0 corresponds to a *sentence A* token,1054 - 1 corresponds to a *sentence B* token.1055 This parameter can only be used when the model is initialized with `type_vocab_size` parameter with value1056 >= 2. All the value in this tensor should be always < type_vocab_size.1057 1058 [What are token type IDs?](../glossary#token-type-ids)1059 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1060 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1061 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the1062 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1063 """1064 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1065 1066 outputs = self.roberta(1067 input_ids,1068 attention_mask=attention_mask,1069 token_type_ids=token_type_ids,1070 position_ids=position_ids,1071 head_mask=head_mask,1072 inputs_embeds=inputs_embeds,1073 encoder_hidden_states=encoder_hidden_states,1074 encoder_attention_mask=encoder_attention_mask,1075 output_attentions=output_attentions,1076 output_hidden_states=output_hidden_states,1077 return_dict=return_dict,1078 )1079 sequence_output = outputs[0]1080 prediction_scores = self.lm_head(sequence_output)1081 1082 masked_lm_loss = None1083 if labels is not None:1084 # move labels to correct device to enable model parallelism1085 labels = labels.to(prediction_scores.device)1086 loss_fct = CrossEntropyLoss()1087 masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))1088 1089 if not return_dict:1090 output = (prediction_scores,) + outputs[2:]1091 return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output1092 1093 return MaskedLMOutput(1094 loss=masked_lm_loss,1095 logits=prediction_scores,1096 hidden_states=outputs.hidden_states,1097 attentions=outputs.attentions,1098 )1099 1100 1101# Copied from transformers.models.roberta.modeling_roberta.RobertaLMHead1102class XLMRobertaLMHead(nn.Module):1103 """Roberta Head for masked language modeling."""1104 1105 def __init__(self, config):1106 super().__init__()1107 self.dense = nn.Linear(config.hidden_size, config.hidden_size)1108 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1109 1110 self.decoder = nn.Linear(config.hidden_size, config.vocab_size)1111 self.bias = nn.Parameter(torch.zeros(config.vocab_size))1112 self.decoder.bias = self.bias1113 1114 def forward(self, features, **kwargs):1115 x = self.dense(features)1116 x = gelu(x)1117 x = self.layer_norm(x)1118 1119 # project back to size of vocabulary with bias1120 x = self.decoder(x)1121 1122 return x1123 1124 def _tie_weights(self):1125 # To tie those two weights if they get disconnected (on TPU or when the bias is resized)1126 # For accelerate compatibility and to not break backward compatibility1127 if self.decoder.bias.device.type == "meta":1128 self.decoder.bias = self.bias1129 else:1130 self.bias = self.decoder.bias1131 1132 1133@auto_docstring(1134 custom_intro="""1135 XLM-RoBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the1136 pooled output) e.g. for GLUE tasks.1137 """1138)1139# Copied from transformers.models.roberta.modeling_roberta.RobertaForSequenceClassification with Roberta->XLMRoberta, ROBERTA->XLM_ROBERTA1140class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):1141 def __init__(self, config):1142 super().__init__(config)1143 self.num_labels = config.num_labels1144 self.config = config1145 1146 self.roberta = XLMRobertaModel(config, add_pooling_layer=False)1147 self.classifier = XLMRobertaClassificationHead(config)1148 1149 # Initialize weights and apply final processing1150 self.post_init()1151 1152 @auto_docstring1153 def forward(1154 self,1155 input_ids: Optional[torch.LongTensor] = None,1156 attention_mask: Optional[torch.FloatTensor] = None,1157 token_type_ids: Optional[torch.LongTensor] = None,1158 position_ids: Optional[torch.LongTensor] = None,1159 head_mask: Optional[torch.FloatTensor] = None,1160 inputs_embeds: Optional[torch.FloatTensor] = None,1161 labels: Optional[torch.LongTensor] = None,1162 output_attentions: Optional[bool] = None,1163 output_hidden_states: Optional[bool] = None,1164 return_dict: Optional[bool] = None,1165 ) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:1166 r"""1167 token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1168 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1]`:1169 1170 - 0 corresponds to a *sentence A* token,1171 - 1 corresponds to a *sentence B* token.1172 This parameter can only be used when the model is initialized with `type_vocab_size` parameter with value1173 >= 2. All the value in this tensor should be always < type_vocab_size.1174 1175 [What are token type IDs?](../glossary#token-type-ids)1176 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1177 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1178 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1179 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1180 """1181 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1182 1183 outputs = self.roberta(1184 input_ids,1185 attention_mask=attention_mask,1186 token_type_ids=token_type_ids,1187 position_ids=position_ids,1188 head_mask=head_mask,1189 inputs_embeds=inputs_embeds,1190 output_attentions=output_attentions,1191 output_hidden_states=output_hidden_states,1192 return_dict=return_dict,1193 )1194 sequence_output = outputs[0]1195 logits = self.classifier(sequence_output)1196 1197 loss = None1198 if labels is not None:1199 # move labels to correct device to enable model parallelism1200 labels = labels.to(logits.device)