structRFM

Minimal HuggingFace port of structRFM -- a structure-guided RNA foundation model: a BERT encoder pretrained with structure-guided masked language modeling on ~21M RNAcentral sequences paired with ensemble-predicted secondary structures.

Why this port?

heqin-zhu/structRFM ships the weights as a stock BertForMaskedLM, but its tokenizer has no post-processing template or normalizer. Upstream code builds the model input by hand ("[CLS]" + seq.upper().replace("T", "U") + "[SEP]"); calling its AutoTokenizer directly on a sequence silently:

1. adds no [CLS] / [SEP] tokens, and
2. maps T and every lowercase base to [UNK].

This repo's tokenizer applies the upstream preprocessing itself (uppercase, T->U, [CLS] ... [SEP]), so tokenizer(seq) produces exactly the token ids the model was pretrained on. It also adds flash_attention_2 support. The weights are unchanged and bit-exact with the upstream release (see Parity Verification).

Architecture

Parameter	Value
Layers	12
Attention heads	12
Embedding dimension	768
FFN hidden dimension	3072 (GELU)
Vocabulary size	10
Positional encoding	Learned absolute
Normalization	LayerNorm (post-residual, eps=1e-12)
Architecture	Post-LN BERT encoder
Max sequence length	512 nt (514 tokens with [CLS] and [SEP])
Parameters	85.5M (encoder) + 0.6M (MLM head)

Vocabulary: [UNK]=0, [CLS]=1, [SEP]=2, [PAD]=3, [MASK]=4, A=5, U=6, C=7, G=8, N=9. Input is uppercased and T is read as U; any other character (including whitespace) maps to [UNK], as upstream.

Pretraining

• Objective: Structure-guided masked language modeling (SgMLM): each batch uses either standard 15% token masking or structure masking, where masking a paired base also masks its partner; the share of structure-masked batches is scheduled over training
• Data: ~21M RNAcentral sequences (<= 512 nt) with secondary structures from MUSES, a multi-source ensemble of thermodynamic, probabilistic and deep-learning predictors
• Source checkpoint: structRFM_checkpoint.tar.gz from the upstream v0.0.8 release (identical to heqin-zhu/structRFM)

Checkpoint selection

Upstream released two pretrained models: v0.0.3, pretrained on single-predictor (BPfold) structures, and v0.0.8, re-pretrained on MUSES ensemble structures. v0.0.8 replaced v0.0.3 on the Hub and is the base of every released downstream checkpoint, so it is the one ported here.

The released secondary-structure (structRFM_SSP_bpRNA1m.pt, structRFM_SSP_RNAStrAlign.pt), tertiary-structure and IRES checkpoints train a task head on detached encoder features, so their encoder weights are bit-identical to this model; use this repo for their embeddings.

Parity Verification

All 13 representation levels (embedding + 12 transformer blocks) and the MLM logits verified to be bit-exact (max abs diff = 0.00) against the original v0.0.8 weights run through the upstream structRFM code, with eager and sdpa attention, on padded and unpadded batches. The tokenizer reproduces the upstream token ids exactly. Verified on an NVIDIA H100 with PyTorch 2.7 / CUDA 12.9.

Related Models

See the full structRFM collection.

Model	Parameters	Notes
structRFM	86M	This model

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

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

sequences = ["GUCCCAACUCUUGCGGGGAGGGAU", "ATGGCTACGTAG"]   # T is read as U
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)

# Mean over nucleotides only (drop [CLS], [SEP] and padding)
special = torch.isin(enc.input_ids, torch.tensor([tokenizer.cls_token_id, tokenizer.sep_token_id]))
mask = (enc.attention_mask.bool() & ~special).unsqueeze(-1)
mean_emb = (token_emb * mask).sum(1) / mask.sum(1)   # (batch, 768)

# Intermediate layers: hidden_states[0] is the embedding, [i] block i
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]

Sequences longer than 512 nt must be split into windows (e.g. 512 nt chunks) and the per-window embeddings pooled.

MLM logits

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

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

enc = tokenizer(["GUCCCAAC[MASK]CUUGCGG"], return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 10)

Faster attention backends

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

# Flash Attention 2 (requires flash-attn) -- fastest on long sequences
model = AutoModel.from_pretrained("Taykhoom/structRFM", trust_remote_code=True,
                                  attn_implementation="flash_attention_2",
                                  dtype=torch.bfloat16)

Backend	Status
eager	Supported; bit-exact with upstream eager
sdpa	Supported; bit-exact with upstream (the upstream default backend)
flash_attention_2	Supported (added by this port); requires fp16/bf16 weights

output_attentions=True falls back to eager under sdpa and flash_attention_2 and returns post-softmax attention probabilities (batch, heads, seq_len, seq_len).

Fine-tuning

Standard HF conventions. For sequence-level tasks, use the CLS embedding or a masked mean over nucleotide positions as input to a prediction head; for nucleotide- or pair-level tasks use last_hidden_state directly. gradient_checkpointing_enable() is supported.

Implementation Notes

This port reuses the shared BERT-updated runtime (identical to the stock transformers BERT architecture), which adds flash_attention_2 alongside eager and sdpa. The MLM decoder is tied to the input embeddings, as in the upstream checkpoint.

config.pad_token_id is 0 ([UNK]), the value in the upstream config, while the tokenizer pads with [PAD]=3. The config value only sets the embedding's padding_idx (no gradient into that row during fine-tuning) and is kept for fidelity; padding is controlled by attention_mask.

The pooler (pooler_output) is not part of the pretrained checkpoint and is randomly initialized; do not use it without fine-tuning.

Citation

@article{zhu2026_structrfm,
  title   = {An open structure-guided {RNA} foundation model for generalizable prediction},
  author  = {Zhu, Heqin and Li, Ruifeng and Chang, Ao and Chen, Haobin and Zhang, Feng and Tang, Fenghe and Ye, Tong and Li, Xin and Gu, Yunjie and Xiong, Peng and Zhou, S. Kevin},
  journal = {Nature},
  year    = {2026},
  doi     = {10.1038/s41586-026-11210-2}
}

Credits

Original model and code by Zhu et al. Source: GitHub, heqin-zhu/structRFM on the Hub. Hugging Face port maintained by Taykhoom Dalal.

License

MIT, following the original repository.