Initial RNA-MSM Hugging Face port

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,200 @@
+---
+library_name: transformers
+tags:
+- RNA
+- language-model
+- MSA
+license: mit
+---
+
+# RNA-MSM
+
+Multiple sequence alignment-based RNA language model trained on homologous RNA
+sequence alignments from the RNAcmap pipeline.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 10 |
+| Attention heads | 12 |
+| Embedding dimension | 768 |
+| FFN hidden dimension | 3072 (GELU) |
+| Vocabulary size | 12 |
+| Positional encoding | Learned (sequence) + learned scalar (alignment row) |
+| Normalization | Pre-LayerNorm + final LayerNorm |
+| Architecture | Axial MSA Transformer (row + column self-attention) |
+| Max sequence length | 1024 tokens (including the prepended `<cls>`) |
+| Max alignment depth | 1024 rows |
+
+**Input format:** RNA-MSM takes 3D input `(batch, num_alignments, seqlen)`. Each
+alignment is a set of homologous RNA sequences of equal length (an MSA). The model
+applies row self-attention (across sequence positions) and column self-attention
+(across alignment rows) at each of the 10 transformer layers.
+
+### Vocabulary
+
+| Token | ID | Token | ID |
+|---|---|---|---|
+| `<cls>` | 0 | `U` | 7 |
+| `<pad>` | 1 | `X` | 8 |
+| `<eos>` | 2 | `N` | 9 |
+| `<unk>` | 3 | `-` | 10 |
+| `A` | 4 | `<mask>` | 11 |
+| `G` | 5 | | |
+| `C` | 6 | | |
+
+Each sequence is prepended with `<cls>` (id 0). No `<eos>` token is appended.
+
+## Pretraining
+
+- **Objective:** Masked language modeling on RNA MSAs (masking ~15% of tokens)
+- **Data:** RNA homologous sequences searched by RNAcmap from non-redundant RNA
+  databases
+- **Source checkpoint:** `RNA_MSM_pretrained.ckpt`
+  ([original Google Drive link](https://drive.google.com/file/d/11A-S13qAb5wiBi1YLs3EOrnixSDq7Q0q/view))
+
+### Checkpoint selection
+
+There is one publicly released RNA-MSM pretrained checkpoint. This is that checkpoint,
+converted from the original PyTorch Lightning `.ckpt` format.
+
+## Parity Verification
+
+Hidden-state representations verified identical (max abs diff = 0.00, exact match) to
+the reference implementation at all 11 representation levels (embedding + 10 transformer
+layers), both on padded and unpadded batches. Verified on GPU with PyTorch 2.7 /
+CUDA 12.9.
+
+## Related Models
+
+See the full [RNA-MSM collection](https://huggingface.co/collections/Taykhoom/rna-msm-6a18b5c2b0181ebbc71ff777).
+
+## Usage
+
+RNA-MSM is an **MSA model** -- it performs best when given multiple homologous
+sequences as input. For single-sequence embedding, each sequence is treated as a
+1-row MSA.
+
+### Single-sequence embedding
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
+model.eval()
+
+sequences = ["AGCUAGCUAGCU", "GCUAGCUA"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+# enc["input_ids"]: (2, 1, seqlen)  -- each sequence treated as 1-row MSA
+
+with torch.no_grad():
+    out = model(**enc)
+
+# last_hidden_state: (batch, num_alignments, seqlen, 768)
+lhs = out.last_hidden_state   # (2, 1, seqlen, 768)
+
+# Per-token embeddings for the query sequence (row 0), excluding CLS
+token_emb = lhs[:, 0, 1:, :]  # (2, seqlen-1, 768)
+
+# Mean-pool over non-padding positions for sequence-level embedding
+mask = enc["attention_mask"][:, 0, 1:].unsqueeze(-1).float()  # (2, seqlen-1, 1)
+seq_emb = (token_emb * mask).sum(1) / mask.sum(1).clamp(min=1)  # (2, 768)
+```
+
+### MSA embedding
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
+model.eval()
+
+# One MSA: 3 aligned homologous sequences of equal length
+msa = [
+    "AGCUAGCUAGCU",
+    "AGCUAGCUAGC-",
+    "AGCU--CUAGCU",
+]
+enc = tokenizer.encode_msa([msa], return_tensors="pt", padding=True)
+# enc["input_ids"]: (1, 3, seqlen)
+
+with torch.no_grad():
+    out = model(**enc)
+
+# last_hidden_state: (1, 3, seqlen, 768)
+# Use row 0 (query sequence) for downstream tasks
+query_emb = out.last_hidden_state[:, 0, 1:, :]  # (1, seqlen-1, 768)
+```
+
+### Intermediate layers
+
+```python
+with torch.no_grad():
+    out = model(**enc, output_hidden_states=True)
+
+# hidden_states: tuple of 11 tensors, each (batch, num_alignments, seqlen, 768)
+# Index 0 = embedding, 1..10 = transformer layer outputs
+layer5_emb = out.hidden_states[5][:, 0, :, :]  # (batch, seqlen, 768)
+```
+
+### MLM logits
+
+```python
+from transformers import AutoModelForMaskedLM
+
+mlm = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
+mlm.eval()
+
+enc = tokenizer(["AGCU<mask>AGCU"], return_tensors="pt", padding=True)
+with torch.no_grad():
+    logits = mlm(**enc).logits  # (1, 1, seqlen, 12)
+```
+
+### Fine-tuning
+
+For sequence-level downstream tasks (e.g., solvent accessibility), extract the
+embedding from the query row (row 0) of the last hidden state, then apply a
+prediction head. The model's attention maps (row attention) are also useful for
+2D structural tasks (e.g., secondary structure prediction).
+
+## Implementation Notes
+
+RNA-MSM uses **axial attention**: each transformer layer applies row self-attention
+(attending across sequence positions, summed over alignment rows) followed by column
+self-attention (attending across alignment rows per position). This custom attention
+pattern is not compatible with `attn_implementation="sdpa"` or
+`attn_implementation="flash_attention_2"` -- only `"eager"` is supported.
+
+`last_hidden_state` has shape `(batch, num_alignments, seqlen, embed_dim)` -- note
+the 4D output, reflecting the MSA structure. For single-sequence use (1-row MSA),
+this is `(batch, 1, seqlen, embed_dim)`.
+
+## Citation
+
+```bibtex
+@article{zhang2024_rnamsm,
+  title   = {Multiple sequence alignment-based {RNA} language model and its application to structural inference},
+  author  = {Zhang, Yikun and Lang, Mei and Jiang, Jiuhong and Gao, Zhiqiang and Xu, Fan and Litfin, Thomas and Chen, Ke and Singh, Jaswinder and Huang, Xiansong and Song, Guoli and Tian, Yonghong and Zhan, Jian and Chen, Jie and Zhou, Yaoqi},
+  journal = {Nucleic Acids Research},
+  volume  = {52},
+  number  = {1},
+  pages   = {e3},
+  year    = {2024},
+  doi     = {10.1093/nar/gkad1031}
+}
+```
+
+## Credits
+
+Original model and code by Zhang et al. Source: [GitHub](https://github.com/yikunpku/RNA-MSM).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+MIT, following the original repository.

config.json

@@ -0,0 +1,29 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_rnamsm.RNAMSMConfig",
+    "AutoModel": "modeling_rnamsm.RNAMSMModel",
+    "AutoModelForMaskedLM": "modeling_rnamsm.RNAMSMForMaskedLM"
+  },
+  "activation_dropout": 0.1,
+  "architectures": [
+    "RNAMSMForMaskedLM"
+  ],
+  "attention_dropout": 0.1,
+  "cls_idx": 0,
+  "dropout": 0.1,
+  "embed_dim": 768,
+  "embed_positions_msa": true,
+  "eos_idx": 2,
+  "ffn_embed_dim": 3072,
+  "mask_idx": 11,
+  "max_alignments": 1024,
+  "max_positions": 1024,
+  "max_tokens_per_msa": 16384,
+  "model_type": "rnamsm",
+  "num_attention_heads": 12,
+  "num_layers": 10,
+  "padding_idx": 1,
+  "transformers_version": "4.57.6",
+  "vocab_size": 12,
+  "model_max_length": 1024
+}

configuration_rnamsm.py

@@ -0,0 +1,48 @@
+from transformers import PretrainedConfig
+
+
+class RNAMSMConfig(PretrainedConfig):
+    model_type = "rnamsm"
+
+    auto_map = {
+        "AutoConfig": "configuration_rnamsm.RNAMSMConfig",
+        "AutoModel": "modeling_rnamsm.RNAMSMModel",
+        "AutoModelForMaskedLM": "modeling_rnamsm.RNAMSMForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size=12,
+        num_layers=10,
+        embed_dim=768,
+        num_attention_heads=12,
+        ffn_embed_dim=3072,
+        padding_idx=1,
+        mask_idx=11,
+        cls_idx=0,
+        eos_idx=2,
+        dropout=0.1,
+        attention_dropout=0.1,
+        activation_dropout=0.1,
+        max_positions=1024,
+        max_alignments=1024,
+        max_tokens_per_msa=16384,
+        embed_positions_msa=True,
+        **kwargs,
+    ):
+        super().__init__(padding_idx=padding_idx, **kwargs)
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.num_attention_heads = num_attention_heads
+        self.ffn_embed_dim = ffn_embed_dim
+        self.mask_idx = mask_idx
+        self.cls_idx = cls_idx
+        self.eos_idx = eos_idx
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.max_positions = max_positions
+        self.max_alignments = max_alignments
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.embed_positions_msa = embed_positions_msa

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3998fb8289d98cf53944fc14da4157a40c03dffecf0efefd7e76044ed16a0095
+size 383678288

modeling_rnamsm.py

@@ -0,0 +1,420 @@
+import math
+from typing import Optional, Tuple
+
+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_rnamsm import RNAMSMConfig
+except ImportError:
+    from configuration_rnamsm import RNAMSMConfig
+
+
+def gelu(x):
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+class RNAMSMLMHead(nn.Module):
+    def __init__(self, config: RNAMSMConfig, embed_tokens_weight: nn.Parameter):
+        super().__init__()
+        self.dense = nn.Linear(config.embed_dim, config.embed_dim)
+        self.layer_norm = nn.LayerNorm(config.embed_dim)
+        self.weight = embed_tokens_weight
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+    def forward(self, x):
+        x = self.dense(x)
+        x = gelu(x)
+        x = self.layer_norm(x)
+        return F.linear(x, self.weight) + self.bias
+
+
+class LearnedPositionalEmbedding(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, tokens: torch.Tensor) -> torch.Tensor:
+        mask = tokens.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 NormalizedResidualBlock(nn.Module):
+    def __init__(self, layer: nn.Module, embedding_dim: int, dropout: float):
+        super().__init__()
+        self.layer = layer
+        self.layer_norm = nn.LayerNorm(embedding_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def forward(self, x, *args, **kwargs):
+        residual = x
+        x = self.layer_norm(x)
+        outputs = self.layer(x, *args, **kwargs)
+        if isinstance(outputs, tuple):
+            x, *out = outputs
+        else:
+            x, out = outputs, None
+        x = self.dropout_module(x)
+        x = residual + x
+        if out is not None:
+            return (x,) + tuple(out)
+        return x
+
+
+class FeedForwardNetwork(nn.Module):
+    def __init__(self, embedding_dim: int, ffn_embedding_dim: int,
+                 activation_dropout: float, max_tokens_per_msa: int):
+        super().__init__()
+        self.fc1 = nn.Linear(embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, embedding_dim)
+        self.activation_fn = nn.GELU()
+        self.activation_dropout_module = nn.Dropout(activation_dropout)
+        self.max_tokens_per_msa = max_tokens_per_msa
+
+    def forward(self, x):
+        x = self.activation_fn(self.fc1(x))
+        x = self.activation_dropout_module(x)
+        return self.fc2(x)
+
+
+class RowSelfAttention(nn.Module):
+    """Self-attention across columns (sequence positions), summed over MSA rows."""
+
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float, max_tokens_per_msa: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.attn_shape = "hnij"
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def align_scaling(self, q):
+        return self.scaling / math.sqrt(q.size(0))
+
+    def compute_attention_weights(self, x, scaling, padding_mask=None):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        q = q * scaling
+        if padding_mask is not None:
+            q = q * (1 - padding_mask.permute(1, 2, 0).unsqueeze(3).unsqueeze(4).to(q))
+        attn_weights = torch.einsum(f"rinhd,rjnhd->{self.attn_shape}", q, k)
+        if padding_mask is not None:
+            attn_weights = attn_weights.masked_fill(
+                padding_mask[:, 0].unsqueeze(0).unsqueeze(2), -10000.0)
+        return attn_weights
+
+    def compute_attention_update(self, x, attn_probs):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        context = torch.einsum(f"{self.attn_shape},rjnhd->rinhd", attn_probs, v)
+        context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+        return self.out_proj(context)
+
+    def _batched_forward(self, x, padding_mask=None):
+        num_rows, num_cols, batch_size, _ = x.size()
+        max_rows = max(1, self.max_tokens_per_msa // num_cols)
+        scaling = self.align_scaling(x)
+        attns = 0
+        for start in range(0, num_rows, max_rows):
+            pm = padding_mask[:, start:start + max_rows] if padding_mask is not None else None
+            attns = attns + self.compute_attention_weights(x[start:start + max_rows], scaling, pm)
+        attn_probs = attns.softmax(-1)
+        attn_probs = self.dropout_module(attn_probs)
+        outputs = [self.compute_attention_update(x[start:start + max_rows], attn_probs)
+                   for start in range(0, num_rows, max_rows)]
+        return torch.cat(outputs, 0), attn_probs
+
+    def forward(self, x, self_attn_mask=None, self_attn_padding_mask=None):
+        num_rows, num_cols, batch_size, _ = x.size()
+        if num_rows * num_cols > self.max_tokens_per_msa and not torch.is_grad_enabled():
+            return self._batched_forward(x, self_attn_padding_mask)
+        scaling = self.align_scaling(x)
+        attn_weights = self.compute_attention_weights(x, scaling, self_attn_padding_mask)
+        attn_probs = attn_weights.softmax(-1)
+        attn_probs = self.dropout_module(attn_probs)
+        output = self.compute_attention_update(x, attn_probs)
+        return output, attn_probs
+
+
+class ColumnSelfAttention(nn.Module):
+    """Self-attention across MSA rows (alignment depth) per sequence position."""
+
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float, max_tokens_per_msa: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def compute_attention_update(self, x, self_attn_padding_mask=None):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        if num_rows == 1:
+            attn_probs = torch.ones(self.num_heads, num_cols, batch_size, 1, 1,
+                                    device=x.device, dtype=x.dtype)
+            output = self.out_proj(self.v_proj(x))
+        else:
+            q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            q = q * self.scaling
+            attn_weights = torch.einsum("icnhd,jcnhd->hcnij", q, k)
+            if self_attn_padding_mask is not None:
+                attn_weights = attn_weights.masked_fill(
+                    self_attn_padding_mask.permute(2, 0, 1).unsqueeze(0).unsqueeze(3), -10000.0)
+            attn_probs = attn_weights.softmax(-1)
+            attn_probs = self.dropout_module(attn_probs)
+            context = torch.einsum("hcnij,jcnhd->icnhd", attn_probs, v)
+            context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+            output = self.out_proj(context)
+        return output, attn_probs
+
+    def _batched_forward(self, x, self_attn_padding_mask=None):
+        num_rows, num_cols, batch_size, _ = x.size()
+        max_cols = max(1, self.max_tokens_per_msa // num_rows)
+        outputs, attns = [], []
+        for start in range(0, num_cols, max_cols):
+            pm = (self_attn_padding_mask[:, :, start:start + max_cols]
+                  if self_attn_padding_mask is not None else None)
+            out, attn = self.compute_attention_update(x[:, start:start + max_cols], pm)
+            outputs.append(out)
+            attns.append(attn)
+        return torch.cat(outputs, 1), torch.cat(attns, 1)
+
+    def forward(self, x, self_attn_mask=None, self_attn_padding_mask=None):
+        num_rows, num_cols, batch_size, _ = x.size()
+        if num_rows * num_cols > self.max_tokens_per_msa and not torch.is_grad_enabled():
+            return self._batched_forward(x, self_attn_padding_mask)
+        return self.compute_attention_update(x, self_attn_padding_mask)
+
+
+class AxialTransformerLayer(nn.Module):
+    def __init__(self, config: RNAMSMConfig):
+        super().__init__()
+        self.row_self_attention = NormalizedResidualBlock(
+            RowSelfAttention(config.embed_dim, config.num_attention_heads,
+                             config.attention_dropout, config.max_tokens_per_msa),
+            config.embed_dim, config.dropout,
+        )
+        self.column_self_attention = NormalizedResidualBlock(
+            ColumnSelfAttention(config.embed_dim, config.num_attention_heads,
+                                config.attention_dropout, config.max_tokens_per_msa),
+            config.embed_dim, config.dropout,
+        )
+        self.feed_forward_layer = NormalizedResidualBlock(
+            FeedForwardNetwork(config.embed_dim, config.ffn_embed_dim,
+                               config.activation_dropout, config.max_tokens_per_msa),
+            config.embed_dim, config.dropout,
+        )
+
+    def forward(self, x, padding_mask=None, output_attentions=False):
+        x, row_attn = self.row_self_attention(x, self_attn_padding_mask=padding_mask)
+        x, col_attn = self.column_self_attention(x, self_attn_padding_mask=padding_mask)
+        x = self.feed_forward_layer(x)
+        return x, row_attn, col_attn
+
+
+class RNAMSMPreTrainedModel(PreTrainedModel):
+    config_class = RNAMSMConfig
+    base_model_prefix = "rnamsm"
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+
+
+class RNAMSMModel(RNAMSMPreTrainedModel):
+    """
+    RNA-MSM backbone: MSA Transformer that processes multiple-sequence-aligned RNA
+    sequences and produces per-position embeddings for each alignment row.
+
+    Input: input_ids of shape (batch, num_alignments, seqlen)
+    Output: last_hidden_state of shape (batch, num_alignments, seqlen, embed_dim)
+    """
+
+    def __init__(self, config: RNAMSMConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim,
+                                         padding_idx=config.padding_idx)
+        self.embed_positions = LearnedPositionalEmbedding(
+            config.max_positions, config.embed_dim, config.padding_idx)
+        if config.embed_positions_msa:
+            self.msa_position_embedding = nn.Parameter(
+                0.01 * torch.randn(1, config.max_alignments, 1, 1))
+        else:
+            self.register_parameter("msa_position_embedding", None)
+        self.dropout_module = nn.Dropout(config.dropout)
+        self.emb_layer_norm_before = nn.LayerNorm(config.embed_dim)
+        self.emb_layer_norm_after = nn.LayerNorm(config.embed_dim)
+        self.layers = nn.ModuleList([AxialTransformerLayer(config)
+                                     for _ in range(config.num_layers)])
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = 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
+
+        assert input_ids.ndim == 3, (
+            "RNA-MSM expects 3D input_ids of shape (batch, num_alignments, seqlen). "
+            "For single sequences, use tokenizer which produces (batch, 1, seqlen).")
+
+        batch_size, num_alignments, seqlen = input_ids.size()
+
+        # HF convention: attention_mask 1=attend, 0=pad -> padding_mask True=padding
+        if attention_mask is not None:
+            padding_mask = attention_mask.eq(0)
+        else:
+            padding_mask = input_ids.eq(self.config.padding_idx)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        # (B, R, C) -> embed: (B, R, C, D)
+        x = self.embed_tokens(input_ids)
+        x = x + self.embed_positions(
+            input_ids.view(batch_size * num_alignments, seqlen)
+        ).view(batch_size, num_alignments, seqlen, self.config.embed_dim)
+
+        if self.msa_position_embedding is not None:
+            if num_alignments > self.config.max_alignments:
+                raise RuntimeError(
+                    f"MSA depth {num_alignments} exceeds max_alignments "
+                    f"{self.config.max_alignments}.")
+            x = x + self.msa_position_embedding[:, :num_alignments]
+
+        x = self.emb_layer_norm_before(x)
+        x = self.dropout_module(x)
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).to(x))
+
+        all_hidden_states = []
+        all_row_attentions = []
+        all_col_attentions = []
+
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        # (B, R, C, D) -> (R, C, B, D) for axial attention
+        x = x.permute(1, 2, 0, 3)
+
+        for layer in self.layers:
+            x, row_attn, col_attn = layer(x, padding_mask=padding_mask,
+                                           output_attentions=output_attentions)
+            if output_hidden_states:
+                all_hidden_states.append(x.permute(2, 0, 1, 3))
+            if output_attentions:
+                all_row_attentions.append(row_attn)
+                all_col_attentions.append(col_attn)
+
+        x = self.emb_layer_norm_after(x)
+        x = x.permute(2, 0, 1, 3)  # (R, C, B, D) -> (B, R, C, D)
+
+        if output_hidden_states:
+            all_hidden_states[-1] = x
+
+        if not return_dict:
+            return tuple(v for v in [
+                x,
+                tuple(all_hidden_states) if output_hidden_states else None,
+                tuple(all_row_attentions) if output_attentions else None,
+            ] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
+            attentions=tuple(all_row_attentions) if output_attentions else None,
+        )
+
+
+class RNAMSMForMaskedLM(RNAMSMPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: RNAMSMConfig):
+        super().__init__(config)
+        self.rnamsm = RNAMSMModel(config)
+        self.lm_head = RNAMSMLMHead(config, self.rnamsm.embed_tokens.weight)
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        out = self.rnamsm(
+            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[0] if not return_dict else 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,
+            )
+
+        if not return_dict:
+            output = (logits,) + out[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=out.hidden_states,
+            attentions=out.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenization_rnamsm.py

@@ -0,0 +1,242 @@
+import json
+import os
+from typing import Dict, List, Optional, Union
+
+import torch
+from transformers import PreTrainedTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+
+
+_VOCAB = {
+    "<cls>":  0,
+    "<pad>":  1,
+    "<eos>":  2,
+    "<unk>":  3,
+    "A":      4,
+    "G":      5,
+    "C":      6,
+    "U":      7,
+    "X":      8,
+    "N":      9,
+    "-":      10,
+    "<mask>": 11,
+}
+
+
+class RNAMSMTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizer for RNA-MSM.
+
+    Vocabulary: <cls>(0) <pad>(1) <eos>(2) <unk>(3) A(4) G(5) C(6) U(7) X(8) N(9) -(10) <mask>(11)
+
+    RNA-MSM is an MSA Transformer: it always expects 3D input
+    (batch, num_alignments, seqlen). This tokenizer treats each input string
+    as a single-sequence MSA (1 alignment row), so the standard __call__ API:
+
+        enc = tokenizer(["AGCU", "GAUC"], return_tensors="pt", padding=True)
+        # enc.input_ids: (2, 1, T)  -- batch of 2 single-sequence MSAs
+
+    For real MSAs (multiple aligned sequences), use encode_msa():
+
+        enc = tokenizer.encode_msa([["AGCU--", "AGCUUU"]], return_tensors="pt")
+        # enc["input_ids"]: (1, 2, T)  -- 1 MSA with 2 alignment rows
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        cls_token="<cls>",
+        pad_token="<pad>",
+        eos_token="<eos>",
+        unk_token="<unk>",
+        mask_token="<mask>",
+        **kwargs,
+    ):
+        if vocab_file and os.path.isfile(vocab_file):
+            with open(vocab_file) as f:
+                self._vocab = json.load(f)
+        else:
+            self._vocab = dict(_VOCAB)
+        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
+        super().__init__(
+            cls_token=cls_token,
+            pad_token=pad_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            mask_token=mask_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def _tokenize(self, text):
+        return list(text)
+
+    def _convert_token_to_id(self, token):
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index):
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        path = os.path.join(save_directory, fname)
+        with open(path, "w") as f:
+            json.dump(self._vocab, f, indent=2)
+        return (path,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0
+        return cls + token_ids_0 + cls + token_ids_1
+
+    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None,
+                                already_has_special_tokens=False):
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0, token_ids_1, already_has_special_tokens=True)
+        mask = [1] + [0] * len(token_ids_0)
+        if token_ids_1 is not None:
+            mask += [1] + [0] * len(token_ids_1)
+        return mask
+
+    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
+        if token_ids_1 is None:
+            return [0] + token_ids_0
+        return [0] + token_ids_0 + [0] + token_ids_1
+
+    def __call__(
+        self,
+        text,
+        text_pair=None,
+        add_special_tokens=True,
+        padding=False,
+        truncation=False,
+        max_length=None,
+        return_tensors=None,
+        **kwargs,
+    ):
+        """
+        Tokenize one or more sequences, each treated as a 1-row MSA.
+
+        text: str or List[str]
+        Returns dict with input_ids of shape (batch, 1, seqlen) and
+        attention_mask of shape (batch, 1, seqlen).
+        """
+        if isinstance(text, str):
+            sequences = [text]
+        else:
+            sequences = list(text)
+
+        encoded = []
+        for seq in sequences:
+            ids = self._tokenize_single(seq, add_special_tokens)
+            encoded.append(ids)
+
+        if padding and len(encoded) > 1:
+            max_len = max(len(ids) for ids in encoded)
+            pad_id = self.pad_token_id
+            encoded = [ids + [pad_id] * (max_len - len(ids)) for ids in encoded]
+
+        input_ids = [[ids] for ids in encoded]
+        attention_mask = [[[1 if t != self.pad_token_id else 0 for t in ids]]
+                          for ids in encoded]
+
+        if return_tensors == "pt":
+            input_ids = torch.tensor(input_ids, dtype=torch.long)
+            attention_mask = torch.tensor(attention_mask, dtype=torch.long)
+            return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})
+
+        return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})
+
+    def _tokenize_single(self, sequence, add_special_tokens=True):
+        tokens = list(sequence)
+        ids = [self._convert_token_to_id(t) for t in tokens]
+        if add_special_tokens:
+            ids = [self.cls_token_id] + ids
+        return ids
+
+    def encode_msa(
+        self,
+        msas,
+        add_special_tokens=True,
+        padding=False,
+        return_tensors=None,
+    ):
+        """
+        Tokenize a batch of MSAs.
+
+        msas: List[List[str]]
+            Each inner list is one MSA (multiple aligned sequences of equal length).
+            All sequences within an MSA must have the same length.
+
+        Returns dict with:
+            input_ids: (batch, max_alignments, max_seqlen)
+            attention_mask: (batch, max_alignments, max_seqlen)
+        """
+        if isinstance(msas[0], str):
+            msas = [msas]
+
+        max_rows = max(len(msa) for msa in msas)
+        max_seqlen = max(
+            len(self._tokenize_single(seq, add_special_tokens))
+            for msa in msas for seq in msa
+        )
+
+        pad_id = self.pad_token_id
+        batch_ids = []
+        batch_mask = []
+
+        for msa in msas:
+            msa_ids = []
+            msa_mask = []
+            for seq in msa:
+                ids = self._tokenize_single(seq, add_special_tokens)
+                if padding:
+                    pad_len = max_seqlen - len(ids)
+                    mask = [1] * len(ids) + [0] * pad_len
+                    ids = ids + [pad_id] * pad_len
+                else:
+                    mask = [1] * len(ids)
+                msa_ids.append(ids)
+                msa_mask.append(mask)
+
+            if padding:
+                pad_row = [pad_id] * max_seqlen
+                pad_mask_row = [0] * max_seqlen
+                while len(msa_ids) < max_rows:
+                    msa_ids.append(pad_row)
+                    msa_mask.append(pad_mask_row)
+
+            batch_ids.append(msa_ids)
+            batch_mask.append(msa_mask)
+
+        if return_tensors == "pt":
+            batch_ids = torch.tensor(batch_ids, dtype=torch.long)
+            batch_mask = torch.tensor(batch_mask, dtype=torch.long)
+            return BatchEncoding({"input_ids": batch_ids, "attention_mask": batch_mask})
+
+        return BatchEncoding({"input_ids": batch_ids, "attention_mask": batch_mask})
+
+    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+        tokens = [self._convert_id_to_token(i) for i in token_ids]
+        if skip_special_tokens:
+            special = {self.cls_token, self.pad_token, self.eos_token,
+                       self.unk_token, self.mask_token}
+            tokens = [t for t in tokens if t not in special]
+        return "".join(tokens)
+
+    def num_special_tokens_to_add(self, pair=False):
+        return 1

tokenizer_config.json

@@ -0,0 +1,56 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<cls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_rnamsm.RNAMSMTokenizer", null]
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "model_max_length": 1024,
+  "pad_token": "<pad>",
+  "tokenizer_class": "RNAMSMTokenizer",
+  "unk_token": "<unk>"
+}

vocab.json

@@ -0,0 +1,14 @@
+{
+  "<cls>": 0,
+  "<pad>": 1,
+  "<eos>": 2,
+  "<unk>": 3,
+  "A": 4,
+  "G": 5,
+  "C": 6,
+  "U": 7,
+  "X": 8,
+  "N": 9,
+  "-": 10,
+  "<mask>": 11
+}