CodonBERT

Minimal HuggingFace port of the CodonBERT checkpoint from CodonBERT -- a BERT-based RNA language model pretrained on codon-level representations of more than 10 million mRNA coding sequences.

Architecture

Vocabulary

The tokenizer operates at the codon level. Unlike the original tokenizer, this port accepts raw nucleotide strings and performs codon splitting automatically. The 64 codons cover all combinations of {A, U, G, C}^3 in RNA space, including the three stop codons. Special tokens follow standard BERT convention: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, [MASK]=4.

Pretraining

• Objective: 15% masked language modeling (MLM) plus paired-sequence taxonomy prediction (STP)
• Data: >10 million mRNA coding sequences from mammals, bacteria, human viruses, and yeast
• Source checkpoint: pytorch_model.bin from the official Sanofi CodonBERT archive, mirrored as lhallee/CodonBERT

Checkpoint selection

There is a single publicly released checkpoint from the original authors. The backbone weights (bert.* prefix) and complete MLM prediction head are mapped directly. Only the two-tensor STP/sequence-relationship head is discarded.

Parity Verification

All verified on GPU with PyTorch 2.7.1 / CUDA 12.9:

• Hidden states (eager): all 13 levels match the original under torch.allclose(atol=1e-5, rtol=1e-5) (observed max abs diff 1.15e-5 on a padded three-sequence batch)
• MLM logits and loss: converted logits match original BertForPreTraining logits under the same tolerance (observed max abs diff 1.13e-5; mixed 0/-100 label loss difference 2.87e-6)
• SDPA (evaluation): final hidden states agree with eager FP32 to 3.58e-6 max abs difference at non-padding positions
• Flash attention 2 (evaluation): verified against eager BF16 at non-padding positions (all-layer differences up to 0.25, expected BF16 accumulation across 12 layers)

Related Models

See the full CodonBERT collection.

Usage

CodonBERT operates on CDS sequences. The tokenizer handles T->U conversion and codon splitting automatically — pass raw nucleotide strings directly.

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

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

# Raw CDS nucleotide strings — T or U both accepted
cds_sequences = ["ATGAAAGGCCCTTAA", "ATGTTTGGG"]

enc = tokenizer(cds_sequences, return_tensors="pt", padding=True)

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

cls_emb   = out.last_hidden_state[:, 0, :]  # (batch, 768) -- CLS token
content_mask = enc["attention_mask"].clone()
content_mask[:, 0] = 0  # exclude CLS
content_mask[torch.arange(content_mask.size(0)), enc["attention_mask"].sum(1) - 1] = 0
mean_emb = (out.last_hidden_state * content_mask.unsqueeze(-1)).sum(1) / \
           content_mask.sum(1, keepdim=True)  # mean over codons only

# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]  # (batch, seq_len, 768)

CDS-aware encoding (full mRNA input)

For full mRNA sequences where the CDS region must be extracted first:

import numpy as np

# cds: binary array with 1 at the first nucleotide of each codon
enc, chunk_counts = tokenizer.batch_encode_with_cds(
    mrna_sequences,
    cds_tracks,       # list of numpy arrays
    return_tensors="pt",
    padding=True,
)
with torch.no_grad():
    out = model(**enc)

Faster attention backends

# Evaluation/inference only; use eager for training (see Implementation Notes).
model_sdpa = AutoModel.from_pretrained(
    "Taykhoom/CodonBERT", trust_remote_code=True, attn_implementation="sdpa"
)
model_flash = AutoModel.from_pretrained(
    "Taykhoom/CodonBERT", trust_remote_code=True,
    attn_implementation="flash_attention_2", dtype=torch.bfloat16
)

MLM logits

from transformers import AutoModelForMaskedLM

model_mlm = AutoModelForMaskedLM.from_pretrained("Taykhoom/CodonBERT", trust_remote_code=True)
model_mlm.eval()

seq = "AUG [MASK] GGG"
enc = tokenizer(seq, return_tensors="pt")
with torch.no_grad():
    logits = model_mlm(**enc).logits  # (1, seq_len, 69)

The MLM prediction transform (dense + GELU + LayerNorm), decoder weight, and output bias are all converted from the original checkpoint. The decoder tensor is initialized exactly from the word embedding tensor. In this adapter it is a separate parameter, not a runtime-tied alias; see the fine-tuning limitations below.

Fine-tuning

For sequence-level tasks, use the CLS token embedding as input to a classification/regression head. Train this checkpoint with the eager backend: the current SDPA and Flash Attention 2 paths do not apply the configured 0.1 attention-probability dropout during training.

The input embedding and MLM decoder start value-identical but are not storage-tied. tie_weights() does not tie them, and resize_token_embeddings() leaves the decoder at 69 outputs; vocabulary resizing is therefore unsupported. The model forward API accepts input_ids, attention_mask, and token_type_ids, but not the stock BERT inputs_embeds, position_ids, or head_mask arguments.

Implementation Notes

Two key differences from the original CodonBERT release:

1. Integrated codon tokenization. The original repository requires users to manually pre-process sequences into space-separated codons before passing them to the tokenizer. This port ships CodonBertTokenizer, a BertTokenizer subclass whose _tokenize method automatically normalizes sequences (T->U, uppercase) and splits them into codon 3-mers. Users can pass raw nucleotide strings directly: tokenizer("AUGAAAGGG") works without any pre-processing. A batch_encode_with_cds(sequences, cds_tracks) method handles full mRNA input with CDS extraction and codon-boundary-aligned chunking.

2. SDPA and Flash Attention 2 support. This port inherits from Taykhoom/BERT-updated, a minimal BERT re-implementation with all three backends (eager, sdpa, flash_attention_2). Evaluation parity against the original eager implementation is verified at every layer. For training, use eager as noted above. Eager attentions requested in training mode are post-dropout tensors, so their rows do not sum to one. Although the shared config accepts hidden_act, the backend always applies GELU; this checkpoint is configured for GELU and is unaffected.

Citation

@article{li2024_codonbert,
  title   = {{CodonBERT} large language model for {mRNA} vaccines},
  author  = {Li, Sizhen and Moayedpour, Saeed and Li, Ruijiang and Bailey, Michael and Riahi, Saleh and Kogler-Anele, Lorenzo and Miladi, Milad and Miner, Jacob and Pertuy, Fabien and Zheng, Dinghai and Wang, Jun and Balsubramani, Akshay and Tran, Khang and Zacharia, Minnie and Wu, Monica and Gu, Xiaobo and Clinton, Ryan and Asquith, Carla and Skaleski, Joseph and Boeglin, Lianne and Chivukula, Sudha and Dias, Anusha and Strugnell, Tod and Ulloa Montoya, Fernando and Agarwal, Vikram and Bar-Joseph, Ziv and Jager, Sven},
  journal = {Genome Research},
  volume  = {34},
  number  = {7},
  pages   = {1027--1035},
  year    = {2024},
  doi     = {10.1101/gr.278870.123}
}

Credits

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

License

Academic/non-commercial use only, following the original artifact license:

Permission is hereby granted, free of charge, for academic research purposes only and for non-commercial use only, to any person from an academic research or non-profit organization obtaining a copy of these models, software, datasets and/or algorithms (including, but not limited to, machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements) and/or associated documentation files (collectively the "Materials") to use, copy, modify, or merge the Materials, subject to the following conditions: this IP License Notice shall be included in all copies of the Materials or of substantial portions of the Materials. For purposes of this notice, "non-commercial use" excludes uses foreseeably resulting in a commercial benefit or monetary gain. All other rights are reserved. The Materials are provided "as is," without warranty of any kind, express or implied, including the warranties of noninfringement.