Upload folder using huggingface_hub

README.md

@@ -0,0 +1,146 @@
+---
+language:
+- rna
+library_name: transformers
+tags:
+- RNA
+- language-model
+license: mit
+---
+
+# mRNA-FM
+
+A 12-layer BERT-style transformer pre-trained on 45 million mRNA coding sequences (CDS) using
+codon (3-mer) tokenisation.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 12 |
+| Attention heads | 20 |
+| Embedding dimension | 1280 |
+| FFN dimension | 5120 |
+| Vocabulary size | 73 |
+| Positional encoding | Learned |
+| Architecture | ESM-1b-style pre-LN Transformer |
+| Max sequence length | 1024 codon tokens (1022 codons = 3066 nucleotides) |
+
+Vocabulary: `<cls>` (0), `<pad>` (1), `<eos>` (2), `<unk>` (3), 64 standard RNA codons
+(indices 4-67), 4 null-padding tokens (68-71), `<mask>` (72).
+
+## Pretraining
+
+- **Objective:** Masked language modelling (codon-level, 15% masking rate)
+- **Data:** RefSeq -- 45 million mRNA coding sequences
+- **Source checkpoint:** `mRNA-FM_pretrained.pth` from [cuhkaih/rnafm](https://huggingface.co/cuhkaih/rnafm)
+
+### Tokenisation
+
+mRNA-FM uses **codon (3-mer) tokenisation**: the input sequence is split into consecutive
+non-overlapping codons (triplets) and each codon is mapped to a single token. The model
+therefore only accepts sequences whose length is a **multiple of 3**.
+
+Input sequences must use **RNA notation (U, not T)**. Convert before tokenising:
+```python
+seq = seq.replace("T", "U")
+```
+
+## Parity Verification
+
+Hidden-state representations verified identical (max abs diff = 0.00) to the original
+implementation at all 13 representation levels (embedding + 12 transformer layers).
+Verified on GPU (CUDA) with PyTorch 2.7 / transformers 4.57.6. SDPA numerical
+differences are expected (~3e-4 max diff over 12 layers) and are not a correctness issue.
+
+## Related Models
+
+See the full [RNA-FM collection](https://huggingface.co/collections/Taykhoom/rna-fm-TODO).
+
+| Model | Training data | Embedding dim | Notes |
+|---|---|---|---|
+| [RNA-FM](https://huggingface.co/Taykhoom/RNA-FM) | 23.7 M ncRNA | 640 | Character tokenisation |
+| **[mRNA-FM](https://huggingface.co/Taykhoom/mRNA-FM)** | 45 M CDS | 1280 | This model |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/mRNA-FM", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/mRNA-FM", trust_remote_code=True)
+model.eval()
+
+# Sequences must be RNA (U not T) and length divisible by 3 (codons)
+sequences = [
+    "AUGGGGUGCGAUCAUACCAGCACUAAUGCCCUCCUGGGAAGUCCUCGUGUUGCA",
+    "AUGCUAGCUAGCUAGCUAUG",
+]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 1280) -- CLS token
+token_emb = out.last_hidden_state             # (batch, n_codons+2, 1280) -- per-codon
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer6_emb = out_all.hidden_states[6]
+```
+
+### MLM logits
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/mRNA-FM", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/mRNA-FM", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["AUG<mask>GCUAUG"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, n_codons+2, 73)
+```
+
+### Fine-tuning
+
+Standard HF conventions. Use the CLS token embedding (`out.last_hidden_state[:, 0, :]`) as
+input to a classification or regression head for sequence-level tasks. Mean-pool over codon
+positions (excluding CLS and EOS) for codon-level aggregation.
+
+## Implementation Notes
+
+The original implementation uses `F.multi_head_attention_forward` (eager). This HF port adds
+`attn_implementation="sdpa"` and `attn_implementation="flash_attention_2"` support, which were
+not part of the original codebase.
+
+Each codon token represents exactly one triplet of nucleotides. The tokeniser splits the raw
+sequence into non-overlapping codons; any trailing nucleotides that do not form a complete codon
+are silently discarded.
+
+## Citation
+
+```bibtex
+@article{chen2022_rnafm,
+  title   = {Interpretable {RNA} Foundation Model from Unannotated Data for Highly Accurate {RNA} Structure and Function Predictions},
+  author  = {Chen, Jiayang and Hu, Zhihang and Sun, Siqi and Tan, Qingxiong and Wang, Yixuan and Yu, Qinze and Zong, Licheng and Hong, Liang and Xiao, Jin and Shen, Tao and King, Irwin and Li, Yu},
+  journal = {arXiv preprint arXiv:2204.00300},
+  year    = {2022},
+  doi     = {10.48550/arXiv.2204.00300}
+}
+```
+
+## Credits
+
+Original model and code by Chen et al. Source: [GitHub](https://github.com/ml4bio/RNA-FM).
+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,21 @@
+{
+  "architectures": [
+    "RnaFmForMaskedLM"
+  ],
+  "attention_heads": 20,
+  "cls_idx": 0,
+  "dtype": "float32",
+  "emb_layer_norm_before": false,
+  "embed_dim": 1280,
+  "eos_idx": 2,
+  "ffn_embed_dim": 5120,
+  "mask_idx": 72,
+  "model_max_length": 1024,
+  "model_type": "rnafm",
+  "model_variant": "rna-3mer",
+  "num_layers": 12,
+  "padding_idx": 1,
+  "token_dropout": false,
+  "transformers_version": "4.57.6",
+  "vocab_size": 73
+}

configuration_rnafm.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig
+
+
+class RnaFmConfig(PretrainedConfig):
+    model_type = "rnafm"
+
+    auto_map = {
+        "AutoConfig": "configuration_rnafm.RnaFmConfig",
+        "AutoModel": "modeling_rnafm.RnaFmModel",
+        "AutoModelForMaskedLM": "modeling_rnafm.RnaFmForMaskedLM",
+        "AutoTokenizer": ["tokenization_rnafm.RnaFmTokenizer", None],
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 25,
+        num_layers: int = 12,
+        embed_dim: int = 640,
+        ffn_embed_dim: int = 5120,
+        attention_heads: int = 20,
+        padding_idx: int = 1,
+        mask_idx: int = 24,
+        cls_idx: int = 0,
+        eos_idx: int = 2,
+        token_dropout: bool = False,
+        emb_layer_norm_before: bool = True,
+        model_max_length: int = 1024,
+        model_variant: str = "rna",
+        **kwargs,
+    ):
+        super().__init__(padding_idx=padding_idx, **kwargs)
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.ffn_embed_dim = ffn_embed_dim
+        self.attention_heads = attention_heads
+        self.mask_idx = mask_idx
+        self.cls_idx = cls_idx
+        self.eos_idx = eos_idx
+        self.token_dropout = token_dropout
+        self.emb_layer_norm_before = emb_layer_norm_before
+        self.model_max_length = model_max_length
+        self.model_variant = model_variant

model.safetensors

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

modeling_rnafm.py

@@ -0,0 +1,355 @@
+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,
+        )

special_tokens_map.json

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

tokenization_rnafm.py

@@ -0,0 +1,111 @@
+import json
+import os
+
+from transformers import PreTrainedTokenizer
+
+
+_RNA_VOCAB = {
+    "<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
+    "A": 4, "C": 5, "G": 6, "U": 7,
+    "R": 8, "Y": 9, "K": 10, "M": 11,
+    "S": 12, "W": 13, "B": 14, "D": 15,
+    "H": 16, "V": 17, "N": 18, "-": 19,
+    "<null_1>": 20, "<null_2>": 21, "<null_3>": 22, "<null_4>": 23,
+    "<mask>": 24,
+}
+
+_MRNA_VOCAB = {
+    "<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
+    "GAG": 4, "AAG": 5, "GAA": 6, "CUG": 7, "CAG": 8, "GAU": 9,
+    "AAA": 10, "GUG": 11, "GAC": 12, "AUG": 13, "GCC": 14, "AAC": 15,
+    "GCU": 16, "AAU": 17, "AUC": 18, "UUC": 19, "GGA": 20, "AUU": 21,
+    "GGC": 22, "UUU": 23, "CCA": 24, "AGC": 25, "GCA": 26, "UCU": 27,
+    "CUC": 28, "ACC": 29, "CAA": 30, "CCU": 31, "UCC": 32, "ACA": 33,
+    "UUG": 34, "GUU": 35, "CUU": 36, "UAC": 37, "ACU": 38, "CCC": 39,
+    "UCA": 40, "GUC": 41, "GGU": 42, "CAC": 43, "AGU": 44, "UAU": 45,
+    "AGA": 46, "CAU": 47, "GGG": 48, "UGG": 49, "UGC": 50, "AGG": 51,
+    "UGU": 52, "AUA": 53, "CGC": 54, "UUA": 55, "GCG": 56, "CGG": 57,
+    "CCG": 58, "GUA": 59, "CUA": 60, "ACG": 61, "UCG": 62, "CGA": 63,
+    "CGU": 64, "UGA": 65, "UAA": 66, "UAG": 67,
+    "<null_1>": 68, "<null_2>": 69, "<null_3>": 70, "<null_4>": 71,
+    "<mask>": 72,
+}
+
+
+class RnaFmTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        k_mer: int = 1,
+        cls_token="<cls>",
+        pad_token="<pad>",
+        eos_token="<eos>",
+        unk_token="<unk>",
+        mask_token="<mask>",
+        **kwargs,
+    ):
+        self.k_mer = k_mer
+        if vocab_file and os.path.isfile(vocab_file):
+            with open(vocab_file) as f:
+                self._vocab = json.load(f)
+        else:
+            self._vocab = dict(_MRNA_VOCAB if k_mer == 3 else _RNA_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,
+            k_mer=k_mer,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def _tokenize(self, text):
+        if self.k_mer == 1:
+            return list(text)
+        return [text[i:i + self.k_mer] for i in range(0, len(text), self.k_mer)]
+
+    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]
+        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, 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) + [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, token_ids_1=None):
+        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,54 @@
+{
+  "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
+    },
+    "72": {
+      "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": {},
+  "k_mer": 3,
+  "mask_token": "<mask>",
+  "model_max_length": 1024,
+  "pad_token": "<pad>",
+  "tokenizer_class": "RnaFmTokenizer",
+  "unk_token": "<unk>"
+}

vocab.json

@@ -0,0 +1,75 @@
+{
+  "<cls>": 0,
+  "<pad>": 1,
+  "<eos>": 2,
+  "<unk>": 3,
+  "GAG": 4,
+  "AAG": 5,
+  "GAA": 6,
+  "CUG": 7,
+  "CAG": 8,
+  "GAU": 9,
+  "AAA": 10,
+  "GUG": 11,
+  "GAC": 12,
+  "AUG": 13,
+  "GCC": 14,
+  "AAC": 15,
+  "GCU": 16,
+  "AAU": 17,
+  "AUC": 18,
+  "UUC": 19,
+  "GGA": 20,
+  "AUU": 21,
+  "GGC": 22,
+  "UUU": 23,
+  "CCA": 24,
+  "AGC": 25,
+  "GCA": 26,
+  "UCU": 27,
+  "CUC": 28,
+  "ACC": 29,
+  "CAA": 30,
+  "CCU": 31,
+  "UCC": 32,
+  "ACA": 33,
+  "UUG": 34,
+  "GUU": 35,
+  "CUU": 36,
+  "UAC": 37,
+  "ACU": 38,
+  "CCC": 39,
+  "UCA": 40,
+  "GUC": 41,
+  "GGU": 42,
+  "CAC": 43,
+  "AGU": 44,
+  "UAU": 45,
+  "AGA": 46,
+  "CAU": 47,
+  "GGG": 48,
+  "UGG": 49,
+  "UGC": 50,
+  "AGG": 51,
+  "UGU": 52,
+  "AUA": 53,
+  "CGC": 54,
+  "UUA": 55,
+  "GCG": 56,
+  "CGG": 57,
+  "CCG": 58,
+  "GUA": 59,
+  "CUA": 60,
+  "ACG": 61,
+  "UCG": 62,
+  "CGA": 63,
+  "CGU": 64,
+  "UGA": 65,
+  "UAA": 66,
+  "UAG": 67,
+  "<null_1>": 68,
+  "<null_2>": 69,
+  "<null_3>": 70,
+  "<null_4>": 71,
+  "<mask>": 72
+}