import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput try: from .configuration_rnafm import RnaFmConfig except ImportError: from configuration_rnafm import RnaFmConfig def gelu(x): return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) class RnaFmLearnedPositionalEmbedding(nn.Embedding): def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int): num_embeddings_ = num_embeddings + padding_idx + 1 super().__init__(num_embeddings_, embedding_dim, padding_idx) self.max_positions = num_embeddings def forward(self, input: torch.Tensor): mask = input.ne(self.padding_idx).int() positions = (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + self.padding_idx return F.embedding( positions, self.weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse, ) class RnaFmAttention(nn.Module): def __init__(self, config: RnaFmConfig): super().__init__() self.embed_dim = config.embed_dim self.num_heads = config.attention_heads self.head_dim = config.embed_dim // config.attention_heads self.scaling = self.head_dim ** -0.5 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) def _project(self, x): tgt_len, bsz, _ = x.size() q = self.q_proj(x) * self.scaling k = self.k_proj(x) v = self.v_proj(x) q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1) k = k.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1) v = v.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1) return q, k, v, tgt_len, bsz def forward(self, x, key_padding_mask=None, output_attentions=False): q, k, v, tgt_len, bsz = self._project(x) attn_weights = torch.bmm(q, k.transpose(1, 2)) if key_padding_mask is not None: attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, tgt_len) attn_weights = attn_weights.masked_fill( key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf") ) attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, tgt_len) attn_weights_float = F.softmax(attn_weights, dim=-1, dtype=torch.float32) attn_probs = attn_weights_float.type_as(attn_weights) attn = torch.bmm(attn_probs, v) attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim) attn = self.out_proj(attn) if output_attentions: weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, tgt_len) return attn, weights return attn, None class RnaFmSdpaAttention(RnaFmAttention): def forward(self, x, key_padding_mask=None, output_attentions=False): if output_attentions: return super().forward(x, key_padding_mask, output_attentions=True) tgt_len, bsz, _ = x.size() q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) q = q.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3) k = k.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3) v = v.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3) attn_mask = None if key_padding_mask is not None: attn_mask = torch.zeros(bsz, 1, 1, tgt_len, dtype=q.dtype, device=q.device) attn_mask = attn_mask.masked_fill( key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf") ) out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) out = out.permute(2, 0, 1, 3).contiguous().view(tgt_len, bsz, self.embed_dim) return self.out_proj(out), None class RnaFmFlashAttention2(RnaFmAttention): def forward(self, x, key_padding_mask=None, output_attentions=False): if output_attentions: return super().forward(x, key_padding_mask, output_attentions=True) try: from flash_attn import flash_attn_func from flash_attn.bert_padding import pad_input, unpad_input except ImportError as e: raise ImportError( "flash_attn is required for attn_implementation='flash_attention_2'. " "Install with: pip install flash-attn --no-build-isolation" ) from e tgt_len, bsz, _ = x.size() q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) q = q.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 0, 2, 3) k = k.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 0, 2, 3) v = v.view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 0, 2, 3) orig_dtype = q.dtype if q.dtype not in (torch.float16, torch.bfloat16): q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) softmax_scale = self.head_dim ** -0.5 if key_padding_mask is not None and key_padding_mask.any(): attention_mask_bool = ~key_padding_mask q_unpad, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attention_mask_bool) k_unpad, *_ = unpad_input(k, attention_mask_bool) v_unpad, *_ = unpad_input(v, attention_mask_bool) from flash_attn import flash_attn_varlen_func out_unpad = flash_attn_varlen_func( q_unpad, k_unpad, v_unpad, cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, softmax_scale=softmax_scale, causal=False, ) out = pad_input(out_unpad, indices, bsz, tgt_len) else: out = flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=False) out = out.to(orig_dtype).permute(1, 0, 2, 3).contiguous().view(tgt_len, bsz, self.embed_dim) return self.out_proj(out), None RNAFM_ATTENTION_CLASSES = { "eager": RnaFmAttention, "sdpa": RnaFmSdpaAttention, "flash_attention_2": RnaFmFlashAttention2, } class RnaFmLayer(nn.Module): def __init__(self, config: RnaFmConfig): super().__init__() attn_cls = RNAFM_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")] self.self_attn = attn_cls(config) self.self_attn_layer_norm = nn.LayerNorm(config.embed_dim) self.fc1 = nn.Linear(config.embed_dim, config.ffn_embed_dim) self.fc2 = nn.Linear(config.ffn_embed_dim, config.embed_dim) self.final_layer_norm = nn.LayerNorm(config.embed_dim) def forward(self, x, key_padding_mask=None, output_attentions=False): residual = x x = self.self_attn_layer_norm(x) x, attn = self.self_attn(x, key_padding_mask=key_padding_mask, output_attentions=output_attentions) x = residual + x residual = x x = self.final_layer_norm(x) x = gelu(self.fc1(x)) x = self.fc2(x) x = residual + x return x, attn class RnaFmPreTrainedModel(PreTrainedModel): config_class = RnaFmConfig base_model_prefix = "rnafm" _supports_sdpa = True _supports_flash_attn_2 = True def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.constant_(module.bias, 0.0) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class RnaFmModel(RnaFmPreTrainedModel): def __init__(self, config: RnaFmConfig): super().__init__(config) self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.padding_idx) self.embed_positions = RnaFmLearnedPositionalEmbedding(config.model_max_length, config.embed_dim, config.padding_idx) self.emb_layer_norm_before = nn.LayerNorm(config.embed_dim) if config.emb_layer_norm_before else None self.layers = nn.ModuleList([RnaFmLayer(config) for _ in range(config.num_layers)]) self.emb_layer_norm_after = nn.LayerNorm(config.embed_dim) self.post_init() def forward( self, input_ids, attention_mask=None, output_hidden_states=None, output_attentions=None, return_dict=None, ): output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions return_dict = return_dict if return_dict is not None else self.config.use_return_dict if attention_mask is not None: padding_mask = attention_mask.eq(0) else: padding_mask = input_ids.eq(self.config.padding_idx) x = self.embed_tokens(input_ids) if self.config.token_dropout: x.masked_fill_((input_ids == self.config.mask_idx).unsqueeze(-1), 0.0) mask_ratio_train = 0.15 * 0.8 src_lengths = (~padding_mask).sum(-1) mask_ratio_observed = (input_ids == self.config.mask_idx).sum(-1).to(x.dtype) / src_lengths.to(x.dtype) x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None] x = x + self.embed_positions(input_ids) if self.emb_layer_norm_before is not None: x = self.emb_layer_norm_before(x) if padding_mask.any(): x = x * (1 - padding_mask.unsqueeze(-1).to(x.dtype)) else: padding_mask = None all_hidden_states = [] all_attentions = [] if output_hidden_states: all_hidden_states.append(x) x = x.transpose(0, 1) for layer in self.layers: x, attn = layer(x, key_padding_mask=padding_mask, output_attentions=output_attentions) if output_hidden_states: all_hidden_states.append(x.transpose(0, 1)) if output_attentions and attn is not None: all_attentions.append(attn) x = self.emb_layer_norm_after(x) x = x.transpose(0, 1) if output_hidden_states: all_hidden_states[-1] = x return BaseModelOutput( last_hidden_state=x, hidden_states=tuple(all_hidden_states) if output_hidden_states else None, attentions=tuple(all_attentions) if output_attentions else None, ) class RnaFmLMHead(nn.Module): def __init__(self, config: RnaFmConfig): super().__init__() self.dense = nn.Linear(config.embed_dim, config.embed_dim) self.layer_norm = nn.LayerNorm(config.embed_dim) self.decoder = nn.Linear(config.embed_dim, config.vocab_size, bias=True) def forward(self, features): x = self.dense(features) x = gelu(x) x = self.layer_norm(x) x = self.decoder(x) return x class RnaFmForMaskedLM(RnaFmPreTrainedModel): _tied_weights_keys = ["lm_head.decoder.weight"] def __init__(self, config: RnaFmConfig): super().__init__(config) self.rnafm = RnaFmModel(config) self.lm_head = RnaFmLMHead(config) self.post_init() def get_input_embeddings(self): return self.rnafm.embed_tokens def set_input_embeddings(self, value): self.rnafm.embed_tokens = value def get_output_embeddings(self): return self.lm_head.decoder def set_output_embeddings(self, new_embeddings): self.lm_head.decoder = new_embeddings def forward( self, input_ids, attention_mask=None, labels=None, output_hidden_states=None, output_attentions=None, return_dict=None, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict out = self.rnafm( input_ids, attention_mask=attention_mask, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, ) logits = self.lm_head(out.last_hidden_state) loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100, ) return MaskedLMOutput( loss=loss, logits=logits, hidden_states=out.hidden_states, attentions=out.attentions, )