RNABERT

Minimal HuggingFace port of the RNABERT variant of RNABERT -- a BERT-style RNA language model pretrained on non-coding RNA sequences from Rfam 14.3 for RNA clustering and structural alignment.

Architecture

Vocabulary: <pad> (0), <mask> (1), A (2), U (3), G (4), and C (5).

No CLS or EOS tokens are added. Sequences are tokenized character-by-character; T is silently converted to U.

Pretraining

• Objective: Masked Language Modeling (MLM) + Structural Alignment Learning (SAL, a pairwise structural alignment contrastive objective)
• Data: Full Rfam 14.3 dataset (~400 nt sequences, as described by the upstream repository)
• Source checkpoint: bert_mul_2.pth (distributed inside RNABERT_pretrained.pth zip, Google Drive)
• Checkpoint selection: The original repository publishes one pretrained checkpoint; this is it.

Parity Verification

All 7 backbone representation levels (embedding + 6 transformer blocks) matched the original bert_mul_2.pth weights (max abs diff = 3.19e-6, atol=1e-5), with and without padding, for eager and SDPA. Verified on GPU with PyTorch 2.7.1 / CUDA 12.9.

Related Models

See the full RNABERT collection.

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True)
model.eval()

sequences = ["AUGCAUGCAUGC", "GCUAGCUAGCUA"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

with torch.no_grad():
    out = model(**enc)

# Token-level embeddings
token_emb = out.last_hidden_state   # (batch, seq_len, 120)

# Mean-pool over non-padding positions
mask = enc["attention_mask"].unsqueeze(-1).float()
mean_emb = (token_emb * mask).sum(1) / mask.sum(1)  # (batch, 120)

# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer3_emb = out_all.hidden_states[3]   # (batch, seq_len, 120)

MLM logits

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True)
model.eval()

enc = tokenizer(["AUG<mask>AUG"], return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 6)

Faster attention backends

import torch
from transformers import AutoModel

# SDPA (PyTorch 2.0+) -- recommended for production
model = AutoModel.from_pretrained(
    "Taykhoom/RNABERT",
    trust_remote_code=True,
    attn_implementation="sdpa",
)

# Flash Attention 2 (requires flash-attn)
model = AutoModel.from_pretrained(
    "Taykhoom/RNABERT",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
    dtype=torch.float16,
)

Fine-tuning

The model has no CLS token, so use mean pooling over non-padding positions for sequence-level tasks.

import torch.nn as nn
from transformers import AutoModel

model = AutoModel.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True)

class RNAClassifier(nn.Module):
    def __init__(self, base, num_labels):
        super().__init__()
        self.base = base
        self.head = nn.Linear(120, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.base(input_ids, attention_mask=attention_mask)
        mask = attention_mask.unsqueeze(-1).float()
        pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1)
        return self.head(pooled)

Implementation Notes

The model config reuses the shared BERT-updated implementation (model_type: "bert_updated"), while tokenization remains RNABERT-specific. trust_remote_code=True is required for both the tokenizer and the model.

The original implementation uses standard scaled dot-product attention (post-LN BERT). This HF port adds attn_implementation="sdpa" and attn_implementation="flash_attention_2" support, which were not part of the original codebase.

Citation

@article{akiyama2022_rnabert,
  title   = {Informative {RNA} base embedding for {RNA} structural alignment and clustering by deep representation learning},
  author  = {Akiyama, Manato and Sakakibara, Yasubumi},
  journal = {NAR Genomics and Bioinformatics},
  volume  = {4},
  number  = {1},
  pages   = {lqac012},
  year    = {2022},
  doi     = {10.1093/nargab/lqac012}
}

Credits

Original model and code by Akiyama and Sakakibara. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

No license is specified in the original repository. Please contact the authors before redistributing or using in commercial settings.