Taykhoom/SpliceBERT-510nt

SpliceBERT-510nt

Minimal HuggingFace port of the vertebrate-510nt variant of SpliceBERT -- a BERT-based RNA language model trained with masked language modeling on fixed-length 510 nt fragments from vertebrate primary RNA sequences.

WARNING: This model was trained on exactly 510 nt of input (excluding [CLS] and [SEP]). Sequences of other lengths were not validated upstream and may not work properly without fine-tuning. For general-purpose RNA embedding, use SpliceBERT-1024nt instead.

Architecture

Vocabulary: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, [MASK]=4, N=5, A=6, C=7, G=8, T=9. Input U is normalized to T.

Pretraining

• —Objective: Masked language modeling (MLM)
• —Data: >2 million vertebrate primary RNA sequences from 72 species
• —Sequence format: Single-nucleotide tokenization with spaces; U converted to T; fixed 510 nt fragments
• —Source checkpoint: SpliceBERT.510nt/pytorch_model.bin (from zenodo:7995778)

Checkpoint selection

The 510nt vertebrate variant is intended for splice site prediction tasks where exact 510 nt windows are used (e.g., centered on a splice site). For variable-length sequences use SpliceBERT-1024nt.

Parity Verification

Hidden-state representations verified (max abs diff < 1e-5) against the original checkpoint at all 7 representation levels (embedding + 6 transformer layers), for both eager and sdpa attention backends. Verified on GPU with PyTorch 2.7.1 / CUDA 12.9.

Related Models

See the full SpliceBERT collection.

Usage

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/SpliceBERT-510nt", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/SpliceBERT-510nt", trust_remote_code=True)
model.eval()

# The model was trained on exactly 510 nt; tokenizer handles U->T automatically
seq = ("ATCGATCG" * 64)[:510]  # exactly 510 nt
enc = tokenizer(seq, return_tensors="pt")

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

hidden = out.last_hidden_state[0]  # (512, 512)
token_emb = hidden[1:-1]           # strip [CLS] and [SEP] -> (510, 512)
mean_emb = token_emb.mean(dim=0)   # (512,)

Fine-tuning

Standard HF conventions. For splice site prediction, token-level classification using all 510 token positions (excluding special tokens) is the typical setup.

Implementation Notes

The original checkpoint was saved as BertForMaskedLM with transformers==4.20.1. This port uses BERT-updated, which adds attn_implementation="sdpa" and attn_implementation="flash_attention_2" support not present in the original codebase.

The pooler weights (pooler.dense) are not present in the original checkpoint and are not included in the saved model.safetensors. add_pooling_layer=True (the default) allocates the pooler layer but its weights are randomly initialized -- do not use pooler_output without fine-tuning.

Citation

@article{chen2024_splicebert,
  title   = {Self-supervised learning on millions of primary {RNA} sequences from 72 vertebrates improves sequence-based {RNA} splicing prediction},
  author  = {Chen, Ken and Zhou, Yue and Ding, Maolin and Wang, Yu and Ren, Zhixiang and Yang, Yuedong},
  journal = {Briefings in Bioinformatics},
  volume  = {25},
  number  = {3},
  pages   = {bbae163},
  year    = {2024},
  doi     = {10.1093/bib/bbae163}
}

Credits

Original model and code by Chen et al. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

The checkpoint weights are distributed under CC BY 4.0 by the upstream Zenodo record. The original SpliceBERT source code is BSD 3-Clause licensed.