Initial structRFM HF port

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,210 @@
+---
+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.

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "model_type": "bert_updated",
+  "auto_map": {
+    "AutoConfig": "configuration_bert_updated.BertUpdatedConfig",
+    "AutoModel": "modeling_bert.BertModel",
+    "AutoModelForMaskedLM": "modeling_bert.BertForMaskedLM"
+  },
+  "vocab_size": 10,
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 12,
+  "intermediate_size": 3072,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "attention_probs_dropout_prob": 0.1,
+  "max_position_embeddings": 514,
+  "type_vocab_size": 2,
+  "initializer_range": 0.02,
+  "layer_norm_eps": 1e-12,
+  "pad_token_id": 0,
+  "model_max_length": 514,
+  "transformers_version": "4.57.6",
+  "_name_or_path": "Taykhoom/structRFM"
+}

configuration_bert_updated.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig
+
+
+class BertUpdatedConfig(PretrainedConfig):
+    model_type = "bert_updated"
+
+    auto_map = {
+        "AutoConfig": "configuration_bert_updated.BertUpdatedConfig",
+        "AutoModel": "modeling_bert.BertModel",
+        "AutoModelForMaskedLM": "modeling_bert.BertForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        mlm_hidden_act=None,
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.mlm_hidden_act = mlm_hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d8e8dcc064b1b584b8ba6ee756cbfbc7953002e8ce580dfcfe85dc27a216a36
+size 344262720

modeling_bert.py

@@ -0,0 +1,547 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput
+
+from .configuration_bert_updated import BertUpdatedConfig
+
+
+class BertSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        return x.view(B, T, self.num_attention_heads, self.attention_head_size).permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        q = self._split_heads(self.query(hidden_states))
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+
+        scale = math.sqrt(self.attention_head_size)
+        scores = torch.matmul(q, k.transpose(-1, -2)) / scale
+        if key_padding_mask is not None:
+            scores = scores.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
+        probs = torch.nan_to_num(
+            F.softmax(scores.float(), dim=-1),
+            nan=0.0,
+        )
+        context_probs = self.dropout(probs)
+        if head_mask is not None:
+            context_probs = context_probs * head_mask
+        context = torch.matmul(context_probs.to(v.dtype), v)
+
+        B, _, T, _ = context.shape
+        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+
+        if output_attentions:
+            return context, probs
+        return context, None
+
+
+class BertSdpaSelfAttention(BertSelfAttention):
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions or head_mask is not None:
+            return super().forward(
+                hidden_states,
+                key_padding_mask,
+                output_attentions=output_attentions,
+                head_mask=head_mask,
+            )
+
+        B, T, _ = hidden_states.shape
+        q = self._split_heads(self.query(hidden_states))
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+
+        attn_mask = None
+        if key_padding_mask is not None:
+            attn_mask = torch.zeros(B, 1, 1, T, dtype=q.dtype, device=q.device)
+            attn_mask = attn_mask.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
+
+        context = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout.p if self.training else 0.0,
+        )
+        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+        return context, None
+
+
+class BertFlashSelfAttention(BertSelfAttention):
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if (
+            output_attentions
+            or head_mask is not None
+            or (
+                key_padding_mask is not None
+                and (~key_padding_mask).sum(dim=-1).eq(0).any()
+            )
+        ):
+            return super().forward(
+                hidden_states,
+                key_padding_mask,
+                output_attentions=output_attentions,
+                head_mask=head_mask,
+            )
+
+        try:
+            from flash_attn import flash_attn_func, flash_attn_varlen_func
+            from flash_attn.bert_padding import pad_input, unpad_input
+        except ImportError as e:
+            raise ImportError(
+                "flash_attn is required for attn_implementation='flash_attention_2'. "
+                "Install with: pip install flash-attn --no-build-isolation"
+            ) from e
+
+        B, T, _ = hidden_states.shape
+        q = self._split_heads(self.query(hidden_states)).permute(0, 2, 1, 3)
+        k = self._split_heads(self.key(hidden_states)).permute(0, 2, 1, 3)
+        v = self._split_heads(self.value(hidden_states)).permute(0, 2, 1, 3)
+
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                "flash_attention_2 requires float16 or bfloat16 model weights. "
+                f"Received {q.dtype}."
+            )
+
+        if key_padding_mask is not None and key_padding_mask.any():
+            attend = ~key_padding_mask
+            q_u, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, attend)
+            k_u, _, _, _, _ = unpad_input(k, attend)
+            v_u, _, _, _, _ = unpad_input(v, attend)
+            out_u = flash_attn_varlen_func(
+                q_u, k_u, v_u,
+                cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+                dropout_p=self.dropout.p if self.training else 0.0,
+                causal=False,
+            )
+            out = pad_input(out_u, indices, B, T)
+        else:
+            out = flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=self.dropout.p if self.training else 0.0,
+                causal=False,
+            )
+
+        out = out.reshape(B, T, self.all_head_size)
+        return out, None
+
+
+BERT_SELF_ATTENTION_CLASSES = {
+    "eager": BertSelfAttention,
+    "sdpa": BertSdpaSelfAttention,
+    "flash_attention_2": BertFlashSelfAttention,
+}
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dropout(self.dense(hidden_states))
+        return self.LayerNorm(hidden_states + input_tensor)
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        attn_cls = BERT_SELF_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
+        self.self = attn_cls(config)
+        self.output = BertSelfOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        self_out, attn_weights = self.self(
+            hidden_states,
+            key_padding_mask,
+            output_attentions,
+            head_mask,
+        )
+        return self.output(self_out, hidden_states), attn_weights
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if config.hidden_act in (None, "identity", "linear"):
+            self.intermediate_act_fn = lambda value: value
+        elif isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.intermediate_act_fn(self.dense(hidden_states))
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dropout(self.dense(hidden_states))
+        return self.LayerNorm(hidden_states + input_tensor)
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = BertAttention(config)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attn_out, attn_weights = self.attention(
+            hidden_states,
+            key_padding_mask,
+            output_attentions,
+            head_mask,
+        )
+        return self.output(self.intermediate(attn_out), attn_out), attn_weights
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        head_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple:
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for layer_index, layer in enumerate(self.layer):
+            layer_head_mask = (
+                head_mask[layer_index] if head_mask is not None else None
+            )
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attn_weights = torch.utils.checkpoint.checkpoint(
+                    layer.__call__,
+                    hidden_states,
+                    key_padding_mask,
+                    output_attentions,
+                    layer_head_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, attn_weights = layer(
+                    hidden_states,
+                    key_padding_mask,
+                    output_attentions,
+                    layer_head_mask,
+                )
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
+
+        return hidden_states, all_hidden_states, all_attentions
+
+
+class BertEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        B, T = (
+            input_ids.shape
+            if input_ids is not None
+            else inputs_embeds.shape[:2]
+        )
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(
+                (B, T),
+                dtype=torch.long,
+                device=(
+                    input_ids.device
+                    if input_ids is not None
+                    else inputs_embeds.device
+                ),
+            )
+        if position_ids is None:
+            position_ids = self.position_ids[:, :T]
+        x = (
+            self.word_embeddings(input_ids)
+            if inputs_embeds is None
+            else inputs_embeds
+        )
+        x = x + self.token_type_embeddings(token_type_ids)
+        x = x + self.position_embeddings(position_ids)
+        return self.dropout(self.LayerNorm(x))
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.activation(self.dense(hidden_states[:, 0]))
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        hidden_act = (
+            config.mlm_hidden_act
+            if getattr(config, "mlm_hidden_act", None) is not None
+            else config.hidden_act
+        )
+        if hidden_act in (None, "identity", "linear"):
+            self.transform_act_fn = lambda value: value
+        elif isinstance(hidden_act, str):
+            self.transform_act_fn = ACT2FN[hidden_act]
+        else:
+            self.transform_act_fn = hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.LayerNorm(self.transform_act_fn(self.dense(hidden_states)))
+
+
+class BertModel(PreTrainedModel):
+    config_class = BertUpdatedConfig
+    base_model_prefix = "bert"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_missing = [r"pooler\."]
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        input_shape = (
+            input_ids.shape
+            if input_ids is not None
+            else inputs_embeds.shape[:2]
+        )
+        input_device = (
+            input_ids.device
+            if input_ids is not None
+            else inputs_embeds.device
+        )
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=input_device)
+        key_padding_mask = attention_mask.eq(0)
+        if not key_padding_mask.any():
+            key_padding_mask = None
+        head_mask = self.get_head_mask(
+            head_mask,
+            self.config.num_hidden_layers,
+        )
+        if all(mask is None for mask in head_mask):
+            head_mask = None
+
+        x = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+        )
+        last_hidden_state, all_hidden_states, all_attentions = self.encoder(
+            x, key_padding_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            head_mask=head_mask,
+        )
+        pooled = self.pooler(last_hidden_state) if self.pooler is not None else None
+
+        if not return_dict:
+            output = (last_hidden_state, pooled)
+            if output_hidden_states:
+                output += (all_hidden_states,)
+            if output_attentions:
+                output += (all_attentions,)
+            return output
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class BertForMaskedLM(PreTrainedModel):
+    config_class = BertUpdatedConfig
+    base_model_prefix = "bert"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_unexpected = [r"bert\.pooler\."]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.transform = BertPredictionHeadTransform(config)
+        self.cls = nn.Linear(config.hidden_size, config.vocab_size)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.bert.embeddings.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.cls
+
+    def set_output_embeddings(self, value):
+        self.cls = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, MaskedLMOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        logits = self.cls(self.transform(outputs.last_hidden_state))
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)
+
+        if not return_dict:
+            output = (logits,)
+            if output_hidden_states:
+                output += (outputs.hidden_states,)
+            if output_attentions:
+                output += (outputs.attentions,)
+            return (loss,) + output if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss, logits=logits,
+            hidden_states=outputs.hidden_states, attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]"
+}

tokenization_structrfm.py

@@ -0,0 +1,104 @@
+import json
+import os
+
+from transformers import PreTrainedTokenizer
+
+_DEFAULT_VOCAB = {
+    "[UNK]": 0,
+    "[CLS]": 1,
+    "[SEP]": 2,
+    "[PAD]": 3,
+    "[MASK]": 4,
+    "A": 5,
+    "U": 6,
+    "C": 7,
+    "G": 8,
+    "N": 9,
+}
+
+
+class StructRFMTokenizer(PreTrainedTokenizer):
+    """Single-nucleotide tokenizer for structRFM.
+
+    Matches the upstream preprocessing (`seq.upper().replace("T", "U")`,
+    then one token per character) and adds [CLS] ... [SEP] when
+    `add_special_tokens=True`. Characters outside the vocabulary
+    (including whitespace) map to [UNK], as in the upstream tokenizer.
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        cls_token="[CLS]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        mask_token="[MASK]",
+        unk_token="[UNK]",
+        **kwargs,
+    ):
+        self._vocab = dict(_DEFAULT_VOCAB)
+        if vocab_file and os.path.isfile(vocab_file):
+            with open(vocab_file) as f:
+                self._vocab = json.load(f)
+        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
+        super().__init__(
+            cls_token=cls_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            unk_token=unk_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def _tokenize(self, text):
+        return list(text.upper().replace("T", "U"))
+
+    def _convert_token_to_id(self, token):
+        return self._vocab.get(token, self._vocab["[UNK]"])
+
+    def _convert_id_to_token(self, index):
+        return self._ids_to_tokens.get(index, "[UNK]")
+
+    def convert_tokens_to_string(self, tokens):
+        return "".join(tokens)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        path = os.path.join(save_directory, fname)
+        with open(path, "w") as f:
+            json.dump(self._vocab, f, indent=2)
+        return (path,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        cls = [self.cls_token_id]
+        sep = [self.sep_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0 + sep
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None,
+                                already_has_special_tokens=False):
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0, token_ids_1, already_has_special_tokens=True
+            )
+        mask = [1] + [0] * len(token_ids_0) + [1]
+        if token_ids_1 is not None:
+            mask += [0] * len(token_ids_1) + [1]
+        return mask
+
+    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
+        if token_ids_1 is None:
+            return [0] * (len(token_ids_0) + 2)
+        return [0] * (len(token_ids_0) + 2) + [1] * (len(token_ids_1) + 1)

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_structrfm.StructRFMTokenizer",
+      null
+    ]
+  },
+  "tokenizer_class": "StructRFMTokenizer",
+  "model_max_length": 514,
+  "padding_side": "right",
+  "truncation_side": "right",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]"
+}

vocab.json

@@ -0,0 +1,12 @@
+{
+  "[UNK]": 0,
+  "[CLS]": 1,
+  "[SEP]": 2,
+  "[PAD]": 3,
+  "[MASK]": 4,
+  "A": 5,
+  "U": 6,
+  "C": 7,
+  "G": 8,
+  "N": 9
+}