Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright Studio Ousia and 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 LUKE model."""16 17import math18from dataclasses import dataclass19from typing import Optional, Union20 21import torch22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ...activations import ACT2FN, gelu26from ...modeling_layers import GradientCheckpointingLayer27from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling28from ...modeling_utils import PreTrainedModel29from ...pytorch_utils import apply_chunking_to_forward30from ...utils import ModelOutput, auto_docstring, logging31from .configuration_luke import LukeConfig32 33 34logger = logging.get_logger(__name__)35 36 37@dataclass38@auto_docstring(39 custom_intro="""40 Base class for outputs of the LUKE model.41 """42)43class BaseLukeModelOutputWithPooling(BaseModelOutputWithPooling):44 r"""45 pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):46 Last layer hidden-state of the first token of the sequence (classification token) further processed by a47 Linear layer and a Tanh activation function.48 entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`):49 Sequence of entity hidden-states at the output of the last layer of the model.50 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):51 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of52 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each53 layer plus the initial entity embedding outputs.54 """55 56 entity_last_hidden_state: Optional[torch.FloatTensor] = None57 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None58 59 60@dataclass61@auto_docstring(62 custom_intro="""63 Base class for model's outputs, with potential hidden states and attentions.64 """65)66class BaseLukeModelOutput(BaseModelOutput):67 r"""68 entity_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, entity_length, hidden_size)`):69 Sequence of entity hidden-states at the output of the last layer of the model.70 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):71 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of72 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each73 layer plus the initial entity embedding outputs.74 """75 76 entity_last_hidden_state: Optional[torch.FloatTensor] = None77 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None78 79 80@dataclass81@auto_docstring(82 custom_intro="""83 Base class for model's outputs, with potential hidden states and attentions.84 """85)86class LukeMaskedLMOutput(ModelOutput):87 r"""88 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):89 The sum of masked language modeling (MLM) loss and entity prediction loss.90 mlm_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):91 Masked language modeling (MLM) loss.92 mep_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):93 Masked entity prediction (MEP) loss.94 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):95 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).96 entity_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):97 Prediction scores of the entity prediction head (scores for each entity vocabulary token before SoftMax).98 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):99 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of100 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each101 layer plus the initial entity embedding outputs.102 """103 104 loss: Optional[torch.FloatTensor] = None105 mlm_loss: Optional[torch.FloatTensor] = None106 mep_loss: Optional[torch.FloatTensor] = None107 logits: Optional[torch.FloatTensor] = None108 entity_logits: Optional[torch.FloatTensor] = None109 hidden_states: Optional[tuple[torch.FloatTensor]] = None110 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None111 attentions: Optional[tuple[torch.FloatTensor, ...]] = None112 113 114@dataclass115@auto_docstring(116 custom_intro="""117 Outputs of entity classification models.118 """119)120class EntityClassificationOutput(ModelOutput):121 r"""122 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):123 Classification loss.124 logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):125 Classification scores (before SoftMax).126 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):127 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of128 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each129 layer plus the initial entity embedding outputs.130 """131 132 loss: Optional[torch.FloatTensor] = None133 logits: Optional[torch.FloatTensor] = None134 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None135 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None136 attentions: Optional[tuple[torch.FloatTensor, ...]] = None137 138 139@dataclass140@auto_docstring(141 custom_intro="""142 Outputs of entity pair classification models.143 """144)145class EntityPairClassificationOutput(ModelOutput):146 r"""147 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):148 Classification loss.149 logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):150 Classification scores (before SoftMax).151 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):152 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of153 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each154 layer plus the initial entity embedding outputs.155 """156 157 loss: Optional[torch.FloatTensor] = None158 logits: Optional[torch.FloatTensor] = None159 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None160 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None161 attentions: Optional[tuple[torch.FloatTensor, ...]] = None162 163 164@dataclass165@auto_docstring(166 custom_intro="""167 Outputs of entity span classification models.168 """169)170class EntitySpanClassificationOutput(ModelOutput):171 r"""172 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):173 Classification loss.174 logits (`torch.FloatTensor` of shape `(batch_size, entity_length, config.num_labels)`):175 Classification scores (before SoftMax).176 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):177 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of178 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each179 layer plus the initial entity embedding outputs.180 """181 182 loss: Optional[torch.FloatTensor] = None183 logits: Optional[torch.FloatTensor] = None184 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None185 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None186 attentions: Optional[tuple[torch.FloatTensor, ...]] = None187 188 189@dataclass190@auto_docstring(191 custom_intro="""192 Outputs of sentence classification models.193 """194)195class LukeSequenceClassifierOutput(ModelOutput):196 r"""197 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):198 Classification (or regression if config.num_labels==1) loss.199 logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):200 Classification (or regression if config.num_labels==1) scores (before SoftMax).201 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):202 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of203 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each204 layer plus the initial entity embedding outputs.205 """206 207 loss: Optional[torch.FloatTensor] = None208 logits: Optional[torch.FloatTensor] = None209 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None210 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None211 attentions: Optional[tuple[torch.FloatTensor, ...]] = None212 213 214@dataclass215@auto_docstring(216 custom_intro="""217 Base class for outputs of token classification models.218 """219)220class LukeTokenClassifierOutput(ModelOutput):221 r"""222 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):223 Classification loss.224 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):225 Classification scores (before SoftMax).226 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):227 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of228 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each229 layer plus the initial entity embedding outputs.230 """231 232 loss: Optional[torch.FloatTensor] = None233 logits: Optional[torch.FloatTensor] = None234 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None235 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None236 attentions: Optional[tuple[torch.FloatTensor, ...]] = None237 238 239@dataclass240@auto_docstring(241 custom_intro="""242 Outputs of question answering models.243 """244)245class LukeQuestionAnsweringModelOutput(ModelOutput):246 r"""247 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):248 Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.249 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):250 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of251 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each252 layer plus the initial entity embedding outputs.253 """254 255 loss: Optional[torch.FloatTensor] = None256 start_logits: Optional[torch.FloatTensor] = None257 end_logits: Optional[torch.FloatTensor] = None258 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None259 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None260 attentions: Optional[tuple[torch.FloatTensor, ...]] = None261 262 263@dataclass264@auto_docstring(265 custom_intro="""266 Outputs of multiple choice models.267 """268)269class LukeMultipleChoiceModelOutput(ModelOutput):270 r"""271 loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided):272 Classification loss.273 logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`):274 *num_choices* is the second dimension of the input tensors. (see *input_ids* above).275 276 Classification scores (before SoftMax).277 entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):278 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of279 shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each280 layer plus the initial entity embedding outputs.281 """282 283 loss: Optional[torch.FloatTensor] = None284 logits: Optional[torch.FloatTensor] = None285 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None286 entity_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None287 attentions: Optional[tuple[torch.FloatTensor, ...]] = None288 289 290class LukeEmbeddings(nn.Module):291 """292 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.293 """294 295 def __init__(self, config):296 super().__init__()297 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)298 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)299 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)300 301 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load302 # any TensorFlow checkpoint file303 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)304 self.dropout = nn.Dropout(config.hidden_dropout_prob)305 306 # End copy307 self.padding_idx = config.pad_token_id308 self.position_embeddings = nn.Embedding(309 config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx310 )311 312 def forward(313 self,314 input_ids=None,315 token_type_ids=None,316 position_ids=None,317 inputs_embeds=None,318 ):319 if position_ids is None:320 if input_ids is not None:321 # Create the position ids from the input token ids. Any padded tokens remain padded.322 position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx).to(input_ids.device)323 else:324 position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)325 326 if input_ids is not None:327 input_shape = input_ids.size()328 else:329 input_shape = inputs_embeds.size()[:-1]330 331 if token_type_ids is None:332 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)333 334 if inputs_embeds is None:335 inputs_embeds = self.word_embeddings(input_ids)336 337 position_embeddings = self.position_embeddings(position_ids)338 token_type_embeddings = self.token_type_embeddings(token_type_ids)339 340 embeddings = inputs_embeds + position_embeddings + token_type_embeddings341 embeddings = self.LayerNorm(embeddings)342 embeddings = self.dropout(embeddings)343 return embeddings344 345 def create_position_ids_from_inputs_embeds(self, inputs_embeds):346 """347 We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.348 349 Args:350 inputs_embeds: torch.Tensor351 352 Returns: torch.Tensor353 """354 input_shape = inputs_embeds.size()[:-1]355 sequence_length = input_shape[1]356 357 position_ids = torch.arange(358 self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device359 )360 return position_ids.unsqueeze(0).expand(input_shape)361 362 363class LukeEntityEmbeddings(nn.Module):364 def __init__(self, config: LukeConfig):365 super().__init__()366 self.config = config367 368 self.entity_embeddings = nn.Embedding(config.entity_vocab_size, config.entity_emb_size, padding_idx=0)369 if config.entity_emb_size != config.hidden_size:370 self.entity_embedding_dense = nn.Linear(config.entity_emb_size, config.hidden_size, bias=False)371 372 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)373 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)374 375 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)376 self.dropout = nn.Dropout(config.hidden_dropout_prob)377 378 def forward(379 self,380 entity_ids: torch.LongTensor,381 position_ids: torch.LongTensor,382 token_type_ids: Optional[torch.LongTensor] = None,383 ):384 if token_type_ids is None:385 token_type_ids = torch.zeros_like(entity_ids)386 387 entity_embeddings = self.entity_embeddings(entity_ids)388 if self.config.entity_emb_size != self.config.hidden_size:389 entity_embeddings = self.entity_embedding_dense(entity_embeddings)390 391 position_embeddings = self.position_embeddings(position_ids.clamp(min=0))392 position_embedding_mask = (position_ids != -1).type_as(position_embeddings).unsqueeze(-1)393 position_embeddings = position_embeddings * position_embedding_mask394 position_embeddings = torch.sum(position_embeddings, dim=-2)395 position_embeddings = position_embeddings / position_embedding_mask.sum(dim=-2).clamp(min=1e-7)396 397 token_type_embeddings = self.token_type_embeddings(token_type_ids)398 399 embeddings = entity_embeddings + position_embeddings + token_type_embeddings400 embeddings = self.LayerNorm(embeddings)401 embeddings = self.dropout(embeddings)402 403 return embeddings404 405 406class LukeSelfAttention(nn.Module):407 def __init__(self, config):408 super().__init__()409 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):410 raise ValueError(411 f"The hidden size {config.hidden_size} is not a multiple of the number of attention "412 f"heads {config.num_attention_heads}."413 )414 415 self.num_attention_heads = config.num_attention_heads416 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)417 self.all_head_size = self.num_attention_heads * self.attention_head_size418 self.use_entity_aware_attention = config.use_entity_aware_attention419 420 self.query = nn.Linear(config.hidden_size, self.all_head_size)421 self.key = nn.Linear(config.hidden_size, self.all_head_size)422 self.value = nn.Linear(config.hidden_size, self.all_head_size)423 424 if self.use_entity_aware_attention:425 self.w2e_query = nn.Linear(config.hidden_size, self.all_head_size)426 self.e2w_query = nn.Linear(config.hidden_size, self.all_head_size)427 self.e2e_query = nn.Linear(config.hidden_size, self.all_head_size)428 429 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)430 431 def transpose_for_scores(self, x):432 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)433 x = x.view(*new_x_shape)434 return x.permute(0, 2, 1, 3)435 436 def forward(437 self,438 word_hidden_states,439 entity_hidden_states,440 attention_mask=None,441 head_mask=None,442 output_attentions=False,443 ):444 word_size = word_hidden_states.size(1)445 446 if entity_hidden_states is None:447 concat_hidden_states = word_hidden_states448 else:449 concat_hidden_states = torch.cat([word_hidden_states, entity_hidden_states], dim=1)450 451 key_layer = self.transpose_for_scores(self.key(concat_hidden_states))452 value_layer = self.transpose_for_scores(self.value(concat_hidden_states))453 454 if self.use_entity_aware_attention and entity_hidden_states is not None:455 # compute query vectors using word-word (w2w), word-entity (w2e), entity-word (e2w), entity-entity (e2e)456 # query layers457 w2w_query_layer = self.transpose_for_scores(self.query(word_hidden_states))458 w2e_query_layer = self.transpose_for_scores(self.w2e_query(word_hidden_states))459 e2w_query_layer = self.transpose_for_scores(self.e2w_query(entity_hidden_states))460 e2e_query_layer = self.transpose_for_scores(self.e2e_query(entity_hidden_states))461 462 # compute w2w, w2e, e2w, and e2e key vectors used with the query vectors computed above463 w2w_key_layer = key_layer[:, :, :word_size, :]464 e2w_key_layer = key_layer[:, :, :word_size, :]465 w2e_key_layer = key_layer[:, :, word_size:, :]466 e2e_key_layer = key_layer[:, :, word_size:, :]467 468 # compute attention scores based on the dot product between the query and key vectors469 w2w_attention_scores = torch.matmul(w2w_query_layer, w2w_key_layer.transpose(-1, -2))470 w2e_attention_scores = torch.matmul(w2e_query_layer, w2e_key_layer.transpose(-1, -2))471 e2w_attention_scores = torch.matmul(e2w_query_layer, e2w_key_layer.transpose(-1, -2))472 e2e_attention_scores = torch.matmul(e2e_query_layer, e2e_key_layer.transpose(-1, -2))473 474 # combine attention scores to create the final attention score matrix475 word_attention_scores = torch.cat([w2w_attention_scores, w2e_attention_scores], dim=3)476 entity_attention_scores = torch.cat([e2w_attention_scores, e2e_attention_scores], dim=3)477 attention_scores = torch.cat([word_attention_scores, entity_attention_scores], dim=2)478 479 else:480 query_layer = self.transpose_for_scores(self.query(concat_hidden_states))481 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))482 483 attention_scores = attention_scores / math.sqrt(self.attention_head_size)484 if attention_mask is not None:485 # Apply the attention mask is (precomputed for all layers in LukeModel forward() function)486 attention_scores = attention_scores + attention_mask487 488 # Normalize the attention scores to probabilities.489 attention_probs = nn.functional.softmax(attention_scores, dim=-1)490 491 # This is actually dropping out entire tokens to attend to, which might492 # seem a bit unusual, but is taken from the original Transformer paper.493 attention_probs = self.dropout(attention_probs)494 495 # Mask heads if we want to496 if head_mask is not None:497 attention_probs = attention_probs * head_mask498 499 context_layer = torch.matmul(attention_probs, value_layer)500 501 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()502 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)503 context_layer = context_layer.view(*new_context_layer_shape)504 505 output_word_hidden_states = context_layer[:, :word_size, :]506 if entity_hidden_states is None:507 output_entity_hidden_states = None508 else:509 output_entity_hidden_states = context_layer[:, word_size:, :]510 511 if output_attentions:512 outputs = (output_word_hidden_states, output_entity_hidden_states, attention_probs)513 else:514 outputs = (output_word_hidden_states, output_entity_hidden_states)515 516 return outputs517 518 519# Copied from transformers.models.bert.modeling_bert.BertSelfOutput520class LukeSelfOutput(nn.Module):521 def __init__(self, config):522 super().__init__()523 self.dense = nn.Linear(config.hidden_size, config.hidden_size)524 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)525 self.dropout = nn.Dropout(config.hidden_dropout_prob)526 527 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:528 hidden_states = self.dense(hidden_states)529 hidden_states = self.dropout(hidden_states)530 hidden_states = self.LayerNorm(hidden_states + input_tensor)531 return hidden_states532 533 534class LukeAttention(nn.Module):535 def __init__(self, config):536 super().__init__()537 self.self = LukeSelfAttention(config)538 self.output = LukeSelfOutput(config)539 self.pruned_heads = set()540 541 def prune_heads(self, heads):542 raise NotImplementedError("LUKE does not support the pruning of attention heads")543 544 def forward(545 self,546 word_hidden_states,547 entity_hidden_states,548 attention_mask=None,549 head_mask=None,550 output_attentions=False,551 ):552 word_size = word_hidden_states.size(1)553 self_outputs = self.self(554 word_hidden_states,555 entity_hidden_states,556 attention_mask,557 head_mask,558 output_attentions,559 )560 if entity_hidden_states is None:561 concat_self_outputs = self_outputs[0]562 concat_hidden_states = word_hidden_states563 else:564 concat_self_outputs = torch.cat(self_outputs[:2], dim=1)565 concat_hidden_states = torch.cat([word_hidden_states, entity_hidden_states], dim=1)566 567 attention_output = self.output(concat_self_outputs, concat_hidden_states)568 569 word_attention_output = attention_output[:, :word_size, :]570 if entity_hidden_states is None:571 entity_attention_output = None572 else:573 entity_attention_output = attention_output[:, word_size:, :]574 575 # add attentions if we output them576 outputs = (word_attention_output, entity_attention_output) + self_outputs[2:]577 578 return outputs579 580 581# Copied from transformers.models.bert.modeling_bert.BertIntermediate582class LukeIntermediate(nn.Module):583 def __init__(self, config):584 super().__init__()585 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)586 if isinstance(config.hidden_act, str):587 self.intermediate_act_fn = ACT2FN[config.hidden_act]588 else:589 self.intermediate_act_fn = config.hidden_act590 591 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:592 hidden_states = self.dense(hidden_states)593 hidden_states = self.intermediate_act_fn(hidden_states)594 return hidden_states595 596 597# Copied from transformers.models.bert.modeling_bert.BertOutput598class LukeOutput(nn.Module):599 def __init__(self, config):600 super().__init__()601 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)602 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)603 self.dropout = nn.Dropout(config.hidden_dropout_prob)604 605 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:606 hidden_states = self.dense(hidden_states)607 hidden_states = self.dropout(hidden_states)608 hidden_states = self.LayerNorm(hidden_states + input_tensor)609 return hidden_states610 611 612class LukeLayer(GradientCheckpointingLayer):613 def __init__(self, config):614 super().__init__()615 self.chunk_size_feed_forward = config.chunk_size_feed_forward616 self.seq_len_dim = 1617 self.attention = LukeAttention(config)618 self.intermediate = LukeIntermediate(config)619 self.output = LukeOutput(config)620 621 def forward(622 self,623 word_hidden_states,624 entity_hidden_states,625 attention_mask=None,626 head_mask=None,627 output_attentions=False,628 ):629 word_size = word_hidden_states.size(1)630 631 self_attention_outputs = self.attention(632 word_hidden_states,633 entity_hidden_states,634 attention_mask,635 head_mask,636 output_attentions=output_attentions,637 )638 if entity_hidden_states is None:639 concat_attention_output = self_attention_outputs[0]640 else:641 concat_attention_output = torch.cat(self_attention_outputs[:2], dim=1)642 643 outputs = self_attention_outputs[2:] # add self attentions if we output attention weights644 645 layer_output = apply_chunking_to_forward(646 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, concat_attention_output647 )648 word_layer_output = layer_output[:, :word_size, :]649 if entity_hidden_states is None:650 entity_layer_output = None651 else:652 entity_layer_output = layer_output[:, word_size:, :]653 654 outputs = (word_layer_output, entity_layer_output) + outputs655 656 return outputs657 658 def feed_forward_chunk(self, attention_output):659 intermediate_output = self.intermediate(attention_output)660 layer_output = self.output(intermediate_output, attention_output)661 return layer_output662 663 664class LukeEncoder(nn.Module):665 def __init__(self, config):666 super().__init__()667 self.config = config668 self.layer = nn.ModuleList([LukeLayer(config) for _ in range(config.num_hidden_layers)])669 self.gradient_checkpointing = False670 671 def forward(672 self,673 word_hidden_states,674 entity_hidden_states,675 attention_mask=None,676 head_mask=None,677 output_attentions=False,678 output_hidden_states=False,679 return_dict=True,680 ):681 all_word_hidden_states = () if output_hidden_states else None682 all_entity_hidden_states = () if output_hidden_states else None683 all_self_attentions = () if output_attentions else None684 685 for i, layer_module in enumerate(self.layer):686 if output_hidden_states:687 all_word_hidden_states = all_word_hidden_states + (word_hidden_states,)688 all_entity_hidden_states = all_entity_hidden_states + (entity_hidden_states,)689 690 layer_head_mask = head_mask[i] if head_mask is not None else None691 layer_outputs = layer_module(692 word_hidden_states,693 entity_hidden_states,694 attention_mask,695 layer_head_mask,696 output_attentions,697 )698 699 word_hidden_states = layer_outputs[0]700 701 if entity_hidden_states is not None:702 entity_hidden_states = layer_outputs[1]703 704 if output_attentions:705 all_self_attentions = all_self_attentions + (layer_outputs[2],)706 707 if output_hidden_states:708 all_word_hidden_states = all_word_hidden_states + (word_hidden_states,)709 all_entity_hidden_states = all_entity_hidden_states + (entity_hidden_states,)710 711 if not return_dict:712 return tuple(713 v714 for v in [715 word_hidden_states,716 all_word_hidden_states,717 all_self_attentions,718 entity_hidden_states,719 all_entity_hidden_states,720 ]721 if v is not None722 )723 return BaseLukeModelOutput(724 last_hidden_state=word_hidden_states,725 hidden_states=all_word_hidden_states,726 attentions=all_self_attentions,727 entity_last_hidden_state=entity_hidden_states,728 entity_hidden_states=all_entity_hidden_states,729 )730 731 732# Copied from transformers.models.bert.modeling_bert.BertPooler733class LukePooler(nn.Module):734 def __init__(self, config):735 super().__init__()736 self.dense = nn.Linear(config.hidden_size, config.hidden_size)737 self.activation = nn.Tanh()738 739 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:740 # We "pool" the model by simply taking the hidden state corresponding741 # to the first token.742 first_token_tensor = hidden_states[:, 0]743 pooled_output = self.dense(first_token_tensor)744 pooled_output = self.activation(pooled_output)745 return pooled_output746 747 748class EntityPredictionHeadTransform(nn.Module):749 def __init__(self, config):750 super().__init__()751 self.dense = nn.Linear(config.hidden_size, config.entity_emb_size)752 if isinstance(config.hidden_act, str):753 self.transform_act_fn = ACT2FN[config.hidden_act]754 else:755 self.transform_act_fn = config.hidden_act756 self.LayerNorm = nn.LayerNorm(config.entity_emb_size, eps=config.layer_norm_eps)757 758 def forward(self, hidden_states):759 hidden_states = self.dense(hidden_states)760 hidden_states = self.transform_act_fn(hidden_states)761 hidden_states = self.LayerNorm(hidden_states)762 return hidden_states763 764 765class EntityPredictionHead(nn.Module):766 def __init__(self, config):767 super().__init__()768 self.config = config769 self.transform = EntityPredictionHeadTransform(config)770 self.decoder = nn.Linear(config.entity_emb_size, config.entity_vocab_size, bias=False)771 self.bias = nn.Parameter(torch.zeros(config.entity_vocab_size))772 773 def forward(self, hidden_states):774 hidden_states = self.transform(hidden_states)775 hidden_states = self.decoder(hidden_states) + self.bias776 777 return hidden_states778 779 780@auto_docstring781class LukePreTrainedModel(PreTrainedModel):782 config: LukeConfig783 base_model_prefix = "luke"784 supports_gradient_checkpointing = True785 _no_split_modules = ["LukeAttention", "LukeEntityEmbeddings"]786 787 def _init_weights(self, module: nn.Module):788 """Initialize the weights"""789 if isinstance(module, nn.Linear):790 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)791 if module.bias is not None:792 module.bias.data.zero_()793 elif isinstance(module, nn.Embedding):794 if module.embedding_dim == 1: # embedding for bias parameters795 module.weight.data.zero_()796 else:797 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)798 if module.padding_idx is not None:799 module.weight.data[module.padding_idx].zero_()800 elif isinstance(module, nn.LayerNorm):801 module.bias.data.zero_()802 module.weight.data.fill_(1.0)803 804 805@auto_docstring(806 custom_intro="""807 The bare LUKE model transformer outputting raw hidden-states for both word tokens and entities without any808 """809)810class LukeModel(LukePreTrainedModel):811 def __init__(self, config: LukeConfig, add_pooling_layer: bool = True):812 r"""813 add_pooling_layer (bool, *optional*, defaults to `True`):814 Whether to add a pooling layer815 """816 super().__init__(config)817 self.config = config818 819 self.embeddings = LukeEmbeddings(config)820 self.entity_embeddings = LukeEntityEmbeddings(config)821 self.encoder = LukeEncoder(config)822 823 self.pooler = LukePooler(config) if add_pooling_layer else None824 825 # Initialize weights and apply final processing826 self.post_init()827 828 def get_input_embeddings(self):829 return self.embeddings.word_embeddings830 831 def set_input_embeddings(self, value):832 self.embeddings.word_embeddings = value833 834 def get_entity_embeddings(self):835 return self.entity_embeddings.entity_embeddings836 837 def set_entity_embeddings(self, value):838 self.entity_embeddings.entity_embeddings = value839 840 def _prune_heads(self, heads_to_prune):841 raise NotImplementedError("LUKE does not support the pruning of attention heads")842 843 @auto_docstring844 def forward(845 self,846 input_ids: Optional[torch.LongTensor] = None,847 attention_mask: Optional[torch.FloatTensor] = None,848 token_type_ids: Optional[torch.LongTensor] = None,849 position_ids: Optional[torch.LongTensor] = None,850 entity_ids: Optional[torch.LongTensor] = None,851 entity_attention_mask: Optional[torch.FloatTensor] = None,852 entity_token_type_ids: Optional[torch.LongTensor] = None,853 entity_position_ids: Optional[torch.LongTensor] = None,854 head_mask: Optional[torch.FloatTensor] = None,855 inputs_embeds: Optional[torch.FloatTensor] = None,856 output_attentions: Optional[bool] = None,857 output_hidden_states: Optional[bool] = None,858 return_dict: Optional[bool] = None,859 ) -> Union[tuple, BaseLukeModelOutputWithPooling]:860 r"""861 entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`):862 Indices of entity tokens in the entity vocabulary.863 864 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and865 [`PreTrainedTokenizer.__call__`] for details.866 entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*):867 Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`:868 869 - 1 for entity tokens that are **not masked**,870 - 0 for entity tokens that are **masked**.871 entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*):872 Segment token indices to indicate first and second portions of the entity token inputs. Indices are873 selected in `[0, 1]`:874 875 - 0 corresponds to a *portion A* entity token,876 - 1 corresponds to a *portion B* entity token.877 entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*):878 Indices of positions of each input entity in the position embeddings. Selected in the range `[0,879 config.max_position_embeddings - 1]`.880 881 Examples:882 883 ```python884 >>> from transformers import AutoTokenizer, LukeModel885 886 >>> tokenizer = AutoTokenizer.from_pretrained("studio-ousia/luke-base")887 >>> model = LukeModel.from_pretrained("studio-ousia/luke-base")888 # Compute the contextualized entity representation corresponding to the entity mention "Beyoncé"889 890 >>> text = "Beyoncé lives in Los Angeles."891 >>> entity_spans = [(0, 7)] # character-based entity span corresponding to "Beyoncé"892 893 >>> encoding = tokenizer(text, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt")894 >>> outputs = model(**encoding)895 >>> word_last_hidden_state = outputs.last_hidden_state896 >>> entity_last_hidden_state = outputs.entity_last_hidden_state897 # Input Wikipedia entities to obtain enriched contextualized representations of word tokens898 899 >>> text = "Beyoncé lives in Los Angeles."900 >>> entities = [901 ... "Beyoncé",902 ... "Los Angeles",903 ... ] # Wikipedia entity titles corresponding to the entity mentions "Beyoncé" and "Los Angeles"904 >>> entity_spans = [905 ... (0, 7),906 ... (17, 28),907 ... ] # character-based entity spans corresponding to "Beyoncé" and "Los Angeles"908 909 >>> encoding = tokenizer(910 ... text, entities=entities, entity_spans=entity_spans, add_prefix_space=True, return_tensors="pt"911 ... )912 >>> outputs = model(**encoding)913 >>> word_last_hidden_state = outputs.last_hidden_state914 >>> entity_last_hidden_state = outputs.entity_last_hidden_state915 ```"""916 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions917 output_hidden_states = (918 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states919 )920 return_dict = return_dict if return_dict is not None else self.config.use_return_dict921 922 if input_ids is not None and inputs_embeds is not None:923 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")924 elif input_ids is not None:925 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)926 input_shape = input_ids.size()927 elif inputs_embeds is not None:928 input_shape = inputs_embeds.size()[:-1]929 else:930 raise ValueError("You have to specify either input_ids or inputs_embeds")931 932 batch_size, seq_length = input_shape933 device = input_ids.device if input_ids is not None else inputs_embeds.device934 935 if attention_mask is None:936 attention_mask = torch.ones((batch_size, seq_length), device=device)937 if token_type_ids is None:938 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)939 if entity_ids is not None:940 entity_seq_length = entity_ids.size(1)941 if entity_attention_mask is None:942 entity_attention_mask = torch.ones((batch_size, entity_seq_length), device=device)943 if entity_token_type_ids is None:944 entity_token_type_ids = torch.zeros((batch_size, entity_seq_length), dtype=torch.long, device=device)945 946 # Prepare head mask if needed947 # 1.0 in head_mask indicate we keep the head948 # attention_probs has shape bsz x n_heads x N x N949 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]950 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]951 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)952 953 # First, compute word embeddings954 word_embedding_output = self.embeddings(955 input_ids=input_ids,956 position_ids=position_ids,957 token_type_ids=token_type_ids,958 inputs_embeds=inputs_embeds,959 )960 961 # Second, compute extended attention mask962 extended_attention_mask = self.get_extended_attention_mask(attention_mask, entity_attention_mask)963 964 # Third, compute entity embeddings and concatenate with word embeddings965 if entity_ids is None:966 entity_embedding_output = None967 else:968 entity_embedding_output = self.entity_embeddings(entity_ids, entity_position_ids, entity_token_type_ids)969 970 # Fourth, send embeddings through the model971 encoder_outputs = self.encoder(972 word_embedding_output,973 entity_embedding_output,974 attention_mask=extended_attention_mask,975 head_mask=head_mask,976 output_attentions=output_attentions,977 output_hidden_states=output_hidden_states,978 return_dict=return_dict,979 )980 981 # Fifth, get the output. LukeModel outputs the same as BertModel, namely sequence_output of shape (batch_size, seq_len, hidden_size)982 sequence_output = encoder_outputs[0]983 984 # Sixth, we compute the pooled_output, word_sequence_output and entity_sequence_output based on the sequence_output985 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None986 987 if not return_dict:988 return (sequence_output, pooled_output) + encoder_outputs[1:]989 990 return BaseLukeModelOutputWithPooling(991 last_hidden_state=sequence_output,992 pooler_output=pooled_output,993 hidden_states=encoder_outputs.hidden_states,994 attentions=encoder_outputs.attentions,995 entity_last_hidden_state=encoder_outputs.entity_last_hidden_state,996 entity_hidden_states=encoder_outputs.entity_hidden_states,997 )998 999 def get_extended_attention_mask(1000 self, word_attention_mask: torch.LongTensor, entity_attention_mask: Optional[torch.LongTensor]1001 ):1002 """1003 Makes broadcastable attention and causal masks so that future and masked tokens are ignored.1004 1005 Arguments:1006 word_attention_mask (`torch.LongTensor`):1007 Attention mask for word tokens with ones indicating tokens to attend to, zeros for tokens to ignore.1008 entity_attention_mask (`torch.LongTensor`, *optional*):1009 Attention mask for entity tokens with ones indicating tokens to attend to, zeros for tokens to ignore.1010 1011 Returns:1012 `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.1013 """1014 attention_mask = word_attention_mask1015 if entity_attention_mask is not None:1016 attention_mask = torch.cat([attention_mask, entity_attention_mask], dim=-1)1017 1018 if attention_mask.dim() == 3:1019 extended_attention_mask = attention_mask[:, None, :, :]1020 elif attention_mask.dim() == 2:1021 extended_attention_mask = attention_mask[:, None, None, :]1022 else:1023 raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape})")1024 1025 extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility1026 extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min1027 return extended_attention_mask1028 1029 1030def create_position_ids_from_input_ids(input_ids, padding_idx):1031 """1032 Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols1033 are ignored. This is modified from fairseq's `utils.make_positions`.1034 1035 Args:1036 x: torch.Tensor x:1037 1038 Returns: torch.Tensor1039 """1040 # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.1041 mask = input_ids.ne(padding_idx).int()1042 incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask)) * mask1043 return incremental_indices.long() + padding_idx1044 1045 1046# Copied from transformers.models.roberta.modeling_roberta.RobertaLMHead1047class LukeLMHead(nn.Module):1048 """Roberta Head for masked language modeling."""1049 1050 def __init__(self, config):1051 super().__init__()1052 self.dense = nn.Linear(config.hidden_size, config.hidden_size)1053 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1054 1055 self.decoder = nn.Linear(config.hidden_size, config.vocab_size)1056 self.bias = nn.Parameter(torch.zeros(config.vocab_size))1057 self.decoder.bias = self.bias1058 1059 def forward(self, features, **kwargs):1060 x = self.dense(features)1061 x = gelu(x)1062 x = self.layer_norm(x)1063 1064 # project back to size of vocabulary with bias1065 x = self.decoder(x)1066 1067 return x1068 1069 def _tie_weights(self):1070 # To tie those two weights if they get disconnected (on TPU or when the bias is resized)1071 # For accelerate compatibility and to not break backward compatibility1072 if self.decoder.bias.device.type == "meta":1073 self.decoder.bias = self.bias1074 else:1075 self.bias = self.decoder.bias1076 1077 1078@auto_docstring(1079 custom_intro="""1080 The LUKE model with a language modeling head and entity prediction head on top for masked language modeling and1081 masked entity prediction.1082 """1083)1084class LukeForMaskedLM(LukePreTrainedModel):1085 _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias", "entity_predictions.decoder.weight"]1086 1087 def __init__(self, config):1088 super().__init__(config)1089 1090 self.luke = LukeModel(config)1091 1092 self.lm_head = LukeLMHead(config)1093 self.entity_predictions = EntityPredictionHead(config)1094 1095 self.loss_fn = nn.CrossEntropyLoss()1096 1097 # Initialize weights and apply final processing1098 self.post_init()1099 1100 def tie_weights(self):1101 super().tie_weights()1102 self._tie_or_clone_weights(self.entity_predictions.decoder, self.luke.entity_embeddings.entity_embeddings)1103 1104 def get_output_embeddings(self):1105 return self.lm_head.decoder1106 1107 def set_output_embeddings(self, new_embeddings):1108 self.lm_head.decoder = new_embeddings1109 1110 @auto_docstring1111 def forward(1112 self,1113 input_ids: Optional[torch.LongTensor] = None,1114 attention_mask: Optional[torch.FloatTensor] = None,1115 token_type_ids: Optional[torch.LongTensor] = None,1116 position_ids: Optional[torch.LongTensor] = None,1117 entity_ids: Optional[torch.LongTensor] = None,1118 entity_attention_mask: Optional[torch.LongTensor] = None,1119 entity_token_type_ids: Optional[torch.LongTensor] = None,1120 entity_position_ids: Optional[torch.LongTensor] = None,1121 labels: Optional[torch.LongTensor] = None,1122 entity_labels: Optional[torch.LongTensor] = None,1123 head_mask: Optional[torch.FloatTensor] = None,1124 inputs_embeds: Optional[torch.FloatTensor] = None,1125 output_attentions: Optional[bool] = None,1126 output_hidden_states: Optional[bool] = None,1127 return_dict: Optional[bool] = None,1128 ) -> Union[tuple, LukeMaskedLMOutput]:1129 r"""1130 entity_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`):1131 Indices of entity tokens in the entity vocabulary.1132 1133 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and1134 [`PreTrainedTokenizer.__call__`] for details.1135 entity_attention_mask (`torch.FloatTensor` of shape `(batch_size, entity_length)`, *optional*):1136 Mask to avoid performing attention on padding entity token indices. Mask values selected in `[0, 1]`:1137 1138 - 1 for entity tokens that are **not masked**,1139 - 0 for entity tokens that are **masked**.1140 entity_token_type_ids (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*):1141 Segment token indices to indicate first and second portions of the entity token inputs. Indices are1142 selected in `[0, 1]`:1143 1144 - 0 corresponds to a *portion A* entity token,1145 - 1 corresponds to a *portion B* entity token.1146 entity_position_ids (`torch.LongTensor` of shape `(batch_size, entity_length, max_mention_length)`, *optional*):1147 Indices of positions of each input entity in the position embeddings. Selected in the range `[0,1148 config.max_position_embeddings - 1]`.1149 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1150 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1151 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the1152 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1153 entity_labels (`torch.LongTensor` of shape `(batch_size, entity_length)`, *optional*):1154 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1155 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the1156 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1157 """1158 1159 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1160 1161 outputs = self.luke(1162 input_ids=input_ids,1163 attention_mask=attention_mask,1164 token_type_ids=token_type_ids,1165 position_ids=position_ids,1166 entity_ids=entity_ids,1167 entity_attention_mask=entity_attention_mask,1168 entity_token_type_ids=entity_token_type_ids,1169 entity_position_ids=entity_position_ids,1170 head_mask=head_mask,1171 inputs_embeds=inputs_embeds,1172 output_attentions=output_attentions,1173 output_hidden_states=output_hidden_states,1174 return_dict=True,1175 )1176 1177 loss = None1178 1179 mlm_loss = None1180 logits = self.lm_head(outputs.last_hidden_state)1181 if labels is not None:1182 # move labels to correct device to enable model parallelism1183 labels = labels.to(logits.device)1184 mlm_loss = self.loss_fn(logits.view(-1, self.config.vocab_size), labels.view(-1))1185 if loss is None:1186 loss = mlm_loss1187 1188 mep_loss = None1189 entity_logits = None1190 if outputs.entity_last_hidden_state is not None:1191 entity_logits = self.entity_predictions(outputs.entity_last_hidden_state)1192 if entity_labels is not None:1193 mep_loss = self.loss_fn(entity_logits.view(-1, self.config.entity_vocab_size), entity_labels.view(-1))1194 if loss is None:1195 loss = mep_loss1196 else:1197 loss = loss + mep_loss1198 1199 if not return_dict:1200 return tuple(