bidirectional.py (38847B)
1 ### All BERT related code copied from the Huggingface Transformers library. 2 import torch 3 import math 4 from typing import List, Optional, Tuple, Union 5 import torch.utils.checkpoint 6 from torch import nn 7 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss 8 from transformers.models.bert.configuration_bert import BertConfig 9 from transformers.utils import ( 10 ModelOutput, 11 add_code_sample_docstrings, 12 add_start_docstrings, 13 add_start_docstrings_to_model_forward, 14 logging, 15 replace_return_docstrings, 16 ) 17 from transformers.modeling_utils import PreTrainedModel 18 from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer 19 from transformers.activations import ACT2FN 20 from transformers.modeling_outputs import ( 21 BaseModelOutputWithPastAndCrossAttentions, 22 BaseModelOutputWithPoolingAndCrossAttentions, 23 CausalLMOutputWithCrossAttentions, 24 MaskedLMOutput, 25 MultipleChoiceModelOutput, 26 NextSentencePredictorOutput, 27 QuestionAnsweringModelOutput, 28 SequenceClassifierOutput, 29 TokenClassifierOutput, 30 ) 31 def load_tf_weights_in_bert(model, config, tf_checkpoint_path): 32 """Load tf checkpoints in a pytorch model.""" 33 try: 34 import re 35 36 import numpy as np 37 import tensorflow as tf 38 except ImportError: 39 logger.error( 40 "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " 41 "https://www.tensorflow.org/install/ for installation instructions." 42 ) 43 raise 44 tf_path = os.path.abspath(tf_checkpoint_path) 45 logger.info(f"Converting TensorFlow checkpoint from {tf_path}") 46 # Load weights from TF model 47 init_vars = tf.train.list_variables(tf_path) 48 names = [] 49 arrays = [] 50 for name, shape in init_vars: 51 logger.info(f"Loading TF weight {name} with shape {shape}") 52 array = tf.train.load_variable(tf_path, name) 53 names.append(name) 54 arrays.append(array) 55 56 for name, array in zip(names, arrays): 57 name = name.split("/") 58 # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v 59 # which are not required for using pretrained model 60 if any( 61 n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] 62 for n in name 63 ): 64 logger.info(f"Skipping {'/'.join(name)}") 65 continue 66 pointer = model 67 for m_name in name: 68 if re.fullmatch(r"[A-Za-z]+_\d+", m_name): 69 scope_names = re.split(r"_(\d+)", m_name) 70 else: 71 scope_names = [m_name] 72 if scope_names[0] == "kernel" or scope_names[0] == "gamma": 73 pointer = getattr(pointer, "weight") 74 elif scope_names[0] == "output_bias" or scope_names[0] == "beta": 75 pointer = getattr(pointer, "bias") 76 elif scope_names[0] == "output_weights": 77 pointer = getattr(pointer, "weight") 78 elif scope_names[0] == "squad": 79 pointer = getattr(pointer, "classifier") 80 else: 81 try: 82 pointer = getattr(pointer, scope_names[0]) 83 except AttributeError: 84 logger.info(f"Skipping {'/'.join(name)}") 85 continue 86 if len(scope_names) >= 2: 87 num = int(scope_names[1]) 88 pointer = pointer[num] 89 if m_name[-11:] == "_embeddings": 90 pointer = getattr(pointer, "weight") 91 elif m_name == "kernel": 92 array = np.transpose(array) 93 try: 94 if pointer.shape != array.shape: 95 raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched") 96 except ValueError as e: 97 e.args += (pointer.shape, array.shape) 98 raise 99 logger.info(f"Initialize PyTorch weight {name}") 100 pointer.data = torch.from_numpy(array) 101 return model 102 103 class BertPreTrainedModel(PreTrainedModel): 104 """ 105 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained 106 models. 107 """ 108 109 config_class = BertConfig 110 load_tf_weights = load_tf_weights_in_bert 111 base_model_prefix = "bert" 112 supports_gradient_checkpointing = True 113 114 def _init_weights(self, module): 115 """Initialize the weights""" 116 if isinstance(module, nn.Linear): 117 # Slightly different from the TF version which uses truncated_normal for initialization 118 # cf https://github.com/pytorch/pytorch/pull/5617 119 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) 120 if module.bias is not None: 121 module.bias.data.zero_() 122 elif isinstance(module, nn.Embedding): 123 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) 124 if module.padding_idx is not None: 125 module.weight.data[module.padding_idx].zero_() 126 elif isinstance(module, nn.LayerNorm): 127 module.bias.data.zero_() 128 module.weight.data.fill_(1.0) 129 130 131 class BertModel(BertPreTrainedModel): 132 """ 133 134 The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of 135 cross-attention is added between the self-attention layers, following the architecture described in [Attention is 136 all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, 137 Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. 138 139 To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set 140 to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and 141 `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. 142 """ 143 144 def __init__(self, config, add_pooling_layer=True): 145 super().__init__(config) 146 self.config = config 147 148 self.embeddings = BertEmbeddings(config) 149 self.encoder = BertEncoder(config) 150 151 self.pooler = BertPooler(config) if add_pooling_layer else None 152 153 # Initialize weights and apply final processing 154 self.post_init() 155 156 def get_input_embeddings(self): 157 return self.embeddings.word_embeddings 158 159 def set_input_embeddings(self, value): 160 self.embeddings.word_embeddings = value 161 162 def _prune_heads(self, heads_to_prune): 163 """ 164 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base 165 class PreTrainedModel 166 """ 167 for layer, heads in heads_to_prune.items(): 168 self.encoder.layer[layer].attention.prune_heads(heads) 169 170 def forward( 171 self, 172 input_ids: Optional[torch.Tensor] = None, 173 attention_mask: Optional[torch.Tensor] = None, 174 token_type_ids: Optional[torch.Tensor] = None, 175 position_ids: Optional[torch.Tensor] = None, 176 head_mask: Optional[torch.Tensor] = None, 177 inputs_embeds: Optional[torch.Tensor] = None, 178 encoder_hidden_states: Optional[torch.Tensor] = None, 179 encoder_attention_mask: Optional[torch.Tensor] = None, 180 past_key_values: Optional[List[torch.FloatTensor]] = None, 181 use_cache: Optional[bool] = None, 182 output_attentions: Optional[bool] = None, 183 output_hidden_states: Optional[bool] = None, 184 return_dict: Optional[bool] = None, 185 special_node_embedding: Optional[torch.Tensor] = None, 186 spt: Optional[torch.Tensor] = None, 187 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: 188 r""" 189 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): 190 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if 191 the model is configured as a decoder. 192 encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): 193 Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in 194 the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: 195 196 - 1 for tokens that are **not masked**, 197 - 0 for tokens that are **masked**. 198 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)`): 199 Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. 200 201 If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that 202 don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all 203 `decoder_input_ids` of shape `(batch_size, sequence_length)`. 204 use_cache (`bool`, *optional*): 205 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see 206 `past_key_values`). 207 """ 208 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions 209 output_hidden_states = ( 210 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states 211 ) 212 return_dict = return_dict if return_dict is not None else self.config.use_return_dict 213 214 if self.config.is_decoder: 215 use_cache = use_cache if use_cache is not None else self.config.use_cache 216 else: 217 use_cache = False 218 219 if input_ids is not None and inputs_embeds is not None: 220 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") 221 elif input_ids is not None: 222 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) 223 input_shape = input_ids.size() 224 elif inputs_embeds is not None: 225 input_shape = inputs_embeds.size()[:-1] 226 else: 227 raise ValueError("You have to specify either input_ids or inputs_embeds") 228 229 batch_size, seq_length = input_shape 230 device = input_ids.device if input_ids is not None else inputs_embeds.device 231 232 # past_key_values_length 233 past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 234 235 if attention_mask is None: 236 attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) 237 238 if token_type_ids is None: 239 if hasattr(self.embeddings, "token_type_ids"): 240 buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] 241 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) 242 token_type_ids = buffered_token_type_ids_expanded 243 else: 244 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) 245 246 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] 247 # ourselves in which case we just need to make it broadcastable to all heads. 248 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) 249 250 # If a 2D or 3D attention mask is provided for the cross-attention 251 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] 252 if self.config.is_decoder and encoder_hidden_states is not None: 253 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() 254 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) 255 if encoder_attention_mask is None: 256 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) 257 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) 258 else: 259 encoder_extended_attention_mask = None 260 261 # Prepare head mask if needed 262 # 1.0 in head_mask indicate we keep the head 263 # attention_probs has shape bsz x n_heads x N x N 264 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] 265 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] 266 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) 267 268 embedding_output = self.embeddings( 269 input_ids=input_ids, 270 position_ids=position_ids, 271 token_type_ids=token_type_ids, 272 inputs_embeds=inputs_embeds, 273 past_key_values_length=past_key_values_length, 274 special_node_embedding=special_node_embedding, 275 spt=spt 276 ) 277 encoder_outputs = self.encoder( 278 embedding_output, 279 attention_mask=extended_attention_mask, 280 head_mask=head_mask, 281 encoder_hidden_states=encoder_hidden_states, 282 encoder_attention_mask=encoder_extended_attention_mask, 283 past_key_values=past_key_values, 284 use_cache=use_cache, 285 output_attentions=output_attentions, 286 output_hidden_states=output_hidden_states, 287 return_dict=return_dict, 288 ) 289 sequence_output = encoder_outputs[0] 290 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None 291 292 if not return_dict: 293 return (sequence_output, pooled_output) + encoder_outputs[1:] 294 295 return BaseModelOutputWithPoolingAndCrossAttentions( 296 last_hidden_state=sequence_output, 297 pooler_output=pooled_output, 298 past_key_values=encoder_outputs.past_key_values, 299 hidden_states=encoder_outputs.hidden_states, 300 attentions=encoder_outputs.attentions, 301 cross_attentions=encoder_outputs.cross_attentions, 302 ) 303 304 class BertEmbeddings(nn.Module): 305 """Construct the embeddings from word, position and token_type embeddings.""" 306 307 def __init__(self, config): 308 super().__init__() 309 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) 310 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) 311 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) 312 313 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load 314 # any TensorFlow checkpoint file 315 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) 316 self.dropout = nn.Dropout(config.hidden_dropout_prob) 317 # position_ids (1, len position emb) is contiguous in memory and exported when serialized 318 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute") 319 self.register_buffer( 320 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False 321 ) 322 self.register_buffer( 323 "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False 324 ) 325 326 def forward( 327 self, 328 input_ids: Optional[torch.LongTensor] = None, 329 token_type_ids: Optional[torch.LongTensor] = None, 330 position_ids: Optional[torch.LongTensor] = None, 331 inputs_embeds: Optional[torch.FloatTensor] = None, 332 past_key_values_length: int = 0, 333 special_node_embedding: Optional[torch.LongTensor] = None, 334 spt: Optional[torch.LongTensor] = None, 335 ) -> torch.Tensor: 336 if input_ids is not None: 337 input_shape = input_ids.size() 338 else: 339 input_shape = inputs_embeds.size()[:-1] 340 341 seq_length = input_shape[1] 342 343 if position_ids is None: 344 position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length] 345 346 # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs 347 # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves 348 # issue #5664 349 if token_type_ids is None: 350 if hasattr(self, "token_type_ids"): 351 buffered_token_type_ids = self.token_type_ids[:, :seq_length] 352 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length) 353 token_type_ids = buffered_token_type_ids_expanded 354 else: 355 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) 356 357 if inputs_embeds is None: 358 inputs_embeds = self.word_embeddings(input_ids) 359 token_type_embeddings = self.token_type_embeddings(token_type_ids) 360 361 embeddings = inputs_embeds + token_type_embeddings 362 if self.position_embedding_type == "absolute": 363 position_embeddings = self.position_embeddings(position_ids) 364 embeddings += position_embeddings 365 # Add special node embedding 366 #embeddings[:, -2, :] = special_node_embedding 367 embeddings[torch.arange(embeddings.size(0)), spt] = special_node_embedding 368 embeddings = self.LayerNorm(embeddings) 369 embeddings = self.dropout(embeddings) 370 return embeddings 371 372 373 class BertSelfAttention(nn.Module): 374 def __init__(self, config, position_embedding_type=None): 375 super().__init__() 376 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): 377 raise ValueError( 378 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " 379 f"heads ({config.num_attention_heads})" 380 ) 381 382 self.num_attention_heads = config.num_attention_heads 383 self.attention_head_size = int(config.hidden_size / config.num_attention_heads) 384 self.all_head_size = self.num_attention_heads * self.attention_head_size 385 386 self.query = nn.Linear(config.hidden_size, self.all_head_size) 387 self.key = nn.Linear(config.hidden_size, self.all_head_size) 388 self.value = nn.Linear(config.hidden_size, self.all_head_size) 389 390 self.dropout = nn.Dropout(config.attention_probs_dropout_prob) 391 self.position_embedding_type = position_embedding_type or getattr( 392 config, "position_embedding_type", "absolute" 393 ) 394 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query": 395 self.max_position_embeddings = config.max_position_embeddings 396 self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size) 397 398 self.is_decoder = config.is_decoder 399 400 def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: 401 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) 402 x = x.view(new_x_shape) 403 return x.permute(0, 2, 1, 3) 404 405 def forward( 406 self, 407 hidden_states: 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 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 413 output_attentions: Optional[bool] = False, 414 ) -> Tuple[torch.Tensor]: 415 mixed_query_layer = self.query(hidden_states) 416 417 # If this is instantiated as a cross-attention module, the keys 418 # and values come from an encoder; the attention mask needs to be 419 # such that the encoder's padding tokens are not attended to. 420 is_cross_attention = encoder_hidden_states is not None 421 422 if is_cross_attention and past_key_value is not None: 423 # reuse k,v, cross_attentions 424 key_layer = past_key_value[0] 425 value_layer = past_key_value[1] 426 attention_mask = encoder_attention_mask 427 elif is_cross_attention: 428 key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) 429 value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) 430 attention_mask = encoder_attention_mask 431 elif past_key_value is not None: 432 key_layer = self.transpose_for_scores(self.key(hidden_states)) 433 value_layer = self.transpose_for_scores(self.value(hidden_states)) 434 key_layer = torch.cat([past_key_value[0], key_layer], dim=2) 435 value_layer = torch.cat([past_key_value[1], value_layer], dim=2) 436 else: 437 key_layer = self.transpose_for_scores(self.key(hidden_states)) 438 value_layer = self.transpose_for_scores(self.value(hidden_states)) 439 440 query_layer = self.transpose_for_scores(mixed_query_layer) 441 442 use_cache = past_key_value is not None 443 if self.is_decoder: 444 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. 445 # Further calls to cross_attention layer can then reuse all cross-attention 446 # key/value_states (first "if" case) 447 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of 448 # all previous decoder key/value_states. Further calls to uni-directional self-attention 449 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) 450 # if encoder bi-directional self-attention `past_key_value` is always `None` 451 past_key_value = (key_layer, value_layer) 452 453 # Take the dot product between "query" and "key" to get the raw attention scores. 454 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) 455 456 if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query": 457 query_length, key_length = query_layer.shape[2], key_layer.shape[2] 458 if use_cache: 459 position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view( 460 -1, 1 461 ) 462 else: 463 position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1) 464 position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1) 465 distance = position_ids_l - position_ids_r 466 467 positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1) 468 positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility 469 470 if self.position_embedding_type == "relative_key": 471 relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) 472 attention_scores = attention_scores + relative_position_scores 473 elif self.position_embedding_type == "relative_key_query": 474 relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) 475 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding) 476 attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key 477 478 attention_scores = attention_scores / math.sqrt(self.attention_head_size) 479 if attention_mask is not None: 480 # Apply the attention mask is (precomputed for all layers in BertModel forward() function) 481 attention_scores = attention_scores + attention_mask 482 483 # Normalize the attention scores to probabilities. 484 attention_probs = nn.functional.softmax(attention_scores, dim=-1) 485 486 # This is actually dropping out entire tokens to attend to, which might 487 # seem a bit unusual, but is taken from the original Transformer paper. 488 attention_probs = self.dropout(attention_probs) 489 490 # Mask heads if we want to 491 if head_mask is not None: 492 attention_probs = attention_probs * head_mask 493 494 context_layer = torch.matmul(attention_probs, value_layer) 495 496 context_layer = context_layer.permute(0, 2, 1, 3).contiguous() 497 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) 498 context_layer = context_layer.view(new_context_layer_shape) 499 500 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) 501 502 if self.is_decoder: 503 outputs = outputs + (past_key_value,) 504 return outputs 505 506 507 class BertSelfOutput(nn.Module): 508 def __init__(self, config): 509 super().__init__() 510 self.dense = nn.Linear(config.hidden_size, config.hidden_size) 511 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) 512 self.dropout = nn.Dropout(config.hidden_dropout_prob) 513 514 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: 515 hidden_states = self.dense(hidden_states) 516 hidden_states = self.dropout(hidden_states) 517 hidden_states = self.LayerNorm(hidden_states + input_tensor) 518 return hidden_states 519 520 521 class BertAttention(nn.Module): 522 def __init__(self, config, position_embedding_type=None): 523 super().__init__() 524 self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type) 525 self.output = BertSelfOutput(config) 526 self.pruned_heads = set() 527 528 def prune_heads(self, heads): 529 if len(heads) == 0: 530 return 531 heads, index = find_pruneable_heads_and_indices( 532 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads 533 ) 534 535 # Prune linear layers 536 self.self.query = prune_linear_layer(self.self.query, index) 537 self.self.key = prune_linear_layer(self.self.key, index) 538 self.self.value = prune_linear_layer(self.self.value, index) 539 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) 540 541 # Update hyper params and store pruned heads 542 self.self.num_attention_heads = self.self.num_attention_heads - len(heads) 543 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads 544 self.pruned_heads = self.pruned_heads.union(heads) 545 546 def forward( 547 self, 548 hidden_states: torch.Tensor, 549 attention_mask: Optional[torch.FloatTensor] = None, 550 head_mask: Optional[torch.FloatTensor] = None, 551 encoder_hidden_states: Optional[torch.FloatTensor] = None, 552 encoder_attention_mask: Optional[torch.FloatTensor] = None, 553 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 554 output_attentions: Optional[bool] = False, 555 ) -> Tuple[torch.Tensor]: 556 self_outputs = self.self( 557 hidden_states, 558 attention_mask, 559 head_mask, 560 encoder_hidden_states, 561 encoder_attention_mask, 562 past_key_value, 563 output_attentions, 564 ) 565 attention_output = self.output(self_outputs[0], hidden_states) 566 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them 567 return outputs 568 569 570 class BertIntermediate(nn.Module): 571 def __init__(self, config): 572 super().__init__() 573 self.dense = nn.Linear(config.hidden_size, config.intermediate_size) 574 if isinstance(config.hidden_act, str): 575 self.intermediate_act_fn = ACT2FN[config.hidden_act] 576 else: 577 self.intermediate_act_fn = config.hidden_act 578 579 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 580 hidden_states = self.dense(hidden_states) 581 hidden_states = self.intermediate_act_fn(hidden_states) 582 return hidden_states 583 584 585 class BertOutput(nn.Module): 586 def __init__(self, config): 587 super().__init__() 588 self.dense = nn.Linear(config.intermediate_size, config.hidden_size) 589 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) 590 self.dropout = nn.Dropout(config.hidden_dropout_prob) 591 592 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: 593 hidden_states = self.dense(hidden_states) 594 hidden_states = self.dropout(hidden_states) 595 hidden_states = self.LayerNorm(hidden_states + input_tensor) 596 return hidden_states 597 598 599 class BertLayer(nn.Module): 600 def __init__(self, config): 601 super().__init__() 602 self.chunk_size_feed_forward = config.chunk_size_feed_forward 603 self.seq_len_dim = 1 604 self.attention = BertAttention(config) 605 self.is_decoder = config.is_decoder 606 self.add_cross_attention = config.add_cross_attention 607 if self.add_cross_attention: 608 if not self.is_decoder: 609 raise ValueError(f"{self} should be used as a decoder model if cross attention is added") 610 self.crossattention = BertAttention(config, position_embedding_type="absolute") 611 self.intermediate = BertIntermediate(config) 612 self.output = BertOutput(config) 613 614 def forward( 615 self, 616 hidden_states: torch.Tensor, 617 attention_mask: Optional[torch.FloatTensor] = None, 618 head_mask: Optional[torch.FloatTensor] = None, 619 encoder_hidden_states: Optional[torch.FloatTensor] = None, 620 encoder_attention_mask: Optional[torch.FloatTensor] = None, 621 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 622 output_attentions: Optional[bool] = False, 623 ) -> Tuple[torch.Tensor]: 624 # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 625 self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None 626 self_attention_outputs = self.attention( 627 hidden_states, 628 attention_mask, 629 head_mask, 630 output_attentions=output_attentions, 631 past_key_value=self_attn_past_key_value, 632 ) 633 attention_output = self_attention_outputs[0] 634 635 # if decoder, the last output is tuple of self-attn cache 636 if self.is_decoder: 637 outputs = self_attention_outputs[1:-1] 638 present_key_value = self_attention_outputs[-1] 639 else: 640 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights 641 642 cross_attn_present_key_value = None 643 if self.is_decoder and encoder_hidden_states is not None: 644 if not hasattr(self, "crossattention"): 645 raise ValueError( 646 f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers" 647 " by setting `config.add_cross_attention=True`" 648 ) 649 650 # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple 651 cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None 652 cross_attention_outputs = self.crossattention( 653 attention_output, 654 attention_mask, 655 head_mask, 656 encoder_hidden_states, 657 encoder_attention_mask, 658 cross_attn_past_key_value, 659 output_attentions, 660 ) 661 attention_output = cross_attention_outputs[0] 662 outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights 663 664 # add cross-attn cache to positions 3,4 of present_key_value tuple 665 cross_attn_present_key_value = cross_attention_outputs[-1] 666 present_key_value = present_key_value + cross_attn_present_key_value 667 668 layer_output = apply_chunking_to_forward( 669 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output 670 ) 671 outputs = (layer_output,) + outputs 672 673 # if decoder, return the attn key/values as the last output 674 if self.is_decoder: 675 outputs = outputs + (present_key_value,) 676 677 return outputs 678 679 def feed_forward_chunk(self, attention_output): 680 intermediate_output = self.intermediate(attention_output) 681 layer_output = self.output(intermediate_output, attention_output) 682 return layer_output 683 684 685 class BertEncoder(nn.Module): 686 def __init__(self, config): 687 super().__init__() 688 self.config = config 689 self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)]) 690 self.gradient_checkpointing = False 691 692 def forward( 693 self, 694 hidden_states: torch.Tensor, 695 attention_mask: Optional[torch.FloatTensor] = None, 696 head_mask: Optional[torch.FloatTensor] = None, 697 encoder_hidden_states: Optional[torch.FloatTensor] = None, 698 encoder_attention_mask: Optional[torch.FloatTensor] = None, 699 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 700 use_cache: Optional[bool] = None, 701 output_attentions: Optional[bool] = False, 702 output_hidden_states: Optional[bool] = False, 703 return_dict: Optional[bool] = True, 704 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: 705 all_hidden_states = () if output_hidden_states else None 706 all_self_attentions = () if output_attentions else None 707 all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None 708 709 if self.gradient_checkpointing and self.training: 710 if use_cache: 711 logger.warning_once( 712 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." 713 ) 714 use_cache = False 715 716 next_decoder_cache = () if use_cache else None 717 for i, layer_module in enumerate(self.layer): 718 if output_hidden_states: 719 all_hidden_states = all_hidden_states + (hidden_states,) 720 721 layer_head_mask = head_mask[i] if head_mask is not None else None 722 past_key_value = past_key_values[i] if past_key_values is not None else None 723 724 if self.gradient_checkpointing and self.training: 725 layer_outputs = self._gradient_checkpointing_func( 726 layer_module.__call__, 727 hidden_states, 728 attention_mask, 729 layer_head_mask, 730 encoder_hidden_states, 731 encoder_attention_mask, 732 past_key_value, 733 output_attentions, 734 ) 735 else: 736 layer_outputs = layer_module( 737 hidden_states, 738 attention_mask, 739 layer_head_mask, 740 encoder_hidden_states, 741 encoder_attention_mask, 742 past_key_value, 743 output_attentions, 744 ) 745 746 hidden_states = layer_outputs[0] 747 if use_cache: 748 next_decoder_cache += (layer_outputs[-1],) 749 if output_attentions: 750 all_self_attentions = all_self_attentions + (layer_outputs[1],) 751 if self.config.add_cross_attention: 752 all_cross_attentions = all_cross_attentions + (layer_outputs[2],) 753 754 if output_hidden_states: 755 all_hidden_states = all_hidden_states + (hidden_states,) 756 757 if not return_dict: 758 return tuple( 759 v 760 for v in [ 761 hidden_states, 762 next_decoder_cache, 763 all_hidden_states, 764 all_self_attentions, 765 all_cross_attentions, 766 ] 767 if v is not None 768 ) 769 return BaseModelOutputWithPastAndCrossAttentions( 770 last_hidden_state=hidden_states, 771 past_key_values=next_decoder_cache, 772 hidden_states=all_hidden_states, 773 attentions=all_self_attentions, 774 cross_attentions=all_cross_attentions, 775 ) 776 777 778 class BertPooler(nn.Module): 779 def __init__(self, config): 780 super().__init__() 781 self.dense = nn.Linear(config.hidden_size, config.hidden_size) 782 self.activation = nn.Tanh() 783 784 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 785 # We "pool" the model by simply taking the hidden state corresponding 786 # to the first token. 787 first_token_tensor = hidden_states[:, 0] 788 pooled_output = self.dense(first_token_tensor) 789 pooled_output = self.activation(pooled_output) 790 return pooled_output 791 792 793 794 class GreaseLM(torch.nn.Module): 795 def __init__(self, args, input_dim=768, hidden_channels=128, layers=5): 796 super(GreaseLM, self).__init__() 797 self.args = args 798 self.num_relations = 12 799 text_layers = 4 800 801 def forward(self, x, edge_index, edge_type, batch, inds, special_node_idx=None): 802 for layer in self.convs[:-1]: 803 x = F.relu(layer(x, edge_index, edge_type.view(-1))) 804 x = self.convs[-1](x, edge_index, edge_type.view(-1)) 805 if self.args.fusion: 806 special_mask = special_node_idx.view(-1, 1) 807 z = x[special_mask[:, 0]] 808 inds = torch.stack(inds, dim=-1).tolist() 809 for i in range(1, len(inds)): 810 v_c = 0 811 for j in range(i, 0, -1): 812 v_c += batch.tolist().count(j-1) 813 inds[i][0] += v_c 814 inds[i][1] += v_c 815 mask = torch.tensor(inds) 816 head = x[mask[:, 0]] 817 tail = x[mask[:, 1]] 818 avg_n = torch.cat([z, head, tail], dim=-1).view(special_mask.shape[0], 3, -1) 819 pool = torch.mean(avg_n, dim=1) 820 #pool = global_mean_pool(x, batch) 821 return pool, z 822 else: 823 inds = torch.stack(inds, dim=-1).tolist() 824 for i in range(1, len(inds)): 825 v_c = 0 826 for j in range(i, 0, -1): 827 v_c += batch.tolist().count(j-1) 828 inds[i][0] += v_c 829 inds[i][1] += v_c 830 mask = torch.tensor(inds) 831 head = x[mask[:, 0]] 832 tail = x[mask[:, 1]] 833 x = torch.cat([head, tail], dim=-1) 834 return x