Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023-present NAVER Corp, The Microsoft Research Asia LayoutLM Team Authors 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 Bros model."""16 17import math18from dataclasses import dataclass19from typing import Optional, Union20 21import torch22from torch import nn23from torch.nn import CrossEntropyLoss24 25from ...activations import ACT2FN26from ...modeling_layers import GradientCheckpointingLayer27from ...modeling_outputs import (28 BaseModelOutputWithCrossAttentions,29 BaseModelOutputWithPoolingAndCrossAttentions,30 TokenClassifierOutput,31)32from ...modeling_utils import PreTrainedModel33from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer34from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging35from .configuration_bros import BrosConfig36 37 38logger = logging.get_logger(__name__)39 40 41@dataclass42@auto_docstring(43 custom_intro="""44 Base class for outputs of token classification models.45 """46)47class BrosSpadeOutput(ModelOutput):48 r"""49 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):50 Classification loss.51 initial_token_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):52 Classification scores for entity initial tokens (before SoftMax).53 subsequent_token_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, sequence_length+1)`):54 Classification scores for entity sequence tokens (before SoftMax).55 """56 57 loss: Optional[torch.FloatTensor] = None58 initial_token_logits: Optional[torch.FloatTensor] = None59 subsequent_token_logits: Optional[torch.FloatTensor] = None60 hidden_states: Optional[tuple[torch.FloatTensor]] = None61 attentions: Optional[tuple[torch.FloatTensor]] = None62 63 64class BrosPositionalEmbedding1D(nn.Module):65 # Reference: https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py#L1566 67 def __init__(self, config):68 super().__init__()69 70 self.dim_bbox_sinusoid_emb_1d = config.dim_bbox_sinusoid_emb_1d71 72 inv_freq = 1 / (73 10000 ** (torch.arange(0.0, self.dim_bbox_sinusoid_emb_1d, 2.0) / self.dim_bbox_sinusoid_emb_1d)74 )75 self.register_buffer("inv_freq", inv_freq)76 77 def forward(self, pos_seq: torch.Tensor) -> torch.Tensor:78 seq_size = pos_seq.size()79 b1, b2, b3 = seq_size80 sinusoid_inp = pos_seq.view(b1, b2, b3, 1) * self.inv_freq.view(1, 1, 1, self.dim_bbox_sinusoid_emb_1d // 2)81 pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)82 return pos_emb83 84 85class BrosPositionalEmbedding2D(nn.Module):86 def __init__(self, config):87 super().__init__()88 89 self.dim_bbox = config.dim_bbox90 self.x_pos_emb = BrosPositionalEmbedding1D(config)91 self.y_pos_emb = BrosPositionalEmbedding1D(config)92 93 def forward(self, bbox: torch.Tensor) -> torch.Tensor:94 stack = []95 for i in range(self.dim_bbox):96 if i % 2 == 0:97 stack.append(self.x_pos_emb(bbox[..., i]))98 else:99 stack.append(self.y_pos_emb(bbox[..., i]))100 bbox_pos_emb = torch.cat(stack, dim=-1)101 return bbox_pos_emb102 103 104class BrosBboxEmbeddings(nn.Module):105 def __init__(self, config):106 super().__init__()107 self.bbox_sinusoid_emb = BrosPositionalEmbedding2D(config)108 self.bbox_projection = nn.Linear(config.dim_bbox_sinusoid_emb_2d, config.dim_bbox_projection, bias=False)109 110 def forward(self, bbox: torch.Tensor):111 bbox_t = bbox.transpose(0, 1)112 bbox_pos = bbox_t[None, :, :, :] - bbox_t[:, None, :, :]113 bbox_pos_emb = self.bbox_sinusoid_emb(bbox_pos)114 bbox_pos_emb = self.bbox_projection(bbox_pos_emb)115 116 return bbox_pos_emb117 118 119class BrosTextEmbeddings(nn.Module):120 """Construct the embeddings from word, position and token_type embeddings."""121 122 def __init__(self, config):123 super().__init__()124 125 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)126 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)127 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)128 129 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load130 # any TensorFlow checkpoint file131 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)132 self.dropout = nn.Dropout(config.hidden_dropout_prob)133 # position_ids (1, len position emb) is contiguous in memory and exported when serialized134 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")135 self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))136 self.register_buffer(137 "token_type_ids",138 torch.zeros(139 self.position_ids.size(),140 dtype=torch.long,141 device=self.position_ids.device,142 ),143 persistent=False,144 )145 146 def forward(147 self,148 input_ids: Optional[torch.Tensor] = None,149 token_type_ids: Optional[torch.Tensor] = None,150 position_ids: Optional[torch.Tensor] = None,151 inputs_embeds: Optional[torch.Tensor] = None,152 ) -> torch.Tensor:153 if input_ids is not None:154 input_shape = input_ids.size()155 else:156 input_shape = inputs_embeds.size()[:-1]157 158 seq_length = input_shape[1]159 160 if position_ids is None:161 position_ids = self.position_ids[:, :seq_length]162 163 if token_type_ids is None:164 if hasattr(self, "token_type_ids"):165 buffered_token_type_ids = self.token_type_ids[:, :seq_length]166 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)167 token_type_ids = buffered_token_type_ids_expanded168 else:169 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)170 171 if inputs_embeds is None:172 inputs_embeds = self.word_embeddings(input_ids)173 token_type_embeddings = self.token_type_embeddings(token_type_ids)174 175 embeddings = inputs_embeds + token_type_embeddings176 if self.position_embedding_type == "absolute":177 position_embeddings = self.position_embeddings(position_ids)178 embeddings += position_embeddings179 embeddings = self.LayerNorm(embeddings)180 embeddings = self.dropout(embeddings)181 return embeddings182 183 184class BrosSelfAttention(nn.Module):185 def __init__(self, config):186 super().__init__()187 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):188 raise ValueError(189 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "190 f"heads ({config.num_attention_heads})"191 )192 193 self.num_attention_heads = config.num_attention_heads194 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)195 self.all_head_size = self.num_attention_heads * self.attention_head_size196 197 self.query = nn.Linear(config.hidden_size, self.all_head_size)198 self.key = nn.Linear(config.hidden_size, self.all_head_size)199 self.value = nn.Linear(config.hidden_size, self.all_head_size)200 201 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)202 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")203 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":204 self.max_position_embeddings = config.max_position_embeddings205 self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)206 207 self.is_decoder = config.is_decoder208 209 def forward(210 self,211 hidden_states: torch.Tensor,212 bbox_pos_emb: torch.Tensor,213 attention_mask: Optional[torch.Tensor] = None,214 head_mask: Optional[torch.Tensor] = None,215 encoder_hidden_states: Optional[torch.Tensor] = None,216 encoder_attention_mask: Optional[torch.Tensor] = None,217 output_attentions: Optional[torch.Tensor] = False,218 ) -> tuple[torch.Tensor]:219 hidden_shape = (hidden_states.shape[0], -1, self.num_attention_heads, self.attention_head_size)220 query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2)221 222 # If this is instantiated as a cross-attention module, the keys223 # and values come from an encoder; the attention mask needs to be224 # such that the encoder's padding tokens are not attended to.225 is_cross_attention = encoder_hidden_states is not None226 227 if is_cross_attention:228 key_layer = self.key(encoder_hidden_states).view(hidden_shape).transpose(1, 2)229 value_layer = self.value(encoder_hidden_states).view(hidden_shape).transpose(1, 2)230 attention_mask = encoder_attention_mask231 else:232 key_layer = self.key(hidden_states).view(hidden_shape).transpose(1, 2)233 value_layer = self.value(hidden_states).view(hidden_shape).transpose(1, 2)234 235 # Take the dot product between "query" and "key" to get the raw attention scores.236 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))237 238 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":239 seq_length = hidden_states.size()[1]240 position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)241 position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)242 distance = position_ids_l - position_ids_r243 positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)244 positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility245 246 if self.position_embedding_type == "relative_key":247 relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)248 attention_scores = attention_scores + relative_position_scores249 elif self.position_embedding_type == "relative_key_query":250 relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)251 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)252 253 attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key254 255 # bbox positional encoding256 batch_size, n_head, seq_length, d_head = query_layer.shape257 bbox_pos_emb = bbox_pos_emb.view(seq_length, seq_length, batch_size, d_head)258 bbox_pos_emb = bbox_pos_emb.permute([2, 0, 1, 3])259 bbox_pos_scores = torch.einsum("bnid,bijd->bnij", (query_layer, bbox_pos_emb))260 261 attention_scores = attention_scores + bbox_pos_scores262 263 attention_scores = attention_scores / math.sqrt(self.attention_head_size)264 if attention_mask is not None:265 # Apply the attention mask is (precomputed for all layers in BrosModel forward() function)266 attention_scores = attention_scores + attention_mask267 268 # Normalize the attention scores to probabilities.269 attention_probs = nn.Softmax(dim=-1)(attention_scores)270 271 # This is actually dropping out entire tokens to attend to, which might272 # seem a bit unusual, but is taken from the original Transformer paper.273 attention_probs = self.dropout(attention_probs)274 275 # Mask heads if we want to276 if head_mask is not None:277 attention_probs = attention_probs * head_mask278 279 context_layer = torch.matmul(attention_probs, value_layer)280 281 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()282 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)283 context_layer = context_layer.view(*new_context_layer_shape)284 285 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)286 287 if self.is_decoder:288 outputs = outputs + (None,)289 return outputs290 291 292# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->Bros293class BrosSelfOutput(nn.Module):294 def __init__(self, config):295 super().__init__()296 self.dense = nn.Linear(config.hidden_size, config.hidden_size)297 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)298 self.dropout = nn.Dropout(config.hidden_dropout_prob)299 300 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:301 hidden_states = self.dense(hidden_states)302 hidden_states = self.dropout(hidden_states)303 hidden_states = self.LayerNorm(hidden_states + input_tensor)304 return hidden_states305 306 307class BrosAttention(nn.Module):308 def __init__(self, config):309 super().__init__()310 self.self = BrosSelfAttention(config)311 self.output = BrosSelfOutput(config)312 self.pruned_heads = set()313 314 def prune_heads(self, heads):315 if len(heads) == 0:316 return317 heads, index = find_pruneable_heads_and_indices(318 heads,319 self.self.num_attention_heads,320 self.self.attention_head_size,321 self.pruned_heads,322 )323 324 # Prune linear layers325 self.self.query = prune_linear_layer(self.self.query, index)326 self.self.key = prune_linear_layer(self.self.key, index)327 self.self.value = prune_linear_layer(self.self.value, index)328 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)329 330 # Update hyper params and store pruned heads331 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)332 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads333 self.pruned_heads = self.pruned_heads.union(heads)334 335 def forward(336 self,337 hidden_states: torch.Tensor,338 bbox_pos_emb: torch.Tensor,339 attention_mask: Optional[torch.Tensor] = None,340 head_mask: Optional[torch.Tensor] = None,341 encoder_hidden_states: Optional[torch.Tensor] = None,342 encoder_attention_mask: Optional[torch.Tensor] = None,343 output_attentions: Optional[bool] = False,344 ) -> tuple[torch.Tensor]:345 self_outputs = self.self(346 hidden_states=hidden_states,347 bbox_pos_emb=bbox_pos_emb,348 attention_mask=attention_mask,349 head_mask=head_mask,350 encoder_hidden_states=encoder_hidden_states,351 encoder_attention_mask=encoder_attention_mask,352 output_attentions=output_attentions,353 )354 attention_output = self.output(self_outputs[0], hidden_states)355 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them356 return outputs357 358 359# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Bros360class BrosIntermediate(nn.Module):361 def __init__(self, config):362 super().__init__()363 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)364 if isinstance(config.hidden_act, str):365 self.intermediate_act_fn = ACT2FN[config.hidden_act]366 else:367 self.intermediate_act_fn = config.hidden_act368 369 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:370 hidden_states = self.dense(hidden_states)371 hidden_states = self.intermediate_act_fn(hidden_states)372 return hidden_states373 374 375class BrosOutput(nn.Module):376 def __init__(self, config):377 super().__init__()378 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)379 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)380 self.dropout = nn.Dropout(config.hidden_dropout_prob)381 382 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:383 hidden_states = self.dense(hidden_states)384 hidden_states = self.dropout(hidden_states)385 hidden_states = self.LayerNorm(hidden_states + input_tensor)386 return hidden_states387 388 389class BrosLayer(GradientCheckpointingLayer):390 def __init__(self, config):391 super().__init__()392 self.chunk_size_feed_forward = config.chunk_size_feed_forward393 self.seq_len_dim = 1394 self.attention = BrosAttention(config)395 self.is_decoder = config.is_decoder396 self.add_cross_attention = config.add_cross_attention397 if self.add_cross_attention:398 if not self.is_decoder:399 raise Exception(f"{self} should be used as a decoder model if cross attention is added")400 self.crossattention = BrosAttention(config)401 self.intermediate = BrosIntermediate(config)402 self.output = BrosOutput(config)403 404 def forward(405 self,406 hidden_states: torch.Tensor,407 bbox_pos_emb: torch.Tensor,408 attention_mask: Optional[torch.FloatTensor] = None,409 head_mask: Optional[torch.FloatTensor] = None,410 encoder_hidden_states: Optional[torch.FloatTensor] = None,411 encoder_attention_mask: Optional[torch.FloatTensor] = None,412 output_attentions: Optional[bool] = False,413 ) -> tuple[torch.Tensor]:414 self_attention_outputs = self.attention(415 hidden_states,416 bbox_pos_emb=bbox_pos_emb,417 attention_mask=attention_mask,418 head_mask=head_mask,419 output_attentions=output_attentions,420 )421 attention_output = self_attention_outputs[0]422 423 # if decoder, the last output is tuple of self-attn cache424 if self.is_decoder:425 outputs = self_attention_outputs[1:-1]426 else:427 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights428 429 if self.is_decoder and encoder_hidden_states is not None:430 if hasattr(self, "crossattention"):431 raise Exception(432 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"433 )434 435 cross_attention_outputs = self.crossattention(436 attention_output,437 attention_mask=attention_mask,438 head_mask=head_mask,439 encoder_hidden_states=encoder_hidden_states,440 encoder_attention_mask=encoder_attention_mask,441 output_attentions=output_attentions,442 )443 attention_output = cross_attention_outputs[0]444 outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights445 446 layer_output = apply_chunking_to_forward(447 self.feed_forward_chunk,448 self.chunk_size_feed_forward,449 self.seq_len_dim,450 attention_output,451 )452 outputs = (layer_output,) + outputs453 454 # if decoder, return the attn key/values as the last output455 if self.is_decoder:456 outputs = outputs + (None,)457 458 return outputs459 460 def feed_forward_chunk(self, attention_output):461 intermediate_output = self.intermediate(attention_output)462 layer_output = self.output(intermediate_output, attention_output)463 return layer_output464 465 466class BrosEncoder(nn.Module):467 def __init__(self, config):468 super().__init__()469 self.config = config470 self.layer = nn.ModuleList([BrosLayer(config) for _ in range(config.num_hidden_layers)])471 472 @can_return_tuple473 def forward(474 self,475 hidden_states: torch.Tensor,476 bbox_pos_emb: torch.Tensor,477 attention_mask: Optional[torch.FloatTensor] = None,478 head_mask: Optional[torch.FloatTensor] = None,479 encoder_hidden_states: Optional[torch.FloatTensor] = None,480 encoder_attention_mask: Optional[torch.FloatTensor] = None,481 output_attentions: Optional[bool] = False,482 output_hidden_states: Optional[bool] = False,483 return_dict: Optional[bool] = True,484 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithCrossAttentions]:485 all_hidden_states = () if output_hidden_states else None486 all_self_attentions = () if output_attentions else None487 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None488 489 for i, layer_module in enumerate(self.layer):490 if output_hidden_states:491 all_hidden_states = all_hidden_states + (hidden_states,)492 493 layer_head_mask = head_mask[i] if head_mask is not None else None494 495 layer_outputs = layer_module(496 hidden_states=hidden_states,497 bbox_pos_emb=bbox_pos_emb,498 attention_mask=attention_mask,499 head_mask=layer_head_mask,500 encoder_hidden_states=encoder_hidden_states,501 encoder_attention_mask=encoder_attention_mask,502 output_attentions=output_attentions,503 )504 505 hidden_states = layer_outputs[0]506 if output_attentions:507 all_self_attentions = all_self_attentions + (layer_outputs[1],)508 if self.config.add_cross_attention:509 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)510 511 if output_hidden_states:512 all_hidden_states = all_hidden_states + (hidden_states,)513 514 return BaseModelOutputWithCrossAttentions(515 last_hidden_state=hidden_states,516 hidden_states=all_hidden_states,517 attentions=all_self_attentions,518 cross_attentions=all_cross_attentions,519 )520 521 522# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->Bros523class BrosPooler(nn.Module):524 def __init__(self, config):525 super().__init__()526 self.dense = nn.Linear(config.hidden_size, config.hidden_size)527 self.activation = nn.Tanh()528 529 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:530 # We "pool" the model by simply taking the hidden state corresponding531 # to the first token.532 first_token_tensor = hidden_states[:, 0]533 pooled_output = self.dense(first_token_tensor)534 pooled_output = self.activation(pooled_output)535 return pooled_output536 537 538class BrosRelationExtractor(nn.Module):539 def __init__(self, config):540 super().__init__()541 self.n_relations = config.n_relations542 self.backbone_hidden_size = config.hidden_size543 self.head_hidden_size = config.hidden_size544 self.classifier_dropout_prob = config.classifier_dropout_prob545 546 self.drop = nn.Dropout(self.classifier_dropout_prob)547 self.query = nn.Linear(self.backbone_hidden_size, self.n_relations * self.head_hidden_size)548 549 self.key = nn.Linear(self.backbone_hidden_size, self.n_relations * self.head_hidden_size)550 551 self.dummy_node = nn.Parameter(torch.zeros(1, self.backbone_hidden_size))552 553 def forward(self, query_layer: torch.Tensor, key_layer: torch.Tensor):554 query_layer = self.query(self.drop(query_layer))555 556 dummy_vec = self.dummy_node.unsqueeze(0).repeat(1, key_layer.size(1), 1)557 key_layer = torch.cat([key_layer, dummy_vec], axis=0)558 key_layer = self.key(self.drop(key_layer))559 560 query_layer = query_layer.view(561 query_layer.size(0), query_layer.size(1), self.n_relations, self.head_hidden_size562 )563 key_layer = key_layer.view(key_layer.size(0), key_layer.size(1), self.n_relations, self.head_hidden_size)564 565 relation_score = torch.matmul(566 query_layer.permute(2, 1, 0, 3), key_layer.permute(2, 1, 3, 0)567 ) # equivalent to torch.einsum("ibnd,jbnd->nbij", (query_layer, key_layer))568 569 return relation_score570 571 572@auto_docstring573class BrosPreTrainedModel(PreTrainedModel):574 config: BrosConfig575 base_model_prefix = "bros"576 577 def _init_weights(self, module: nn.Module):578 """Initialize the weights"""579 std = self.config.initializer_range580 if isinstance(module, nn.Linear):581 # Slightly different from the TF version which uses truncated_normal for initialization582 # cf https://github.com/pytorch/pytorch/pull/5617583 module.weight.data.normal_(mean=0.0, std=std)584 if module.bias is not None:585 module.bias.data.zero_()586 elif isinstance(module, nn.Embedding):587 module.weight.data.normal_(mean=0.0, std=std)588 if module.padding_idx is not None:589 module.weight.data[module.padding_idx].zero_()590 elif isinstance(module, nn.LayerNorm):591 module.bias.data.zero_()592 module.weight.data.fill_(1.0)593 elif isinstance(module, BrosRelationExtractor):594 nn.init.normal_(module.dummy_node, std=std)595 596 597@auto_docstring598class BrosModel(BrosPreTrainedModel):599 def __init__(self, config, add_pooling_layer=True):600 r"""601 add_pooling_layer (bool, *optional*, defaults to `True`):602 Whether to add a pooling layer603 """604 super().__init__(config)605 self.config = config606 607 self.embeddings = BrosTextEmbeddings(config)608 self.bbox_embeddings = BrosBboxEmbeddings(config)609 self.encoder = BrosEncoder(config)610 611 self.pooler = BrosPooler(config) if add_pooling_layer else None612 613 self.init_weights()614 615 def get_input_embeddings(self):616 return self.embeddings.word_embeddings617 618 def set_input_embeddings(self, value):619 self.embeddings.word_embeddings = value620 621 def _prune_heads(self, heads_to_prune):622 """623 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base624 class PreTrainedModel625 """626 for layer, heads in heads_to_prune.items():627 self.encoder.layer[layer].attention.prune_heads(heads)628 629 @can_return_tuple630 @auto_docstring631 def forward(632 self,633 input_ids: Optional[torch.Tensor] = None,634 bbox: Optional[torch.Tensor] = None,635 attention_mask: Optional[torch.Tensor] = None,636 token_type_ids: Optional[torch.Tensor] = None,637 position_ids: Optional[torch.Tensor] = None,638 head_mask: Optional[torch.Tensor] = None,639 inputs_embeds: Optional[torch.Tensor] = None,640 encoder_hidden_states: Optional[torch.Tensor] = None,641 encoder_attention_mask: Optional[torch.Tensor] = None,642 output_attentions: Optional[bool] = None,643 output_hidden_states: Optional[bool] = None,644 return_dict: Optional[bool] = None,645 ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:646 r"""647 bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'):648 Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values649 (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the650 bounding box.651 652 Examples:653 654 ```python655 >>> import torch656 >>> from transformers import BrosProcessor, BrosModel657 658 >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased")659 660 >>> model = BrosModel.from_pretrained("jinho8345/bros-base-uncased")661 662 >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")663 >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1)664 >>> encoding["bbox"] = bbox665 666 >>> outputs = model(**encoding)667 >>> last_hidden_states = outputs.last_hidden_state668 ```"""669 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions670 output_hidden_states = (671 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states672 )673 return_dict = return_dict if return_dict is not None else self.config.use_return_dict674 675 if input_ids is not None and inputs_embeds is not None:676 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")677 elif input_ids is not None:678 input_shape = input_ids.size()679 elif inputs_embeds is not None:680 input_shape = inputs_embeds.size()[:-1]681 else:682 raise ValueError("You have to specify either input_ids or inputs_embeds")683 684 if bbox is None:685 raise ValueError("You have to specify bbox")686 687 batch_size, seq_length = input_shape688 device = input_ids.device if input_ids is not None else inputs_embeds.device689 690 if attention_mask is None:691 attention_mask = torch.ones(input_shape, device=device)692 693 if token_type_ids is None:694 if hasattr(self.embeddings, "token_type_ids"):695 buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]696 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)697 token_type_ids = buffered_token_type_ids_expanded698 else:699 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)700 701 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]702 # ourselves in which case we just need to make it broadcastable to all heads.703 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)704 705 # If a 2D or 3D attention mask is provided for the cross-attention706 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]707 if self.config.is_decoder and encoder_hidden_states is not None:708 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()709 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)710 if encoder_attention_mask is None:711 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)712 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)713 else:714 encoder_extended_attention_mask = None715 716 # Prepare head mask if needed717 # 1.0 in head_mask indicate we keep the head718 # attention_probs has shape bsz x n_heads x N x N719 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]720 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]721 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)722 723 embedding_output = self.embeddings(724 input_ids=input_ids,725 position_ids=position_ids,726 token_type_ids=token_type_ids,727 inputs_embeds=inputs_embeds,728 )729 730 # if bbox has 2 points (4 float tensors) per token, convert it to 4 points (8 float tensors) per token731 if bbox.shape[-1] == 4:732 bbox = bbox[:, :, [0, 1, 2, 1, 2, 3, 0, 3]]733 scaled_bbox = bbox * self.config.bbox_scale734 bbox_position_embeddings = self.bbox_embeddings(scaled_bbox)735 736 encoder_outputs = self.encoder(737 embedding_output,738 bbox_pos_emb=bbox_position_embeddings,739 attention_mask=extended_attention_mask,740 head_mask=head_mask,741 encoder_hidden_states=encoder_hidden_states,742 encoder_attention_mask=encoder_extended_attention_mask,743 output_attentions=output_attentions,744 output_hidden_states=output_hidden_states,745 return_dict=True,746 )747 sequence_output = encoder_outputs[0]748 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None749 750 return BaseModelOutputWithPoolingAndCrossAttentions(751 last_hidden_state=sequence_output,752 pooler_output=pooled_output,753 hidden_states=encoder_outputs.hidden_states,754 attentions=encoder_outputs.attentions,755 cross_attentions=encoder_outputs.cross_attentions,756 )757 758 759@auto_docstring760class BrosForTokenClassification(BrosPreTrainedModel):761 _keys_to_ignore_on_load_unexpected = [r"pooler"]762 763 def __init__(self, config):764 super().__init__(config)765 self.num_labels = config.num_labels766 767 self.bros = BrosModel(config)768 classifier_dropout = (769 config.classifier_dropout if hasattr(config, "classifier_dropout") else config.hidden_dropout_prob770 )771 self.dropout = nn.Dropout(classifier_dropout)772 self.classifier = nn.Linear(config.hidden_size, config.num_labels)773 774 self.init_weights()775 776 @can_return_tuple777 @auto_docstring778 def forward(779 self,780 input_ids: Optional[torch.Tensor] = None,781 bbox: Optional[torch.Tensor] = None,782 attention_mask: Optional[torch.Tensor] = None,783 bbox_first_token_mask: Optional[torch.Tensor] = None,784 token_type_ids: Optional[torch.Tensor] = None,785 position_ids: Optional[torch.Tensor] = None,786 head_mask: Optional[torch.Tensor] = None,787 inputs_embeds: Optional[torch.Tensor] = None,788 labels: Optional[torch.Tensor] = None,789 output_attentions: Optional[bool] = None,790 output_hidden_states: Optional[bool] = None,791 return_dict: Optional[bool] = None,792 ) -> Union[tuple[torch.Tensor], TokenClassifierOutput]:793 r"""794 bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'):795 Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values796 (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the797 bounding box.798 bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):799 Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`:800 801 - 1 for tokens that are **not masked**,802 - 0 for tokens that are **masked**.803 804 Examples:805 806 ```python807 >>> import torch808 >>> from transformers import BrosProcessor, BrosForTokenClassification809 810 >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased")811 812 >>> model = BrosForTokenClassification.from_pretrained("jinho8345/bros-base-uncased")813 814 >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")815 >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1)816 >>> encoding["bbox"] = bbox817 818 >>> outputs = model(**encoding)819 ```"""820 821 return_dict = return_dict if return_dict is not None else self.config.use_return_dict822 823 outputs = self.bros(824 input_ids,825 bbox=bbox,826 attention_mask=attention_mask,827 token_type_ids=token_type_ids,828 position_ids=position_ids,829 head_mask=head_mask,830 inputs_embeds=inputs_embeds,831 output_attentions=output_attentions,832 output_hidden_states=output_hidden_states,833 return_dict=True,834 )835 836 sequence_output = outputs[0]837 838 sequence_output = self.dropout(sequence_output)839 logits = self.classifier(sequence_output)840 841 loss = None842 if labels is not None:843 loss_fct = CrossEntropyLoss()844 if bbox_first_token_mask is not None:845 bbox_first_token_mask = bbox_first_token_mask.view(-1)846 loss = loss_fct(847 logits.view(-1, self.num_labels)[bbox_first_token_mask], labels.view(-1)[bbox_first_token_mask]848 )849 else:850 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))851 852 return TokenClassifierOutput(853 loss=loss,854 logits=logits,855 hidden_states=outputs.hidden_states,856 attentions=outputs.attentions,857 )858 859 860@auto_docstring(861 custom_intro="""862 Bros Model with a token classification head on top (initial_token_layers and subsequent_token_layer on top of the863 hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. The initial_token_classifier is used to864 predict the first token of each entity, and the subsequent_token_classifier is used to predict the subsequent865 tokens within an entity. Compared to BrosForTokenClassification, this model is more robust to serialization errors866 since it predicts next token from one token.867 """868)869class BrosSpadeEEForTokenClassification(BrosPreTrainedModel):870 _keys_to_ignore_on_load_unexpected = [r"pooler"]871 872 def __init__(self, config):873 super().__init__(config)874 self.config = config875 self.num_labels = config.num_labels876 self.n_relations = config.n_relations877 self.backbone_hidden_size = config.hidden_size878 879 self.bros = BrosModel(config)880 classifier_dropout = (881 config.classifier_dropout if hasattr(config, "classifier_dropout") else config.hidden_dropout_prob882 )883 884 # Initial token classification for Entity Extraction (NER)885 self.initial_token_classifier = nn.Sequential(886 nn.Dropout(classifier_dropout),887 nn.Linear(config.hidden_size, config.hidden_size),888 nn.Dropout(classifier_dropout),889 nn.Linear(config.hidden_size, config.num_labels),890 )891 892 # Subsequent token classification for Entity Extraction (NER)893 self.subsequent_token_classifier = BrosRelationExtractor(config)894 895 self.init_weights()896 897 @can_return_tuple898 @auto_docstring899 def forward(900 self,901 input_ids: Optional[torch.Tensor] = None,902 bbox: Optional[torch.Tensor] = None,903 attention_mask: Optional[torch.Tensor] = None,904 bbox_first_token_mask: Optional[torch.Tensor] = None,905 token_type_ids: Optional[torch.Tensor] = None,906 position_ids: Optional[torch.Tensor] = None,907 head_mask: Optional[torch.Tensor] = None,908 inputs_embeds: Optional[torch.Tensor] = None,909 initial_token_labels: Optional[torch.Tensor] = None,910 subsequent_token_labels: Optional[torch.Tensor] = None,911 output_attentions: Optional[bool] = None,912 output_hidden_states: Optional[bool] = None,913 return_dict: Optional[bool] = None,914 ) -> Union[tuple[torch.Tensor], BrosSpadeOutput]:915 r"""916 bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'):917 Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values918 (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the919 bounding box.920 bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):921 Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`:922 923 - 1 for tokens that are **not masked**,924 - 0 for tokens that are **masked**.925 initial_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):926 Labels for the initial token classification.927 subsequent_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):928 Labels for the subsequent token classification.929 930 Examples:931 932 ```python933 >>> import torch934 >>> from transformers import BrosProcessor, BrosSpadeEEForTokenClassification935 936 >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased")937 938 >>> model = BrosSpadeEEForTokenClassification.from_pretrained("jinho8345/bros-base-uncased")939 940 >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")941 >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1)942 >>> encoding["bbox"] = bbox943 944 >>> outputs = model(**encoding)945 ```"""946 947 return_dict = return_dict if return_dict is not None else self.config.use_return_dict948 949 outputs = self.bros(950 input_ids=input_ids,951 bbox=bbox,952 attention_mask=attention_mask,953 token_type_ids=token_type_ids,954 position_ids=position_ids,955 head_mask=head_mask,956 inputs_embeds=inputs_embeds,957 output_attentions=output_attentions,958 output_hidden_states=output_hidden_states,959 return_dict=True,960 )961 962 last_hidden_states = outputs[0]963 last_hidden_states = last_hidden_states.transpose(0, 1).contiguous()964 initial_token_logits = self.initial_token_classifier(last_hidden_states).transpose(0, 1).contiguous()965 subsequent_token_logits = self.subsequent_token_classifier(last_hidden_states, last_hidden_states).squeeze(0)966 967 # make subsequent token (sequence token classification) mask968 inv_attention_mask = 1 - attention_mask969 batch_size, max_seq_length = inv_attention_mask.shape970 device = inv_attention_mask.device971 invalid_token_mask = torch.cat([inv_attention_mask, torch.zeros([batch_size, 1]).to(device)], axis=1).bool()972 subsequent_token_logits = subsequent_token_logits.masked_fill(973 invalid_token_mask[:, None, :], torch.finfo(subsequent_token_logits.dtype).min974 )975 self_token_mask = torch.eye(max_seq_length, max_seq_length + 1).to(device=device, dtype=torch.bool)976 subsequent_token_logits = subsequent_token_logits.masked_fill(977 self_token_mask[None, :, :], torch.finfo(subsequent_token_logits.dtype).min978 )979 subsequent_token_mask = attention_mask.view(-1).bool()980 981 loss = None982 if initial_token_labels is not None and subsequent_token_labels is not None:983 loss_fct = CrossEntropyLoss()984 985 # get initial token loss986 initial_token_labels = initial_token_labels.view(-1)987 if bbox_first_token_mask is not None:988 bbox_first_token_mask = bbox_first_token_mask.view(-1)989 initial_token_loss = loss_fct(990 initial_token_logits.view(-1, self.num_labels)[bbox_first_token_mask],991 initial_token_labels[bbox_first_token_mask],992 )993 else:994 initial_token_loss = loss_fct(initial_token_logits.view(-1, self.num_labels), initial_token_labels)995 996 subsequent_token_labels = subsequent_token_labels.view(-1)997 subsequent_token_loss = loss_fct(998 subsequent_token_logits.view(-1, max_seq_length + 1)[subsequent_token_mask],999 subsequent_token_labels[subsequent_token_mask],1000 )1001 1002 loss = initial_token_loss + subsequent_token_loss1003 1004 return BrosSpadeOutput(1005 loss=loss,1006 initial_token_logits=initial_token_logits,1007 subsequent_token_logits=subsequent_token_logits,1008 hidden_states=outputs.hidden_states,1009 attentions=outputs.attentions,1010 )1011 1012 1013@auto_docstring(1014 custom_intro="""1015 Bros Model with a token classification head on top (a entity_linker layer on top of the hidden-states output) e.g.1016 for Entity-Linking. The entity_linker is used to predict intra-entity links (one entity to another entity).1017 """1018)1019class BrosSpadeELForTokenClassification(BrosPreTrainedModel):1020 _keys_to_ignore_on_load_unexpected = [r"pooler"]1021 1022 def __init__(self, config):1023 super().__init__(config)1024 self.config = config1025 self.num_labels = config.num_labels1026 self.n_relations = config.n_relations1027 self.backbone_hidden_size = config.hidden_size1028 1029 self.bros = BrosModel(config)1030 (config.classifier_dropout if hasattr(config, "classifier_dropout") else config.hidden_dropout_prob)1031 1032 self.entity_linker = BrosRelationExtractor(config)1033 1034 self.init_weights()1035 1036 @can_return_tuple1037 @auto_docstring1038 def forward(1039 self,1040 input_ids: Optional[torch.Tensor] = None,1041 bbox: Optional[torch.Tensor] = None,1042 attention_mask: Optional[torch.Tensor] = None,1043 bbox_first_token_mask: Optional[torch.Tensor] = None,1044 token_type_ids: Optional[torch.Tensor] = None,1045 position_ids: Optional[torch.Tensor] = None,1046 head_mask: Optional[torch.Tensor] = None,1047 inputs_embeds: Optional[torch.Tensor] = None,1048 labels: Optional[torch.Tensor] = None,1049 output_attentions: Optional[bool] = None,1050 output_hidden_states: Optional[bool] = None,1051 return_dict: Optional[bool] = None,1052 ) -> Union[tuple[torch.Tensor], TokenClassifierOutput]:1053 r"""1054 bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'):1055 Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values1056 (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the1057 bounding box.1058 bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):1059 Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`:1060 1061 - 1 for tokens that are **not masked**,1062 - 0 for tokens that are **masked**.1063 1064 Examples:1065 1066 ```python1067 >>> import torch1068 >>> from transformers import BrosProcessor, BrosSpadeELForTokenClassification1069 1070 >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased")1071 1072 >>> model = BrosSpadeELForTokenClassification.from_pretrained("jinho8345/bros-base-uncased")1073 1074 >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")1075 >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1)1076 >>> encoding["bbox"] = bbox1077 1078 >>> outputs = model(**encoding)1079 ```"""1080 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1081 1082 outputs = self.bros(1083 input_ids=input_ids,1084 bbox=bbox,1085 attention_mask=attention_mask,1086 token_type_ids=token_type_ids,1087 position_ids=position_ids,1088 head_mask=head_mask,1089 inputs_embeds=inputs_embeds,1090 output_attentions=output_attentions,1091 output_hidden_states=output_hidden_states,1092 return_dict=True,1093 )1094 1095 last_hidden_states = outputs[0]1096 last_hidden_states = last_hidden_states.transpose(0, 1).contiguous()1097 1098 logits = self.entity_linker(last_hidden_states, last_hidden_states).squeeze(0)1099 1100 loss = None1101 if labels is not None:1102 loss_fct = CrossEntropyLoss()1103 1104 batch_size, max_seq_length = attention_mask.shape1105 device = attention_mask.device1106 1107 self_token_mask = torch.eye(max_seq_length, max_seq_length + 1).to(device=device, dtype=torch.bool)1108 1109 mask = bbox_first_token_mask.view(-1)1110 bbox_first_token_mask = torch.cat(1111 [1112 ~bbox_first_token_mask,1113 torch.zeros([batch_size, 1], dtype=torch.bool, device=device),1114 ],1115 axis=1,1116 )1117 logits = logits.masked_fill(bbox_first_token_mask[:, None, :], torch.finfo(logits.dtype).min)1118 logits = logits.masked_fill(self_token_mask[None, :, :], torch.finfo(logits.dtype).min)1119 1120 loss = loss_fct(logits.view(-1, max_seq_length + 1)[mask], labels.view(-1)[mask])1121 1122 return TokenClassifierOutput(1123 loss=loss,1124 logits=logits,1125 hidden_states=outputs.hidden_states,1126 attentions=outputs.attentions,1127 )1128 1129 1130__all__ = [1131 "BrosPreTrainedModel",1132 "BrosModel",1133 "BrosForTokenClassification",1134 "BrosSpadeEEForTokenClassification",1135 "BrosSpadeELForTokenClassification",1136]1137