Upload folder using huggingface_hub

README.md

@@ -0,0 +1,162 @@
+---
+language:
+- rna
+library_name: transformers
+tags:
+- RNA
+- language-model
+license: apache-2.0
+---
+
+# RNAErnie
+
+RNAErnie is a BERT-based RNA language model pretrained on RNACentral using a
+motif-aware masking strategy with type-guided fine-tuning. It uses a DNA-style
+vocabulary (T instead of U) and extends the token vocabulary with 28 ncRNA
+type labels to enable type-guided learning.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 12 |
+| Attention heads | 12 |
+| Embedding dimension | 768 |
+| Intermediate size | 3072 |
+| Vocabulary size | 39 |
+| Positional encoding | Absolute learned |
+| Architecture | Post-LN BERT / ERNIE |
+| Max sequence length | 512 |
+
+**Vocabulary:** Special tokens `[PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, [MASK]=4, [DEL]=5, [IND]=6`;
+ncRNA type labels at indices 7-34 (RNaseMRPRNA, RNasePRNA, SRPRNA, YRNA, antisenseRNA,
+autocatalyticallysplicedintron, guideRNA, hammerheadribozyme, lncRNA, miRNA, miscRNA,
+ncRNA, other, piRNA, premiRNA, precursorRNA, rRNA, ribozyme, sRNA, scRNA, scaRNA,
+siRNA, snRNA, snoRNA, tRNA, telomeraseRNA, tmRNA, vaultRNA);
+nucleotides `A=35, T=36, C=37, G=38`.
+
+**Tokenisation note:** Input U is silently converted to T. The model was pretrained
+with DNA-style T notation.
+
+## Pretraining
+
+- **Objective:** Masked language modelling (MLM) with motif-aware masking
+- **Data:** RNACentral (sequences with length <= 512)
+- **Source checkpoint:** `model_state.pdparams` from the original PaddlePaddle repository
+
+### Checkpoint selection
+
+There is a single publicly released RNAErnie checkpoint
+(`output/BERT,ERNIE,MOTIF,PROMPT/checkpoint_final/model_state.pdparams`),
+corresponding to the `BERT,ERNIE,MOTIF,PROMPT` pretraining variant described in the
+paper.
+
+## Parity Verification
+
+Hidden-state representations verified identical (max abs diff < 7e-6) to a
+standalone PyTorch reference implementation built directly from the raw
+PaddlePaddle weights at all 13 representation levels (embedding + 12 layers).
+Verified on GPU with PyTorch 2.7 / CUDA 12.
+
+**Note on weight conversion:** PaddlePaddle stores `nn.Linear` weights as
+`(in_features, out_features)`, the transpose of PyTorch's `(out_features, in_features)`.
+All linear layer weights (attention projections, FFN, pooler, MLM transform) are
+transposed during conversion; embedding tables and bias vectors are copied as-is.
+
+## Implementation Notes
+
+The original implementation uses PaddlePaddle's ERNIE/TransformerEncoderLayer
+backbone. This HF port re-implements the identical Post-LN BERT architecture in
+pure PyTorch and adds `attn_implementation="sdpa"` and
+`attn_implementation="flash_attention_2"` support, which were not part of the
+original codebase.
+
+## Related Models
+
+See the full [RNAErnie collection](https://huggingface.co/collections/Taykhoom/rnaernie-6a219927c11fdcccedb243db).
+
+| Model | Context | Training data | Notes |
+|---|---|---|---|
+| **[RNAErnie](./)** | **512** | **RNACentral (nts<=512)** | **This model; PaddlePaddle ERNIE backbone** |
+| [RNAErnie2](../RNAErnie2) | 2048 | RNACentral v22 (~31M seqs) | Retrained; PyTorch BERT |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNAErnie", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/RNAErnie", trust_remote_code=True)
+model.eval()
+
+sequences = ["AUGCAUGCAUGC", "GCUGCAUGCUAGC"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]  # (batch, 768) -- CLS token
+token_emb = out.last_hidden_state           # (batch, seq_len, 768)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer6_emb = out_all.hidden_states[6]       # (batch, seq_len, 768)
+```
+
+### MLM logits
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNAErnie", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNAErnie", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["ATG[MASK]ATG"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits  # (1, seq_len, 39)
+```
+
+### SDPA / Flash Attention 2
+
+```python
+model = AutoModel.from_pretrained(
+    "Taykhoom/RNAErnie",
+    attn_implementation="sdpa",   # or "flash_attention_2"
+    trust_remote_code=True,
+)
+```
+
+### Fine-tuning
+
+Standard HF conventions. For sequence-level tasks, use the CLS token embedding
+(`last_hidden_state[:, 0, :]`) as input to a classification head. For type-guided
+fine-tuning (as in the paper), prepend the ncRNA type label token to the input.
+
+## Citation
+
+```bibtex
+@article{wang2024_rnaernie,
+  title   = {Multi-purpose {RNA} language modelling with motif-aware pretraining and type-guided fine-tuning},
+  author  = {Wang, Ning and Bian, Jiang and Li, Yuchen and Li, Xuhong and Mumtaz, Shahid and Kong, Linghe and Xiong, Haoyi},
+  journal = {Nature Machine Intelligence},
+  volume  = {6},
+  pages   = {548--557},
+  year    = {2024},
+  doi     = {10.1038/s42256-024-00836-4}
+}
+```
+
+## Credits
+
+Original model and code by Wang et al. Source: [GitHub](https://github.com/CatIIIIIIII/RNAErnie).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+Apache 2.0, following the original repository.

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "RNAErnieForMaskedLM"
+  ],
+  "model_type": "rnaernie",
+  "auto_map": {
+    "AutoConfig": "configuration_rnaernie.RNAErnieConfig",
+    "AutoModel": "modeling_rnaernie.RNAErnieModel",
+    "AutoModelForMaskedLM": "modeling_rnaernie.RNAErnieForMaskedLM"
+  },
+  "vocab_size": 39,
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 12,
+  "intermediate_size": 3072,
+  "hidden_act": "relu",
+  "hidden_dropout_prob": 0.1,
+  "attention_probs_dropout_prob": 0.1,
+  "max_position_embeddings": 513,
+  "type_vocab_size": 2,
+  "layer_norm_eps": 1e-12,
+  "pad_token_id": 0,
+  "initializer_range": 0.02,
+  "transformers_version": "4.57.6"
+}

configuration_rnaernie.py

@@ -0,0 +1,40 @@
+from transformers import PretrainedConfig
+
+
+class RNAErnieConfig(PretrainedConfig):
+    model_type = "rnaernie"
+
+    auto_map = {
+        "AutoConfig": "configuration_rnaernie.RNAErnieConfig",
+        "AutoModel": "modeling_rnaernie.RNAErnieModel",
+        "AutoModelForMaskedLM": "modeling_rnaernie.RNAErnieForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 39,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        hidden_act: str = "relu",
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 513,
+        type_vocab_size: int = 2,
+        layer_norm_eps: float = 1e-12,
+        pad_token_id: int = 0,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.layer_norm_eps = layer_norm_eps

model.safetensors

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

modeling_rnaernie.py

@@ -0,0 +1,407 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput
+
+try:
+    from .configuration_rnaernie import RNAErnieConfig
+except ImportError:
+    from configuration_rnaernie import RNAErnieConfig
+
+
+class RNAErnieSelfAttention(nn.Module):
+
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        return x.view(B, T, self.num_attention_heads, self.attention_head_size).permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        q = self._split_heads(self.query(hidden_states))
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+
+        scale = math.sqrt(self.attention_head_size)
+        scores = torch.matmul(q, k.transpose(-1, -2)) / scale
+        if key_padding_mask is not None:
+            scores = scores.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
+        probs = F.softmax(scores, dim=-1)
+        probs = self.dropout(probs)
+        context = torch.matmul(probs, v)
+
+        B, _, T, _ = context.shape
+        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+
+        if output_attentions:
+            return context, probs
+        return context, None
+
+
+class RNAErnieSdpaSelfAttention(RNAErnieSelfAttention):
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(hidden_states, key_padding_mask, output_attentions=True)
+
+        B, T, _ = hidden_states.shape
+        q = self._split_heads(self.query(hidden_states))
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+
+        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[:, None, None, :], float("-inf"))
+
+        context = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+        return context, None
+
+
+class RNAErnieFlashSelfAttention(RNAErnieSelfAttention):
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(hidden_states, 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, _ = hidden_states.shape
+        q = self._split_heads(self.query(hidden_states))
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        orig_dtype = q.dtype
+        if orig_dtype not in (torch.float16, torch.bfloat16):
+            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
+
+        if key_padding_mask is not None and key_padding_mask.any():
+            attend = ~key_padding_mask
+            q_u, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attend)
+            k_u, _, _, _, _ = unpad_input(k, attend)
+            v_u, _, _, _, _ = unpad_input(v, attend)
+            out_u = flash_attn_varlen_func(
+                q_u, k_u, v_u,
+                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_u, indices, B, T)
+        else:
+            out = flash_attn_func(q, k, v, causal=False)
+
+        out = out.to(orig_dtype).reshape(B, T, self.all_head_size)
+        return out, None
+
+
+RNAERNIE_SELF_ATTENTION_CLASSES = {
+    "eager": RNAErnieSelfAttention,
+    "sdpa": RNAErnieSdpaSelfAttention,
+    "flash_attention_2": RNAErnieFlashSelfAttention,
+}
+
+
+class RNAErnieSelfOutput(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dropout(self.dense(hidden_states))
+        return self.LayerNorm(hidden_states + input_tensor)
+
+
+class RNAErnieAttention(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        attn_cls = RNAERNIE_SELF_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
+        self.self = attn_cls(config)
+        self.output = RNAErnieSelfOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        self_out, attn_weights = self.self(hidden_states, key_padding_mask, output_attentions)
+        return self.output(self_out, hidden_states), attn_weights
+
+
+class RNAErnieIntermediate(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.act = nn.ReLU() if config.hidden_act == "relu" else nn.GELU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.act(self.dense(hidden_states))
+
+
+class RNAErnieOutput(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dropout(self.dense(hidden_states))
+        return self.LayerNorm(hidden_states + input_tensor)
+
+
+class RNAErnieLayer(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.attention = RNAErnieAttention(config)
+        self.intermediate = RNAErnieIntermediate(config)
+        self.output = RNAErnieOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attn_out, attn_weights = self.attention(hidden_states, key_padding_mask, output_attentions)
+        return self.output(self.intermediate(attn_out), attn_out), attn_weights
+
+
+class RNAErnieEncoder(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.layer = nn.ModuleList([RNAErnieLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> Tuple:
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for layer in self.layer:
+            hidden_states, attn_weights = layer(hidden_states, key_padding_mask, output_attentions)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
+
+        return hidden_states, all_hidden_states, all_attentions
+
+
+class RNAErnieEmbeddings(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False)
+
+    def forward(self, input_ids: torch.LongTensor, token_type_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        B, T = input_ids.shape
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+        x = self.word_embeddings(input_ids)
+        x = x + self.position_embeddings(self.position_ids[:, :T])
+        x = x + self.token_type_embeddings(token_type_ids)
+        return self.dropout(self.LayerNorm(x))
+
+
+class RNAErniePooler(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.activation(self.dense(hidden_states[:, 0]))
+
+
+class RNAErniePredictionHeadTransform(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.act = nn.ReLU() if config.hidden_act == "relu" else nn.GELU()
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.LayerNorm(self.act(self.dense(hidden_states)))
+
+
+class RNAErnieLMPredictionHead(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.transform = RNAErniePredictionHeadTransform(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.transform(hidden_states))
+
+
+class RNAErnieOnlyMLMHead(nn.Module):
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__()
+        self.predictions = RNAErnieLMPredictionHead(config)
+
+    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+        return self.predictions(sequence_output)
+
+
+class RNAErnieModel(PreTrainedModel):
+    config_class = RNAErnieConfig
+    base_model_prefix = "bert"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__(config)
+        self.embeddings = RNAErnieEmbeddings(config)
+        self.encoder = RNAErnieEncoder(config)
+        self.pooler = RNAErniePooler(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        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 None:
+            attention_mask = torch.ones_like(input_ids)
+        key_padding_mask = attention_mask.eq(0)
+        if not key_padding_mask.any():
+            key_padding_mask = None
+
+        x = self.embeddings(input_ids, token_type_ids)
+        last_hidden_state, all_hidden_states, all_attentions = self.encoder(
+            x, key_padding_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        pooled = self.pooler(last_hidden_state)
+
+        if not return_dict:
+            return tuple(v for v in [last_hidden_state, pooled, all_hidden_states, all_attentions] if v is not None)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class RNAErnieForMaskedLM(PreTrainedModel):
+    config_class = RNAErnieConfig
+    base_model_prefix = "bert"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: RNAErnieConfig):
+        super().__init__(config)
+        self.bert = RNAErnieModel(config)
+        self.cls = RNAErnieOnlyMLMHead(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, MaskedLMOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
+            output_hidden_states=output_hidden_states, output_attentions=output_attentions,
+            return_dict=True,
+        )
+        logits = self.cls(outputs.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,) + outputs[2:]
+            return (loss,) + output if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss, logits=logits,
+            hidden_states=outputs.hidden_states, attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "mask_token": "[MASK]"
+}

tokenization_rnaernie.py

@@ -0,0 +1,137 @@
+import os
+from typing import Dict, List, Optional, Tuple
+
+from transformers import PreTrainedTokenizer
+
+_VOCAB = {
+    "[PAD]": 0,
+    "[UNK]": 1,
+    "[CLS]": 2,
+    "[SEP]": 3,
+    "[MASK]": 4,
+    "[DEL]": 5,
+    "[IND]": 6,
+    "RNaseMRPRNA": 7,
+    "RNasePRNA": 8,
+    "SRPRNA": 9,
+    "YRNA": 10,
+    "antisenseRNA": 11,
+    "autocatalyticallysplicedintron": 12,
+    "guideRNA": 13,
+    "hammerheadribozyme": 14,
+    "lncRNA": 15,
+    "miRNA": 16,
+    "miscRNA": 17,
+    "ncRNA": 18,
+    "other": 19,
+    "piRNA": 20,
+    "premiRNA": 21,
+    "precursorRNA": 22,
+    "rRNA": 23,
+    "ribozyme": 24,
+    "sRNA": 25,
+    "scRNA": 26,
+    "scaRNA": 27,
+    "siRNA": 28,
+    "snRNA": 29,
+    "snoRNA": 30,
+    "tRNA": 31,
+    "telomeraseRNA": 32,
+    "tmRNA": 33,
+    "vaultRNA": 34,
+    "A": 35,
+    "T": 36,
+    "C": 37,
+    "G": 38,
+}
+
+
+class RNAErnieTokenizer(PreTrainedTokenizer):
+    """Character-level RNA tokenizer for RNAErnie (original ERNIE/PaddlePaddle version).
+
+    Converts U to T before tokenisation (model was pretrained with DNA-style T).
+    Input sequences are uppercased and U->T substituted automatically.
+
+    Vocabulary (39 tokens):
+    - Special: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, [MASK]=4, [DEL]=5, [IND]=6
+    - ncRNA type labels: indices 7-34 (28 labels)
+    - Nucleotides: A=35, T=36, C=37, G=38
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        pad_token="[PAD]",
+        unk_token="[UNK]",
+        cls_token="[CLS]",
+        sep_token="[SEP]",
+        mask_token="[MASK]",
+        **kwargs,
+    ):
+        if vocab_file and os.path.isfile(vocab_file):
+            self._vocab = {}
+            with open(vocab_file, encoding="utf-8") as f:
+                for idx, line in enumerate(f):
+                    token = line.rstrip("\n")
+                    self._vocab[token] = idx
+        else:
+            self._vocab = dict(_VOCAB)
+        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
+        super().__init__(
+            pad_token=pad_token,
+            unk_token=unk_token,
+            cls_token=cls_token,
+            sep_token=sep_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]:
+        return list(text.upper().replace("U", "T"))
+
+    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) -> Tuple[str]:
+        os.makedirs(save_directory, exist_ok=True)
+        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.txt"
+        path = os.path.join(save_directory, fname)
+        with open(path, "w", encoding="utf-8") as f:
+            for token, _ in sorted(self._vocab.items(), key=lambda x: x[1]):
+                f.write(token + "\n")
+        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]
+        sep = [self.sep_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0 + sep
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    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, 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: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]:
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return [0] * len(cls + token_ids_0 + sep)
+        return [0] * len(cls + token_ids_0 + sep) + [1] * len(token_ids_1 + sep)

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_rnaernie.RNAErnieTokenizer", null]
+  },
+  "tokenizer_class": "RNAErnieTokenizer",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "unk_token": "[UNK]",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "padding_side": "right"
+}

vocab.txt

@@ -0,0 +1,39 @@
+[PAD]
+[UNK]
+[CLS]
+[SEP]
+[MASK]
+[DEL]
+[IND]
+RNaseMRPRNA
+RNasePRNA
+SRPRNA
+YRNA
+antisenseRNA
+autocatalyticallysplicedintron
+guideRNA
+hammerheadribozyme
+lncRNA
+miRNA
+miscRNA
+ncRNA
+other
+piRNA
+premiRNA
+precursorRNA
+rRNA
+ribozyme
+sRNA
+scRNA
+scaRNA
+siRNA
+snRNA
+snoRNA
+tRNA
+telomeraseRNA
+tmRNA
+vaultRNA
+A
+T
+C
+G