File size: 4,523 Bytes
e9057bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # coding=utf-8
"""TF encoder components used by RegFM."""
from copy import deepcopy
import torch
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers.modeling_bert import *
from transformers.modeling_bert import GenomicBertModelNew as TransContextModel
class CrossAttention(nn.Module):
def __init__(self, hidden_size, num_heads=4, dropout=0.1):
super(CrossAttention, self).__init__()
self.attention = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=num_heads, dropout=dropout)
self.linear1 = nn.Linear(hidden_size, hidden_size * 4)
self.linear2 = nn.Linear(hidden_size * 4, hidden_size)
self.norm1 = nn.LayerNorm(hidden_size)
self.norm2 = nn.LayerNorm(hidden_size)
self.dropout = nn.Dropout(dropout)
self.activation = nn.ReLU()
def forward(self, query, key, value, key_padding_mask=None):
attn_output, attn_weights = self.attention(query, key, value, key_padding_mask=key_padding_mask)
query = query + self.dropout(attn_output)
query = self.norm1(query)
ff_output = self.linear2(self.dropout(self.activation(self.linear1(query))))
output = query + self.dropout(ff_output)
output = self.norm2(output)
return output, attn_weights
class TransContextForMaskedLM(BertPreTrainedModel):
def __init__(self, config):
super().__init__(config)
tf_config = deepcopy(config)
tf_config.vocab_size = 2108
tf_config.max_position_embeddings = 2112
self.bert = TransContextModel(tf_config, config)
self.cls = BertOnlyMLMHead(config)
self.vocab_size = config.vocab_size
self.config = config
self.tf_config = tf_config
self.init_weights()
def init_weights(self):
for module_ in self.named_modules():
if isinstance(module_[1], (torch.nn.Linear, torch.nn.Embedding)):
module_[1].weight.data.normal_(mean=0.0, std=self.config.initializer_range)
elif isinstance(module_[1], torch.nn.LayerNorm):
module_[1].bias.data.zero_()
module_[1].weight.data.fill_(1.0)
if isinstance(module_[1], torch.nn.Linear) and module_[1].bias is not None:
module_[1].bias.data.zero_()
def get_output_embeddings(self):
return self.cls.predictions.decoder
@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
def forward(
self,
input_ids=None,
trans_ids=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
masked_lm_labels=None,
encoder_hidden_states=None,
encoder_attention_mask=None,
lm_labels=None,
l2_lambda=0.01,
):
outputs = self.bert(
input_ids,
trans_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
)
sequence_output = outputs[0]
prediction_scores = self.cls(sequence_output)
outputs = (prediction_scores, sequence_output)
if masked_lm_labels is not None:
class_counts = torch.bincount(masked_lm_labels[masked_lm_labels != -100], minlength=self.vocab_size)
class_weights = 1.0 / (class_counts.float() + 1e-6)
loss_fct = CrossEntropyLoss(weight=class_weights)
masked_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size), masked_lm_labels.view(-1))
_, predictions = torch.max(prediction_scores, dim=-1)
masked_indices = masked_lm_labels != -100
masked_predictions = predictions[masked_indices]
masked_labels = masked_lm_labels[masked_indices]
accuracy = (masked_predictions == masked_labels).float().mean().item()
outputs = (masked_lm_loss, accuracy) + outputs
if lm_labels is not None:
prediction_scores = prediction_scores[:, :-1, :].contiguous()
lm_labels = lm_labels[:, 1:].contiguous()
loss_fct = CrossEntropyLoss()
ltr_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size), lm_labels.view(-1))
outputs = (ltr_lm_loss,) + outputs
return outputs
|