README.md

---
language:
- rna
library_name: transformers
tags:
- RNA
- language-model
- bert
license: other
---

# RNABERT

A small BERT-style RNA language model pretrained on non-coding RNA sequences from Rfam 14.3, using
Masked Language Modeling (MLM) and Structural Alignment Learning (SAL). Designed for RNA clustering
and structural alignment tasks.

## Architecture

| Parameter | Value |
|---|---|
| Layers | 6 |
| Attention heads | 12 |
| Embedding dimension | 120 |
| FFN intermediate size | 40 |
| Vocabulary size | 6 (PAD, MASK, A, U, G, C) |
| Positional encoding | Learned absolute |
| Architecture | Post-LN BERT encoder |
| Max sequence length | 440 |

**Vocabulary:**

| Token | ID |
|---|---|
| `<pad>` | 0 |
| `<mask>` | 1 |
| A | 2 |
| U | 3 |
| G | 4 |
| 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:** Rfam 14.3 (~440 nt max length sequences)
- **Source checkpoint:** `bert_mul_2.pth` (distributed inside `RNABERT_pretrained.pth` zip,
  [Google Drive](https://drive.google.com/file/d/1sT6jlv9vrpX0npKmnbFeOqZ1JZDrZTQ2/view?usp=sharing))

### Checkpoint selection

There is one published pretrained checkpoint from the original repository. This is it.

## Parity Verification

Hidden-state representations verified identical (max abs diff = 3e-6) to the original
implementation at all 7 representation levels (embedding + 6 transformer layers), with and
without padding, for both eager and SDPA backends. Verified on GPU with PyTorch 2.7 /
transformers 4.57.6.

## Related Models

See the full [RNABERT collection](https://huggingface.co/collections/Taykhoom/rnabert-6a17cc9ca4852bd606ef4dba).

| Model | Notes |
|---|---|
| **[Taykhoom/RNABERT](https://huggingface.co/Taykhoom/RNABERT)** | This model |

## Usage

### Embedding generation

```python
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

```python
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)
```

### Fine-tuning

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

```python
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

This port uses a standalone `RNABertModel` (custom `PreTrainedModel` subclass, `model_type: "rnabert"`).
`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.

```python
# Faster inference with SDPA (default on modern PyTorch)
model = AutoModel.from_pretrained("Taykhoom/RNABERT", trust_remote_code=True,
                                   attn_implementation="sdpa")

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

## Citation

```bibtex
@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](https://github.com/mana438/RNABERT).
The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
and reviewed manually by Taykhoom Dalal.

## License

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