Initial RiNALMo-mega 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,160 @@
+---
+library_name: transformers
+tags:
+- RNA
+- language-model
+- biology
+license: cc-by-4.0
+---
+
+# RiNALMo-mega
+
+Minimal HuggingFace port of the **mega** (148.1M-parameter) variant of
+[RiNALMo](https://doi.org/10.1038/s41467-025-60872-5) -- a general-purpose RNA language model
+pre-trained on 36 million non-coding RNA sequences.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 30 |
+| Attention heads | 20 |
+| Embedding dimension | 640 |
+| FFN hidden dimension | 1706 (SwiGLU, floor(2/3 x 4 x embed)) |
+| Parameters | 148,056,972 |
+| Vocabulary size | 22 |
+| Positional encoding | RoPE (base=10000, non-interleaved) |
+| Normalization | LayerNorm (eps=1e-5) |
+| Architecture | Pre-LN Transformer with SwiGLU FFN |
+| Max sequence length | ~8192 (practical; RoPE has no hard limit) |
+
+**Vocabulary** (index order): `<cls>` (0), `<pad>` (1), `<eos>` (2), `<unk>` (3),
+`<mask>` (4), A (5), C (6), G (7), T (8), I (9), R (10), Y (11), K (12), M (13),
+S (14), W (15), B (16), D (17), H (18), V (19), N (20), `-` (21).
+
+Note: the tokenizer converts **U -> T** before encoding (the model was trained on T).
+
+## Pretraining
+
+- **Objective:** Masked language modeling (BERT-style, 15% mask rate)
+- **Data:** 36 million non-coding RNA sequences from multiple databases
+- **Source checkpoint:** `rinalmo_mega_pretrained.pt` from [Zenodo 15043668](https://zenodo.org/records/15043668)
+
+### Checkpoint selection
+
+The mega variant (148.1M params) offers a strong quality/cost tradeoff. Use micro for
+fast inference; use giga for maximum representation quality.
+
+## Parity Verification
+
+All 31 representation levels (embedding + 30 transformer layers) verified to be
+bit-exact (max abs diff = 0.00) against a pure-PyTorch reference that loads the
+original weights. Weight mapping verified for all 390 per-block tensors.
+Eager and SDPA implementations agree within 4e-6 on padded batches.
+
+## Related Models
+
+See the full [RiNALMo collection](https://huggingface.co/collections/Taykhoom/rinalmo-6a17c182db99be3f7a0d7bac).
+
+| Model | Parameters | Notes |
+|---|---|---|
+| [RiNALMo-micro](https://huggingface.co/Taykhoom/RiNALMo-micro) | 33.5M | Smallest variant |
+| **[RiNALMo-mega](https://huggingface.co/Taykhoom/RiNALMo-mega)** | 148.1M | This model |
+| [RiNALMo-giga](https://huggingface.co/Taykhoom/RiNALMo-giga) | 650.9M | Full model |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True)
+model.eval()
+
+sequences = ["ACUUUGGCCA", "CCCGGU"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 640) -- CLS token
+token_emb = out.last_hidden_state             # (batch, seq_len, 640)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer15_emb = out_all.hidden_states[15]       # after block 15
+```
+
+### MLM logits
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["ACU<mask>UGGCCA"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, seq_len, 22)
+```
+
+### Faster attention backends
+
+```python
+# SDPA (PyTorch 2.0+)
+model = AutoModel.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True,
+                                   attn_implementation="sdpa")
+
+# Flash Attention 2 (requires flash-attn package)
+model = AutoModel.from_pretrained("Taykhoom/RiNALMo-mega", trust_remote_code=True,
+                                   attn_implementation="flash_attention_2",
+                                   dtype=torch.bfloat16)
+```
+
+### Fine-tuning
+
+Standard HF conventions. For sequence-level tasks, pool over non-padding positions
+or use the CLS token embedding as input to a prediction head.
+
+## Implementation Notes
+
+The original RiNALMo uses Flash Attention 2.3.2 during training. This HF port exposes
+eager (standard PyTorch), SDPA, and Flash Attention 2 through HuggingFace's
+`attn_implementation` dispatch. SDPA and the dispatch interface are additions; the
+Flash backend preserves the original non-causal Flash Attention design.
+
+The model uses a non-standard Pre-LN residual: the attention residual connection is
+taken from the *normalized* input (i.e., `x = attn_ln(x); x = x + attn(x)`) rather
+than the original. The FFN uses standard Pre-LN.
+
+TokenDropout rescales embeddings by `(1 - mask_ratio_train) / (1 - mask_ratio_observed)`
+even at inference, consistent with the original training code.
+
+## Citation
+
+```bibtex
+@article{penic2025_rinalmo,
+  title   = {RiNALMo: general-purpose {RNA} language models can generalize well on structure prediction tasks},
+  author  = {Penić, Rafael Josip and Vlašić, Tin and Huber, Roland G. and Wan, Yue and Šikić, Mile},
+  journal = {Nature Communications},
+  volume  = {16},
+  number  = {1},
+  pages   = {5671},
+  year    = {2025},
+  doi     = {10.1038/s41467-025-60872-5}
+}
+```
+
+## Credits
+
+Original model and code by Penić et al. Source: [GitHub lbcb-sci/RiNALMo](https://github.com/lbcb-sci/RiNALMo).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+Apache 2.0 (code) / CC BY 4.0 (model weights), following the original repository.

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "RiNALMoForMaskedLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_rinalmo.RiNALMoConfig",
+    "AutoModel": "modeling_rinalmo.RiNALMoModel",
+    "AutoModelForMaskedLM": "modeling_rinalmo.RiNALMoForMaskedLM"
+  },
+  "attention_dropout": 0.1,
+  "cls_idx": 0,
+  "dtype": "float32",
+  "embed_dim": 640,
+  "eos_idx": 2,
+  "mask_idx": 4,
+  "mask_ratio": 0.15,
+  "mask_tkn_prob": 0.8,
+  "model_max_length": 8192,
+  "model_type": "rinalmo",
+  "num_heads": 20,
+  "num_layers": 30,
+  "padding_idx": 1,
+  "residual_dropout": 0.1,
+  "rope_base": 10000,
+  "token_dropout_active": true,
+  "transformers_version": "4.57.6",
+  "transition_dropout": 0.0,
+  "transition_factor": 4,
+  "unk_idx": 3,
+  "use_rot_emb": true,
+  "vocab_size": 22
+}

configuration_rinalmo.py

@@ -0,0 +1,54 @@
+from transformers import PretrainedConfig
+
+
+class RiNALMoConfig(PretrainedConfig):
+    model_type = "rinalmo"
+
+    auto_map = {
+        "AutoConfig": "configuration_rinalmo.RiNALMoConfig",
+        "AutoModel": "modeling_rinalmo.RiNALMoModel",
+        "AutoModelForMaskedLM": "modeling_rinalmo.RiNALMoForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 22,
+        embed_dim: int = 1280,
+        num_layers: int = 33,
+        num_heads: int = 20,
+        transition_factor: int = 4,
+        padding_idx: int = 1,
+        mask_idx: int = 4,
+        cls_idx: int = 0,
+        eos_idx: int = 2,
+        unk_idx: int = 3,
+        use_rot_emb: bool = True,
+        rope_base: int = 10000,
+        attention_dropout: float = 0.1,
+        transition_dropout: float = 0.0,
+        residual_dropout: float = 0.1,
+        token_dropout_active: bool = True,
+        mask_ratio: float = 0.15,
+        mask_tkn_prob: float = 0.8,
+        model_max_length: int = 8192,
+        **kwargs,
+    ):
+        super().__init__(padding_idx=padding_idx, **kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.transition_factor = transition_factor
+        self.mask_idx = mask_idx
+        self.cls_idx = cls_idx
+        self.eos_idx = eos_idx
+        self.unk_idx = unk_idx
+        self.use_rot_emb = use_rot_emb
+        self.rope_base = rope_base
+        self.attention_dropout = attention_dropout
+        self.transition_dropout = transition_dropout
+        self.residual_dropout = residual_dropout
+        self.token_dropout_active = token_dropout_active
+        self.mask_ratio = mask_ratio
+        self.mask_tkn_prob = mask_tkn_prob
+        self.model_max_length = model_max_length

model.safetensors

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

modeling_rinalmo.py

@@ -0,0 +1,364 @@
+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_rinalmo import RiNALMoConfig
+except ImportError:
+    from configuration_rinalmo import RiNALMoConfig
+
+
+def _rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(q, k, cos, sin):
+    cos = cos.to(q.dtype)
+    sin = sin.to(q.dtype)
+    return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)
+
+
+class RotaryPositionEmbedding(nn.Module):
+    def __init__(self, dim: int, base: int = 10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cache(self, seq_len: int, device, dtype):
+        if seq_len != self._seq_len_cached:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self._cos_cached = emb.cos()[None, None, :, :]
+            self._sin_cached = emb.sin()[None, None, :, :]
+
+    def forward(self, q, k):
+        self._update_cache(q.shape[-2], q.device, q.dtype)
+        return _apply_rotary_pos_emb(q, k, self._cos_cached, self._sin_cached)
+
+
+class RiNALMoAttention(nn.Module):
+    def __init__(self, config: RiNALMoConfig):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.num_heads = config.num_heads
+        self.head_dim = config.embed_dim // config.num_heads
+
+        self.qkv_proj = nn.Linear(config.embed_dim, 3 * config.embed_dim, bias=False)
+        self.out_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)
+        self.attn_dropout = nn.Dropout(p=config.attention_dropout)
+
+        if config.use_rot_emb:
+            self.rotary_emb = RotaryPositionEmbedding(self.head_dim, base=config.rope_base)
+        else:
+            self.rotary_emb = None
+
+    def forward(self, x, key_padding_mask=None, output_attentions=False):
+        B, T, _ = x.shape
+
+        qkv = self.qkv_proj(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if self.rotary_emb is not None:
+            q, k = self.rotary_emb(q, k)
+
+        scale = math.sqrt(self.head_dim)
+        attn = torch.matmul(q, k.transpose(-1, -2)) / scale
+
+        if key_padding_mask is not None:
+            attn = attn.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"))
+
+        attn = attn.softmax(dim=-1)
+        attn_weights = attn if output_attentions else None
+        attn = self.attn_dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
+        out = self.out_proj(out)
+        return out, attn_weights
+
+
+class RiNALMoSdpaAttention(RiNALMoAttention):
+    def forward(self, x, key_padding_mask=None, output_attentions=False):
+        if output_attentions:
+            return super().forward(x, key_padding_mask, output_attentions=True)
+
+        B, T, _ = x.shape
+
+        qkv = self.qkv_proj(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if self.rotary_emb is not None:
+            q, k = self.rotary_emb(q, k)
+
+        attn_mask = None
+        if key_padding_mask is not None:
+            attn_mask = torch.zeros(B, 1, 1, T, 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, dropout_p=0.0)
+        out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
+        out = self.out_proj(out)
+        return out, None
+
+
+class RiNALMoFlashAttention2(RiNALMoAttention):
+    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, flash_attn_varlen_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
+
+        B, T, _ = x.shape
+
+        qkv = self.qkv_proj(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(B, T, self.num_heads, self.head_dim)
+        k = k.view(B, T, self.num_heads, self.head_dim)
+        v = v.view(B, T, self.num_heads, self.head_dim)
+
+        if self.rotary_emb is not None:
+            q_t = q.transpose(1, 2)
+            k_t = k.transpose(1, 2)
+            q_t, k_t = self.rotary_emb(q_t, k_t)
+            q = q_t.transpose(1, 2)
+            k = k_t.transpose(1, 2)
+
+        orig_dtype = q.dtype
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+        if key_padding_mask is not None and key_padding_mask.any():
+            attend_mask = ~key_padding_mask
+            q_unpad, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attend_mask)
+            k_unpad, *_ = unpad_input(k, attend_mask)
+            v_unpad, *_ = unpad_input(v, attend_mask)
+
+            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,
+                causal=False,
+            )
+            out = pad_input(out_unpad.view(-1, self.embed_dim), indices, B, T)
+        else:
+            out = flash_attn_func(q, k, v, causal=False)
+            out = out.view(B, T, self.embed_dim)
+
+        out = out.to(orig_dtype)
+        out = self.out_proj(out)
+        return out, None
+
+
+RINALMO_ATTENTION_CLASSES = {
+    "eager": RiNALMoAttention,
+    "sdpa": RiNALMoSdpaAttention,
+    "flash_attention_2": RiNALMoFlashAttention2,
+}
+
+
+class RiNALMoSwiGLU(nn.Module):
+    def __init__(self, embed_dim: int, ffn_dim: int):
+        super().__init__()
+        self.linear = nn.Linear(embed_dim, ffn_dim, bias=True)
+        self.linear_gate = nn.Linear(embed_dim, ffn_dim, bias=True)
+        self.beta = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        gate = self.linear(x)
+        swish = gate * torch.sigmoid(self.beta * gate)
+        return swish * self.linear_gate(x)
+
+
+class TokenDropout(nn.Module):
+    def __init__(self, active: bool, mask_ratio: float, mask_tkn_prob: float,
+                 mask_idx: int, padding_idx: int):
+        super().__init__()
+        self.active = active
+        self.mask_ratio_train = mask_ratio * mask_tkn_prob
+        self.mask_idx = mask_idx
+        self.padding_idx = padding_idx
+
+    def forward(self, x, tokens):
+        if not self.active:
+            return x
+        pad_mask = tokens.eq(self.padding_idx)
+        src_lens = (~pad_mask).sum(dim=-1).to(x.dtype)
+        x = torch.where((tokens == self.mask_idx).unsqueeze(-1), torch.zeros_like(x), x)
+        mask_ratio_obs = (tokens == self.mask_idx).sum(dim=-1).to(x.dtype) / src_lens
+        scale = (1.0 - self.mask_ratio_train) / (1.0 - mask_ratio_obs)
+        x = x * scale[:, None, None]
+        return x
+
+
+class RiNALMoLayer(nn.Module):
+    def __init__(self, config: RiNALMoConfig):
+        super().__init__()
+        ffn_dim = int(2 / 3 * config.transition_factor * config.embed_dim)
+        attn_cls = RINALMO_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
+
+        self.attn_layer_norm = nn.LayerNorm(config.embed_dim)
+        self.attn = attn_cls(config)
+        self.out_layer_norm = nn.LayerNorm(config.embed_dim)
+        self.ffn = RiNALMoSwiGLU(config.embed_dim, ffn_dim)
+        self.ffn_dropout = nn.Dropout(p=config.transition_dropout)
+        self.ffn_down = nn.Linear(ffn_dim, config.embed_dim, bias=True)
+        self.residual_dropout_1 = nn.Dropout(p=config.residual_dropout)
+        self.residual_dropout_2 = nn.Dropout(p=config.residual_dropout)
+
+    def forward(self, x, key_padding_mask=None, output_attentions=False):
+        x = self.attn_layer_norm(x)
+        attn_out, attn_weights = self.attn(x, key_padding_mask=key_padding_mask,
+                                            output_attentions=output_attentions)
+        x = x + self.residual_dropout_1(attn_out)
+
+        residual = x
+        x = self.out_layer_norm(x)
+        x = residual + self.residual_dropout_2(self.ffn_down(self.ffn_dropout(self.ffn(x))))
+
+        return x, attn_weights
+
+
+class RiNALMoPreTrainedModel(PreTrainedModel):
+    config_class = RiNALMoConfig
+    base_model_prefix = "model"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class RiNALMoModel(RiNALMoPreTrainedModel):
+    def __init__(self, config: RiNALMoConfig):
+        super().__init__(config)
+        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.padding_idx)
+        self.token_dropout = TokenDropout(
+            active=config.token_dropout_active,
+            mask_ratio=config.mask_ratio,
+            mask_tkn_prob=config.mask_tkn_prob,
+            mask_idx=config.mask_idx,
+            padding_idx=config.padding_idx,
+        )
+        self.layers = nn.ModuleList([RiNALMoLayer(config) for _ in range(config.num_layers)])
+        self.final_layer_norm = 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:
+            key_padding_mask = attention_mask.eq(0)
+        else:
+            key_padding_mask = input_ids.eq(self.config.padding_idx)
+
+        x = self.embedding(input_ids)
+        x = self.token_dropout(x, input_ids)
+
+        all_hidden_states = []
+        all_attentions = []
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        for layer in self.layers:
+            x, attn_weights = layer(x, key_padding_mask=key_padding_mask,
+                                    output_attentions=output_attentions)
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            if output_attentions:
+                all_attentions.append(attn_weights)
+
+        x = self.final_layer_norm(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 RiNALMoForMaskedLM(RiNALMoPreTrainedModel):
+    def __init__(self, config: RiNALMoConfig):
+        super().__init__(config)
+        self.model = RiNALMoModel(config)
+        self.lm_head = RiNALMoLMHead(config)
+        self.post_init()
+
+    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.model(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)
+
+
+class RiNALMoLMHead(nn.Module):
+    def __init__(self, config: RiNALMoConfig):
+        super().__init__()
+        self.linear1 = nn.Linear(config.embed_dim, config.embed_dim)
+        self.layer_norm = nn.LayerNorm(config.embed_dim)
+        self.linear2 = nn.Linear(config.embed_dim, config.vocab_size)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = F.gelu(x)
+        x = self.layer_norm(x)
+        x = self.linear2(x)
+        return x

special_tokens_map.json

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

tokenization_rinalmo.py

@@ -0,0 +1,123 @@
+import json
+import os
+from typing import Dict, List, Optional
+
+from transformers import PreTrainedTokenizer
+
+
+_VOCAB = {
+    "<cls>":  0,
+    "<pad>":  1,
+    "<eos>":  2,
+    "<unk>":  3,
+    "<mask>": 4,
+    "A":  5,
+    "C":  6,
+    "G":  7,
+    "T":  8,
+    "I":  9,
+    "R": 10,
+    "Y": 11,
+    "K": 12,
+    "M": 13,
+    "S": 14,
+    "W": 15,
+    "B": 16,
+    "D": 17,
+    "H": 18,
+    "V": 19,
+    "N": 20,
+    "-": 21,
+}
+
+
+class RiNALMoTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizer for RiNALMo. Character-level over a 22-token RNA alphabet.
+    Converts U->T before tokenizing (the model was trained on T, not U).
+    Wraps sequences as <cls> ... <eos>.
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: Optional[str] = None,
+        cls_token: str = "<cls>",
+        pad_token: str = "<pad>",
+        eos_token: str = "<eos>",
+        unk_token: str = "<unk>",
+        mask_token: str = "<mask>",
+        **kwargs,
+    ):
+        if vocab_file is not None 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) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return dict(self._vocab)
+
+    def _tokenize(self, text: str) -> List[str]:
+        text = text.upper().replace("U", "T")
+        return list(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = 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: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        cls = [self.cls_token_id]
+        eos = [self.eos_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0 + eos
+        return cls + token_ids_0 + eos + cls + token_ids_1 + eos
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(token_ids_0, token_ids_1, True)
+        mask = [1] + [0] * len(token_ids_0) + [1]
+        if token_ids_1 is not None:
+            mask += [1] + [0] * len(token_ids_1) + [1]
+        return mask
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+    ) -> List[int]:
+        if token_ids_1 is None:
+            return [0] * (len(token_ids_0) + 2)
+        return [0] * (len(token_ids_0) + 2) + [0] * (len(token_ids_1) + 2)

tokenizer_config.json

@@ -0,0 +1,59 @@
+{
+  "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
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "model_max_length": 8192,
+  "pad_token": "<pad>",
+  "tokenizer_class": "RiNALMoTokenizer",
+  "unk_token": "<unk>",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_rinalmo.RiNALMoTokenizer",
+      null
+    ]
+  }
+}

vocab.json

@@ -0,0 +1,24 @@
+{
+  "<cls>": 0,
+  "<pad>": 1,
+  "<eos>": 2,
+  "<unk>": 3,
+  "<mask>": 4,
+  "A": 5,
+  "C": 6,
+  "G": 7,
+  "T": 8,
+  "I": 9,
+  "R": 10,
+  "Y": 11,
+  "K": 12,
+  "M": 13,
+  "S": 14,
+  "W": 15,
+  "B": 16,
+  "D": 17,
+  "H": 18,
+  "V": 19,
+  "N": 20,
+  "-": 21
+}