DoruC/Grounded-Segment-Anything
0
1# coding=utf-82# Copyright 2021 Tel AViv University, AllenAI and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" PyTorch Splinter model."""16 17 18import math19from dataclasses import dataclass20from typing import List, Optional, Tuple, Union21 22import torch23import torch.utils.checkpoint24from torch import nn25from torch.nn import CrossEntropyLoss26 27from ...activations import ACT2FN28from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, ModelOutput, QuestionAnsweringModelOutput29from ...modeling_utils import PreTrainedModel30from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer31from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging32from .configuration_splinter import SplinterConfig33 34 35logger = logging.get_logger(__name__)36 37_CHECKPOINT_FOR_DOC = "tau/splinter-base"38_CONFIG_FOR_DOC = "SplinterConfig"39 40SPLINTER_PRETRAINED_MODEL_ARCHIVE_LIST = [41 "tau/splinter-base",42 "tau/splinter-base-qass",43 "tau/splinter-large",44 "tau/splinter-large-qass",45 # See all Splinter models at https://huggingface.co/models?filter=splinter46]47 48 49class SplinterEmbeddings(nn.Module):50 """Construct the embeddings from word, position and token_type embeddings."""51 52 def __init__(self, config):53 super().__init__()54 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)55 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)56 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)57 58 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load59 # any TensorFlow checkpoint file60 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)61 self.dropout = nn.Dropout(config.hidden_dropout_prob)62 63 # position_ids (1, len position emb) is contiguous in memory and exported when serialized64 self.register_buffer(65 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False66 )67 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")68 69 def forward(70 self,71 input_ids: Optional[torch.LongTensor] = None,72 token_type_ids: Optional[torch.LongTensor] = None,73 position_ids: Optional[torch.LongTensor] = None,74 inputs_embeds: Optional[torch.FloatTensor] = None,75 past_key_values_length: Optional[int] = 0,76 ) -> Tuple:77 if input_ids is not None:78 input_shape = input_ids.size()79 else:80 input_shape = inputs_embeds.size()[:-1]81 82 seq_length = input_shape[1]83 84 if position_ids is None:85 position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]86 87 if token_type_ids is None:88 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)89 90 if inputs_embeds is None:91 inputs_embeds = self.word_embeddings(input_ids)92 token_type_embeddings = self.token_type_embeddings(token_type_ids)93 94 embeddings = inputs_embeds + token_type_embeddings95 if self.position_embedding_type == "absolute":96 position_embeddings = self.position_embeddings(position_ids)97 embeddings += position_embeddings98 embeddings = self.LayerNorm(embeddings)99 embeddings = self.dropout(embeddings)100 return embeddings101 102 103# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Splinter104class SplinterSelfAttention(nn.Module):105 def __init__(self, config, position_embedding_type=None):106 super().__init__()107 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):108 raise ValueError(109 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "110 f"heads ({config.num_attention_heads})"111 )112 113 self.num_attention_heads = config.num_attention_heads114 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)115 self.all_head_size = self.num_attention_heads * self.attention_head_size116 117 self.query = nn.Linear(config.hidden_size, self.all_head_size)118 self.key = nn.Linear(config.hidden_size, self.all_head_size)119 self.value = nn.Linear(config.hidden_size, self.all_head_size)120 121 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)122 self.position_embedding_type = position_embedding_type or getattr(123 config, "position_embedding_type", "absolute"124 )125 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":126 self.max_position_embeddings = config.max_position_embeddings127 self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)128 129 self.is_decoder = config.is_decoder130 131 def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:132 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)133 x = x.view(new_x_shape)134 return x.permute(0, 2, 1, 3)135 136 def forward(137 self,138 hidden_states: torch.Tensor,139 attention_mask: Optional[torch.FloatTensor] = None,140 head_mask: Optional[torch.FloatTensor] = None,141 encoder_hidden_states: Optional[torch.FloatTensor] = None,142 encoder_attention_mask: Optional[torch.FloatTensor] = None,143 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,144 output_attentions: Optional[bool] = False,145 ) -> Tuple[torch.Tensor]:146 mixed_query_layer = self.query(hidden_states)147 148 # If this is instantiated as a cross-attention module, the keys149 # and values come from an encoder; the attention mask needs to be150 # such that the encoder's padding tokens are not attended to.151 is_cross_attention = encoder_hidden_states is not None152 153 if is_cross_attention and past_key_value is not None:154 # reuse k,v, cross_attentions155 key_layer = past_key_value[0]156 value_layer = past_key_value[1]157 attention_mask = encoder_attention_mask158 elif is_cross_attention:159 key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))160 value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))161 attention_mask = encoder_attention_mask162 elif past_key_value is not None:163 key_layer = self.transpose_for_scores(self.key(hidden_states))164 value_layer = self.transpose_for_scores(self.value(hidden_states))165 key_layer = torch.cat([past_key_value[0], key_layer], dim=2)166 value_layer = torch.cat([past_key_value[1], value_layer], dim=2)167 else:168 key_layer = self.transpose_for_scores(self.key(hidden_states))169 value_layer = self.transpose_for_scores(self.value(hidden_states))170 171 query_layer = self.transpose_for_scores(mixed_query_layer)172 173 use_cache = past_key_value is not None174 if self.is_decoder:175 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.176 # Further calls to cross_attention layer can then reuse all cross-attention177 # key/value_states (first "if" case)178 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of179 # all previous decoder key/value_states. Further calls to uni-directional self-attention180 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)181 # if encoder bi-directional self-attention `past_key_value` is always `None`182 past_key_value = (key_layer, value_layer)183 184 # Take the dot product between "query" and "key" to get the raw attention scores.185 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))186 187 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":188 query_length, key_length = query_layer.shape[2], key_layer.shape[2]189 if use_cache:190 position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(191 -1, 1192 )193 else:194 position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)195 position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)196 distance = position_ids_l - position_ids_r197 198 positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)199 positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility200 201 if self.position_embedding_type == "relative_key":202 relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)203 attention_scores = attention_scores + relative_position_scores204 elif self.position_embedding_type == "relative_key_query":205 relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)206 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)207 attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key208 209 attention_scores = attention_scores / math.sqrt(self.attention_head_size)210 if attention_mask is not None:211 # Apply the attention mask is (precomputed for all layers in SplinterModel forward() function)212 attention_scores = attention_scores + attention_mask213 214 # Normalize the attention scores to probabilities.215 attention_probs = nn.functional.softmax(attention_scores, dim=-1)216 217 # This is actually dropping out entire tokens to attend to, which might218 # seem a bit unusual, but is taken from the original Transformer paper.219 attention_probs = self.dropout(attention_probs)220 221 # Mask heads if we want to222 if head_mask is not None:223 attention_probs = attention_probs * head_mask224 225 context_layer = torch.matmul(attention_probs, value_layer)226 227 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()228 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)229 context_layer = context_layer.view(new_context_layer_shape)230 231 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)232 233 if self.is_decoder:234 outputs = outputs + (past_key_value,)235 return outputs236 237 238# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->Splinter239class SplinterSelfOutput(nn.Module):240 def __init__(self, config):241 super().__init__()242 self.dense = nn.Linear(config.hidden_size, config.hidden_size)243 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)244 self.dropout = nn.Dropout(config.hidden_dropout_prob)245 246 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:247 hidden_states = self.dense(hidden_states)248 hidden_states = self.dropout(hidden_states)249 hidden_states = self.LayerNorm(hidden_states + input_tensor)250 return hidden_states251 252 253# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Splinter254class SplinterAttention(nn.Module):255 def __init__(self, config, position_embedding_type=None):256 super().__init__()257 self.self = SplinterSelfAttention(config, position_embedding_type=position_embedding_type)258 self.output = SplinterSelfOutput(config)259 self.pruned_heads = set()260 261 def prune_heads(self, heads):262 if len(heads) == 0:263 return264 heads, index = find_pruneable_heads_and_indices(265 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads266 )267 268 # Prune linear layers269 self.self.query = prune_linear_layer(self.self.query, index)270 self.self.key = prune_linear_layer(self.self.key, index)271 self.self.value = prune_linear_layer(self.self.value, index)272 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)273 274 # Update hyper params and store pruned heads275 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)276 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads277 self.pruned_heads = self.pruned_heads.union(heads)278 279 def forward(280 self,281 hidden_states: torch.Tensor,282 attention_mask: Optional[torch.FloatTensor] = None,283 head_mask: Optional[torch.FloatTensor] = None,284 encoder_hidden_states: Optional[torch.FloatTensor] = None,285 encoder_attention_mask: Optional[torch.FloatTensor] = None,286 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,287 output_attentions: Optional[bool] = False,288 ) -> Tuple[torch.Tensor]:289 self_outputs = self.self(290 hidden_states,291 attention_mask,292 head_mask,293 encoder_hidden_states,294 encoder_attention_mask,295 past_key_value,296 output_attentions,297 )298 attention_output = self.output(self_outputs[0], hidden_states)299 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them300 return outputs301 302 303# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Splinter304class SplinterIntermediate(nn.Module):305 def __init__(self, config):306 super().__init__()307 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)308 if isinstance(config.hidden_act, str):309 self.intermediate_act_fn = ACT2FN[config.hidden_act]310 else:311 self.intermediate_act_fn = config.hidden_act312 313 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:314 hidden_states = self.dense(hidden_states)315 hidden_states = self.intermediate_act_fn(hidden_states)316 return hidden_states317 318 319# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->Splinter320class SplinterOutput(nn.Module):321 def __init__(self, config):322 super().__init__()323 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)324 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)325 self.dropout = nn.Dropout(config.hidden_dropout_prob)326 327 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:328 hidden_states = self.dense(hidden_states)329 hidden_states = self.dropout(hidden_states)330 hidden_states = self.LayerNorm(hidden_states + input_tensor)331 return hidden_states332 333 334# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Splinter335class SplinterLayer(nn.Module):336 def __init__(self, config):337 super().__init__()338 self.chunk_size_feed_forward = config.chunk_size_feed_forward339 self.seq_len_dim = 1340 self.attention = SplinterAttention(config)341 self.is_decoder = config.is_decoder342 self.add_cross_attention = config.add_cross_attention343 if self.add_cross_attention:344 if not self.is_decoder:345 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")346 self.crossattention = SplinterAttention(config, position_embedding_type="absolute")347 self.intermediate = SplinterIntermediate(config)348 self.output = SplinterOutput(config)349 350 def forward(351 self,352 hidden_states: torch.Tensor,353 attention_mask: Optional[torch.FloatTensor] = None,354 head_mask: Optional[torch.FloatTensor] = None,355 encoder_hidden_states: Optional[torch.FloatTensor] = None,356 encoder_attention_mask: Optional[torch.FloatTensor] = None,357 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,358 output_attentions: Optional[bool] = False,359 ) -> Tuple[torch.Tensor]:360 # decoder uni-directional self-attention cached key/values tuple is at positions 1,2361 self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None362 self_attention_outputs = self.attention(363 hidden_states,364 attention_mask,365 head_mask,366 output_attentions=output_attentions,367 past_key_value=self_attn_past_key_value,368 )369 attention_output = self_attention_outputs[0]370 371 # if decoder, the last output is tuple of self-attn cache372 if self.is_decoder:373 outputs = self_attention_outputs[1:-1]374 present_key_value = self_attention_outputs[-1]375 else:376 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights377 378 cross_attn_present_key_value = None379 if self.is_decoder and encoder_hidden_states is not None:380 if not hasattr(self, "crossattention"):381 raise ValueError(382 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"383 " by setting `config.add_cross_attention=True`"384 )385 386 # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple387 cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None388 cross_attention_outputs = self.crossattention(389 attention_output,390 attention_mask,391 head_mask,392 encoder_hidden_states,393 encoder_attention_mask,394 cross_attn_past_key_value,395 output_attentions,396 )397 attention_output = cross_attention_outputs[0]398 outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights399 400 # add cross-attn cache to positions 3,4 of present_key_value tuple401 cross_attn_present_key_value = cross_attention_outputs[-1]402 present_key_value = present_key_value + cross_attn_present_key_value403 404 layer_output = apply_chunking_to_forward(405 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output406 )407 outputs = (layer_output,) + outputs408 409 # if decoder, return the attn key/values as the last output410 if self.is_decoder:411 outputs = outputs + (present_key_value,)412 413 return outputs414 415 def feed_forward_chunk(self, attention_output):416 intermediate_output = self.intermediate(attention_output)417 layer_output = self.output(intermediate_output, attention_output)418 return layer_output419 420 421# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Splinter422class SplinterEncoder(nn.Module):423 def __init__(self, config):424 super().__init__()425 self.config = config426 self.layer = nn.ModuleList([SplinterLayer(config) for _ in range(config.num_hidden_layers)])427 self.gradient_checkpointing = False428 429 def forward(430 self,431 hidden_states: torch.Tensor,432 attention_mask: Optional[torch.FloatTensor] = None,433 head_mask: Optional[torch.FloatTensor] = None,434 encoder_hidden_states: Optional[torch.FloatTensor] = None,435 encoder_attention_mask: Optional[torch.FloatTensor] = None,436 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,437 use_cache: Optional[bool] = None,438 output_attentions: Optional[bool] = False,439 output_hidden_states: Optional[bool] = False,440 return_dict: Optional[bool] = True,441 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:442 all_hidden_states = () if output_hidden_states else None443 all_self_attentions = () if output_attentions else None444 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None445 446 if self.gradient_checkpointing and self.training:447 if use_cache:448 logger.warning_once(449 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."450 )451 use_cache = False452 453 next_decoder_cache = () if use_cache else None454 for i, layer_module in enumerate(self.layer):455 if output_hidden_states:456 all_hidden_states = all_hidden_states + (hidden_states,)457 458 layer_head_mask = head_mask[i] if head_mask is not None else None459 past_key_value = past_key_values[i] if past_key_values is not None else None460 461 if self.gradient_checkpointing and self.training:462 463 def create_custom_forward(module):464 def custom_forward(*inputs):465 return module(*inputs, past_key_value, output_attentions)466 467 return custom_forward468 469 layer_outputs = torch.utils.checkpoint.checkpoint(470 create_custom_forward(layer_module),471 hidden_states,472 attention_mask,473 layer_head_mask,474 encoder_hidden_states,475 encoder_attention_mask,476 )477 else:478 layer_outputs = layer_module(479 hidden_states,480 attention_mask,481 layer_head_mask,482 encoder_hidden_states,483 encoder_attention_mask,484 past_key_value,485 output_attentions,486 )487 488 hidden_states = layer_outputs[0]489 if use_cache:490 next_decoder_cache += (layer_outputs[-1],)491 if output_attentions:492 all_self_attentions = all_self_attentions + (layer_outputs[1],)493 if self.config.add_cross_attention:494 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)495 496 if output_hidden_states:497 all_hidden_states = all_hidden_states + (hidden_states,)498 499 if not return_dict:500 return tuple(501 v502 for v in [503 hidden_states,504 next_decoder_cache,505 all_hidden_states,506 all_self_attentions,507 all_cross_attentions,508 ]509 if v is not None510 )511 return BaseModelOutputWithPastAndCrossAttentions(512 last_hidden_state=hidden_states,513 past_key_values=next_decoder_cache,514 hidden_states=all_hidden_states,515 attentions=all_self_attentions,516 cross_attentions=all_cross_attentions,517 )518 519 520class SplinterPreTrainedModel(PreTrainedModel):521 """522 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained523 models.524 """525 526 config_class = SplinterConfig527 base_model_prefix = "splinter"528 supports_gradient_checkpointing = True529 530 # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights531 def _init_weights(self, module):532 """Initialize the weights"""533 if isinstance(module, nn.Linear):534 # Slightly different from the TF version which uses truncated_normal for initialization535 # cf https://github.com/pytorch/pytorch/pull/5617536 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)537 if module.bias is not None:538 module.bias.data.zero_()539 elif isinstance(module, nn.Embedding):540 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)541 if module.padding_idx is not None:542 module.weight.data[module.padding_idx].zero_()543 elif isinstance(module, nn.LayerNorm):544 module.bias.data.zero_()545 module.weight.data.fill_(1.0)546 547 def _set_gradient_checkpointing(self, module, value=False):548 if isinstance(module, SplinterEncoder):549 module.gradient_checkpointing = value550 551 552SPLINTER_START_DOCSTRING = r"""553 This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use554 it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and555 behavior.556 557 Parameters:558 config ([`SplinterConfig`]): Model configuration class with all the parameters of the model.559 Initializing with a config file does not load the weights associated with the model, only the560 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.561"""562 563SPLINTER_INPUTS_DOCSTRING = r"""564 Args:565 input_ids (`torch.LongTensor` of shape `({0})`):566 Indices of input sequence tokens in the vocabulary.567 568 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and569 [`PreTrainedTokenizer.__call__`] for details.570 571 [What are input IDs?](../glossary#input-ids)572 attention_mask (`torch.FloatTensor` of shape `{0}`, *optional*):573 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:574 575 - 1 for tokens that are **not masked**,576 - 0 for tokens that are **masked**.577 578 [What are attention masks?](../glossary#attention-mask)579 token_type_ids (`torch.LongTensor` of shape `{0}`, *optional*):580 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,581 1]`:582 583 - 0 corresponds to a *sentence A* token,584 - 1 corresponds to a *sentence B* token.585 586 [What are token type IDs?](../glossary#token-type-ids)587 position_ids (`torch.LongTensor` of shape `{0}`, *optional*):588 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,589 config.max_position_embeddings - 1]`.590 591 [What are position IDs?](../glossary#position-ids)592 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):593 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:594 595 - 1 indicates the head is **not masked**,596 - 0 indicates the head is **masked**.597 598 inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):599 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This600 is useful if you want more control over how to convert *input_ids* indices into associated vectors than the601 model's internal embedding lookup matrix.602 output_attentions (`bool`, *optional*):603 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned604 tensors for more detail.605 output_hidden_states (`bool`, *optional*):606 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for607 more detail.608 return_dict (`bool`, *optional*):609 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.610"""611 612 613@add_start_docstrings(614 "The bare Splinter Model transformer outputting raw hidden-states without any specific head on top.",615 SPLINTER_START_DOCSTRING,616)617class SplinterModel(SplinterPreTrainedModel):618 """619 The model is an encoder (with only self-attention) following the architecture described in [Attention is all you620 need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,621 Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.622 """623 624 def __init__(self, config):625 super().__init__(config)626 self.config = config627 628 self.embeddings = SplinterEmbeddings(config)629 self.encoder = SplinterEncoder(config)630 631 # Initialize weights and apply final processing632 self.post_init()633 634 def get_input_embeddings(self):635 return self.embeddings.word_embeddings636 637 def set_input_embeddings(self, value):638 self.embeddings.word_embeddings = value639 640 def _prune_heads(self, heads_to_prune):641 """642 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base643 class PreTrainedModel644 """645 for layer, heads in heads_to_prune.items():646 self.encoder.layer[layer].attention.prune_heads(heads)647 648 @add_start_docstrings_to_model_forward(SPLINTER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))649 @add_code_sample_docstrings(650 checkpoint=_CHECKPOINT_FOR_DOC,651 output_type=BaseModelOutputWithPastAndCrossAttentions,652 config_class=_CONFIG_FOR_DOC,653 )654 def forward(655 self,656 input_ids: Optional[torch.Tensor] = None,657 attention_mask: Optional[torch.Tensor] = None,658 token_type_ids: Optional[torch.Tensor] = None,659 position_ids: Optional[torch.Tensor] = None,660 head_mask: Optional[torch.Tensor] = None,661 inputs_embeds: Optional[torch.Tensor] = None,662 encoder_hidden_states: Optional[torch.Tensor] = None,663 encoder_attention_mask: Optional[torch.Tensor] = None,664 past_key_values: Optional[List[torch.FloatTensor]] = None,665 use_cache: Optional[bool] = None,666 output_attentions: Optional[bool] = None,667 output_hidden_states: Optional[bool] = None,668 return_dict: Optional[bool] = None,669 ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:670 r"""671 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):672 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if673 the model is configured as a decoder.674 encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):675 Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in676 the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:677 678 - 1 for tokens that are **not masked**,679 - 0 for tokens that are **masked**.680 past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):681 Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.682 If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that683 don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all684 `decoder_input_ids` of shape `(batch_size, sequence_length)`.685 use_cache (`bool`, *optional*):686 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see687 `past_key_values`).688 """689 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions690 output_hidden_states = (691 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states692 )693 return_dict = return_dict if return_dict is not None else self.config.use_return_dict694 695 if self.config.is_decoder:696 use_cache = use_cache if use_cache is not None else self.config.use_cache697 else:698 use_cache = False699 700 if input_ids is not None and inputs_embeds is not None:701 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")702 elif input_ids is not None:703 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)704 input_shape = input_ids.size()705 elif inputs_embeds is not None:706 input_shape = inputs_embeds.size()[:-1]707 else:708 raise ValueError("You have to specify either input_ids or inputs_embeds")709 710 batch_size, seq_length = input_shape711 device = input_ids.device if input_ids is not None else inputs_embeds.device712 713 # past_key_values_length714 past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0715 716 if attention_mask is None:717 attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)718 if token_type_ids is None:719 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)720 721 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]722 # ourselves in which case we just need to make it broadcastable to all heads.723 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)724 725 # If a 2D or 3D attention mask is provided for the cross-attention726 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]727 if self.config.is_decoder and encoder_hidden_states is not None:728 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()729 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)730 if encoder_attention_mask is None:731 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)732 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)733 else:734 encoder_extended_attention_mask = None735 736 # Prepare head mask if needed737 # 1.0 in head_mask indicate we keep the head738 # attention_probs has shape bsz x n_heads x N x N739 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]740 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]741 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)742 743 embedding_output = self.embeddings(744 input_ids=input_ids,745 position_ids=position_ids,746 token_type_ids=token_type_ids,747 inputs_embeds=inputs_embeds,748 past_key_values_length=past_key_values_length,749 )750 encoder_outputs = self.encoder(751 embedding_output,752 attention_mask=extended_attention_mask,753 head_mask=head_mask,754 encoder_hidden_states=encoder_hidden_states,755 encoder_attention_mask=encoder_extended_attention_mask,756 past_key_values=past_key_values,757 use_cache=use_cache,758 output_attentions=output_attentions,759 output_hidden_states=output_hidden_states,760 return_dict=return_dict,761 )762 sequence_output = encoder_outputs[0]763 764 if not return_dict:765 return (sequence_output,) + encoder_outputs[1:]766 767 return BaseModelOutputWithPastAndCrossAttentions(768 last_hidden_state=sequence_output,769 past_key_values=encoder_outputs.past_key_values,770 hidden_states=encoder_outputs.hidden_states,771 attentions=encoder_outputs.attentions,772 cross_attentions=encoder_outputs.cross_attentions,773 )774 775 776class SplinterFullyConnectedLayer(nn.Module):777 def __init__(self, input_dim, output_dim, hidden_act="gelu"):778 super().__init__()779 780 self.input_dim = input_dim781 self.output_dim = output_dim782 783 self.dense = nn.Linear(self.input_dim, self.output_dim)784 self.act_fn = ACT2FN[hidden_act]785 self.LayerNorm = nn.LayerNorm(self.output_dim)786 787 def forward(self, inputs: torch.Tensor) -> torch.Tensor:788 hidden_states = self.dense(inputs)789 hidden_states = self.act_fn(hidden_states)790 hidden_states = self.LayerNorm(hidden_states)791 return hidden_states792 793 794class QuestionAwareSpanSelectionHead(nn.Module):795 """796 Implementation of Question-Aware Span Selection (QASS) head, described in Splinter's paper:797 798 """799 800 def __init__(self, config):801 super().__init__()802 803 self.query_start_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size)804 self.query_end_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size)805 self.start_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size)806 self.end_transform = SplinterFullyConnectedLayer(config.hidden_size, config.hidden_size)807 808 self.start_classifier = nn.Linear(config.hidden_size, config.hidden_size, bias=False)809 self.end_classifier = nn.Linear(config.hidden_size, config.hidden_size, bias=False)810 811 def forward(self, inputs, positions):812 _, _, dim = inputs.size()813 index = positions.unsqueeze(-1).repeat(1, 1, dim) # [batch_size, num_positions, dim]814 gathered_reps = torch.gather(inputs, dim=1, index=index) # [batch_size, num_positions, dim]815 816 query_start_reps = self.query_start_transform(gathered_reps) # [batch_size, num_positions, dim]817 query_end_reps = self.query_end_transform(gathered_reps) # [batch_size, num_positions, dim]818 start_reps = self.start_transform(inputs) # [batch_size, seq_length, dim]819 end_reps = self.end_transform(inputs) # [batch_size, seq_length, dim]820 821 hidden_states = self.start_classifier(query_start_reps) # [batch_size, num_positions, dim]822 start_reps = start_reps.permute(0, 2, 1) # [batch_size, dim, seq_length]823 start_logits = torch.matmul(hidden_states, start_reps)824 825 hidden_states = self.end_classifier(query_end_reps)826 end_reps = end_reps.permute(0, 2, 1)827 end_logits = torch.matmul(hidden_states, end_reps)828 829 return start_logits, end_logits830 831 832@add_start_docstrings(833 """834 Splinter Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear835 layers on top of the hidden-states output to compute `span start logits` and `span end logits`).836 """,837 SPLINTER_START_DOCSTRING,838)839class SplinterForQuestionAnswering(SplinterPreTrainedModel):840 def __init__(self, config):841 super().__init__(config)842 843 self.splinter = SplinterModel(config)844 self.splinter_qass = QuestionAwareSpanSelectionHead(config)845 self.question_token_id = config.question_token_id846 847 # Initialize weights and apply final processing848 self.post_init()849 850 @add_start_docstrings_to_model_forward(SPLINTER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))851 @add_code_sample_docstrings(852 checkpoint=_CHECKPOINT_FOR_DOC,853 output_type=QuestionAnsweringModelOutput,854 config_class=_CONFIG_FOR_DOC,855 )856 def forward(857 self,858 input_ids: Optional[torch.Tensor] = None,859 attention_mask: Optional[torch.Tensor] = None,860 token_type_ids: Optional[torch.Tensor] = None,861 position_ids: Optional[torch.Tensor] = None,862 head_mask: Optional[torch.Tensor] = None,863 inputs_embeds: Optional[torch.Tensor] = None,864 start_positions: Optional[torch.LongTensor] = None,865 end_positions: Optional[torch.LongTensor] = None,866 output_attentions: Optional[bool] = None,867 output_hidden_states: Optional[bool] = None,868 return_dict: Optional[bool] = None,869 question_positions: Optional[torch.LongTensor] = None,870 ) -> Union[Tuple, QuestionAnsweringModelOutput]:871 r"""872 start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):873 Labels for position (index) of the start of the labelled span for computing the token classification loss.874 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence875 are not taken into account for computing the loss.876 end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):877 Labels for position (index) of the end of the labelled span for computing the token classification loss.878 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence879 are not taken into account for computing the loss.880 question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):881 The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size,882 num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be883 the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size,884 sequence_length)`.885 """886 return_dict = return_dict if return_dict is not None else self.config.use_return_dict887 888 question_positions_were_none = False889 if question_positions is None:890 if input_ids is not None:891 question_position_for_each_example = torch.argmax(892 (torch.eq(input_ids, self.question_token_id)).int(), dim=-1893 )894 else:895 question_position_for_each_example = torch.zeros(896 inputs_embeds.size(0), dtype=torch.long, layout=inputs_embeds.layout, device=inputs_embeds.device897 )898 question_positions = question_position_for_each_example.unsqueeze(-1)899 question_positions_were_none = True900 901 outputs = self.splinter(902 input_ids,903 attention_mask=attention_mask,904 token_type_ids=token_type_ids,905 position_ids=position_ids,906 head_mask=head_mask,907 inputs_embeds=inputs_embeds,908 output_attentions=output_attentions,909 output_hidden_states=output_hidden_states,910 return_dict=return_dict,911 )912 913 sequence_output = outputs[0]914 start_logits, end_logits = self.splinter_qass(sequence_output, question_positions)915 916 if question_positions_were_none:917 start_logits, end_logits = start_logits.squeeze(1), end_logits.squeeze(1)918 919 if attention_mask is not None:920 start_logits = start_logits + (1 - attention_mask) * torch.finfo(start_logits.dtype).min921 end_logits = end_logits + (1 - attention_mask) * torch.finfo(end_logits.dtype).min922 923 total_loss = None924 if start_positions is not None and end_positions is not None:925 # If we are on multi-GPU, split add a dimension926 if len(start_positions.size()) > 1:927 start_positions = start_positions.squeeze(-1)928 if len(end_positions.size()) > 1:929 end_positions = end_positions.squeeze(-1)930 # sometimes the start/end positions are outside our model inputs, we ignore these terms931 ignored_index = start_logits.size(1)932 start_positions.clamp_(0, ignored_index)933 end_positions.clamp_(0, ignored_index)934 935 loss_fct = CrossEntropyLoss(ignore_index=ignored_index)936 start_loss = loss_fct(start_logits, start_positions)937 end_loss = loss_fct(end_logits, end_positions)938 total_loss = (start_loss + end_loss) / 2939 940 if not return_dict:941 output = (start_logits, end_logits) + outputs[1:]942 return ((total_loss,) + output) if total_loss is not None else output943 944 return QuestionAnsweringModelOutput(945 loss=total_loss,946 start_logits=start_logits,947 end_logits=end_logits,948 hidden_states=outputs.hidden_states,949 attentions=outputs.attentions,950 )951 952 953@dataclass954class SplinterForPreTrainingOutput(ModelOutput):955 """956 Class for outputs of Splinter as a span selection model.957 958 Args:959 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when start and end positions are provided):960 Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.961 start_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`):962 Span-start scores (before SoftMax).963 end_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`):964 Span-end scores (before SoftMax).965 hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):966 Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +967 one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.968 969 Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.970 attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):971 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,972 sequence_length)`.973 974 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention975 heads.976 """977 978 loss: Optional[torch.FloatTensor] = None979 start_logits: torch.FloatTensor = None980 end_logits: torch.FloatTensor = None981 hidden_states: Optional[Tuple[torch.FloatTensor]] = None982 attentions: Optional[Tuple[torch.FloatTensor]] = None983 984 985@add_start_docstrings(986 """987 Splinter Model for the recurring span selection task as done during the pretraining. The difference to the QA task988 is that we do not have a question, but multiple question tokens that replace the occurrences of recurring spans989 instead.990 """,991 SPLINTER_START_DOCSTRING,992)993class SplinterForPreTraining(SplinterPreTrainedModel):994 def __init__(self, config):995 super().__init__(config)996 997 self.splinter = SplinterModel(config)998 self.splinter_qass = QuestionAwareSpanSelectionHead(config)999 self.question_token_id = config.question_token_id1000 1001 # Initialize weights and apply final processing1002 self.post_init()1003 1004 @add_start_docstrings_to_model_forward(1005 SPLINTER_INPUTS_DOCSTRING.format("batch_size, num_questions, sequence_length")1006 )1007 def forward(1008 self,1009 input_ids: Optional[torch.Tensor] = None,1010 attention_mask: Optional[torch.Tensor] = None,1011 token_type_ids: Optional[torch.Tensor] = None,1012 position_ids: Optional[torch.Tensor] = None,1013 head_mask: Optional[torch.Tensor] = None,1014 inputs_embeds: Optional[torch.Tensor] = None,1015 start_positions: Optional[torch.LongTensor] = None,1016 end_positions: Optional[torch.LongTensor] = None,1017 output_attentions: Optional[bool] = None,1018 output_hidden_states: Optional[bool] = None,1019 return_dict: Optional[bool] = None,1020 question_positions: Optional[torch.LongTensor] = None,1021 ) -> Union[Tuple, SplinterForPreTrainingOutput]:1022 r"""1023 start_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):1024 Labels for position (index) of the start of the labelled span for computing the token classification loss.1025 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1026 are not taken into account for computing the loss.1027 end_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):1028 Labels for position (index) of the end of the labelled span for computing the token classification loss.1029 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1030 are not taken into account for computing the loss.1031 question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):1032 The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size,1033 num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be1034 the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size,1035 sequence_length)`.1036 """1037 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1038 1039 if question_positions is None and start_positions is not None and end_positions is not None:1040 raise TypeError("question_positions must be specified in order to calculate the loss")1041 1042 elif question_positions is None and input_ids is None:1043 raise TypeError("question_positions must be specified when input_embeds is used")1044 1045 elif question_positions is None:1046 question_positions = self._prepare_question_positions(input_ids)1047 1048 outputs = self.splinter(1049 input_ids,1050 attention_mask=attention_mask,1051 token_type_ids=token_type_ids,1052 position_ids=position_ids,1053 head_mask=head_mask,1054 inputs_embeds=inputs_embeds,1055 output_attentions=output_attentions,1056 output_hidden_states=output_hidden_states,1057 return_dict=return_dict,1058 )1059 1060 sequence_output = outputs[0]1061 batch_size, sequence_length, dim = sequence_output.size()1062 # [batch_size, num_questions, sequence_length]1063 start_logits, end_logits = self.splinter_qass(sequence_output, question_positions)1064 1065 num_questions = question_positions.size(1)1066 if attention_mask is not None:1067 attention_mask_for_each_question = attention_mask.unsqueeze(1).expand(1068 batch_size, num_questions, sequence_length1069 )1070 start_logits = start_logits + (1 - attention_mask_for_each_question) * torch.finfo(start_logits.dtype).min1071 end_logits = end_logits + (1 - attention_mask_for_each_question) * torch.finfo(end_logits.dtype).min1072 1073 total_loss = None1074 # [batch_size, num_questions, sequence_length]1075 if start_positions is not None and end_positions is not None:1076 # sometimes the start/end positions are outside our model inputs, we ignore these terms1077 start_positions.clamp_(0, max(0, sequence_length - 1))1078 end_positions.clamp_(0, max(0, sequence_length - 1))1079 1080 # Ignore zero positions in the loss. Splinter never predicts zero1081 # during pretraining and zero is used for padding question1082 # tokens as well as for start and end positions of padded1083 # question tokens.1084 loss_fct = CrossEntropyLoss(ignore_index=self.config.pad_token_id)1085 start_loss = loss_fct(1086 start_logits.view(batch_size * num_questions, sequence_length),1087 start_positions.view(batch_size * num_questions),1088 )1089 end_loss = loss_fct(1090 end_logits.view(batch_size * num_questions, sequence_length),1091 end_positions.view(batch_size * num_questions),1092 )1093 total_loss = (start_loss + end_loss) / 21094 1095 if not return_dict:1096 output = (start_logits, end_logits) + outputs[1:]1097 return ((total_loss,) + output) if total_loss is not None else output1098 1099 return SplinterForPreTrainingOutput(1100 loss=total_loss,1101 start_logits=start_logits,1102 end_logits=end_logits,1103 hidden_states=outputs.hidden_states,1104 attentions=outputs.attentions,1105 )1106 1107 def _prepare_question_positions(self, input_ids: torch.Tensor) -> torch.Tensor:1108 rows, flat_positions = torch.where(input_ids == self.config.question_token_id)1109 num_questions = torch.bincount(rows)1110 positions = torch.full(1111 (input_ids.size(0), num_questions.max()),1112 self.config.pad_token_id,1113 dtype=torch.long,1114 device=input_ids.device,1115 )1116 cols = torch.cat([torch.arange(n) for n in num_questions])1117 positions[rows, cols] = flat_positions1118 return positions1119