| |
| """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 |
|
|
