Initial structRFM HF port

933cd85 1 day ago

8.12 kB

---
library_name: transformers
tags:
- biology
- RNA
- language-model
- bert
license: mit
---
# structRFM

Minimal HuggingFace port of [structRFM](https://huggingface.co/heqin-zhu/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](https://huggingface.co/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](#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](https://github.com/heqin-zhu/structRFM/releases/tag/v0.0.8)
  (identical to [`heqin-zhu/structRFM`](https://huggingface.co/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](https://huggingface.co/collections/Taykhoom/structrfm-6aebd71d994e3b036c62ed23).

| Model | Parameters | Notes |
|---|---:|---|
| **[structRFM](https://huggingface.co/Taykhoom/structRFM)** | 86M | This model |

## Usage

### Embedding generation

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

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

```python
# 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](https://huggingface.co/Taykhoom/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

```bibtex
@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](https://github.com/heqin-zhu/structRFM),
[`heqin-zhu/structRFM` on the Hub](https://huggingface.co/heqin-zhu/structRFM).
Hugging Face port maintained by Taykhoom Dalal.

## License

MIT, following the original repository.

	---
	library_name: transformers
	tags:
	- biology
	- RNA
	- language-model
	- bert
	license: mit
	---

	# structRFM

	Minimal HuggingFace port of [structRFM](https://huggingface.co/heqin-zhu/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](https://huggingface.co/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](#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](https://github.com/heqin-zhu/structRFM/releases/tag/v0.0.8)
	(identical to [`heqin-zhu/structRFM`](https://huggingface.co/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](https://huggingface.co/collections/Taykhoom/structrfm-6aebd71d994e3b036c62ed23).

	\| Model \| Parameters \| Notes \|
	\|---\|---:\|---\|
	\| [structRFM](https://huggingface.co/Taykhoom/structRFM) \| 86M \| This model \|

	## Usage

	### Embedding generation

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

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

	```python
	# 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](https://huggingface.co/Taykhoom/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

	```bibtex
	@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](https://github.com/heqin-zhu/structRFM),
	[`heqin-zhu/structRFM` on the Hub](https://huggingface.co/heqin-zhu/structRFM).
	Hugging Face port maintained by Taykhoom Dalal.

	## License

	MIT, following the original repository.