Initial SpliceBERT-human-510nt Hugging Face 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,137 @@
+---
+library_name: transformers
+tags:
+- rna
+- language-model
+- splicing
+license: cc-by-4.0
+---
+
+# SpliceBERT-human-510nt
+
+Minimal HuggingFace port of the **human-510nt** variant of
+[SpliceBERT](https://github.com/biomed-AI/SpliceBERT) -- a BERT-based RNA
+language model trained with masked language modeling on fixed-length 510 nt
+fragments from human mRNA 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](https://huggingface.co/Taykhoom/SpliceBERT-1024nt) instead.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 6 |
+| Attention heads | 16 |
+| Embedding dimension | 512 |
+| FFN hidden dimension | 2048 (GELU) |
+| Vocabulary size | 10 |
+| Positional encoding | Learned absolute |
+| Normalization | LayerNorm (post-residual, eps=1e-12) |
+| Architecture | Post-LN BERT encoder |
+| Max sequence length | 510 nt (512 tokens; fixed-length training) |
+| Checkpoint size | ~19.5M parameters |
+
+**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:** Human primary RNA sequences
+- **Sequence format:** Single-nucleotide tokenization with spaces; U converted to T; fixed 510 nt fragments
+- **Source checkpoint:** `SpliceBERT-human.510nt/pytorch_model.bin` (from [zenodo:7995778](https://doi.org/10.5281/zenodo.7995778))
+
+### Checkpoint selection
+
+This human-only variant may outperform the multi-species 510nt model on human-specific
+splicing tasks. For cross-species generalization or variable-length sequences, use
+[SpliceBERT-1024nt](https://huggingface.co/Taykhoom/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](https://huggingface.co/collections/Taykhoom/splicebert-6a20b72e9bec05b79ce009aa).
+
+| Model | Context | Training data | Notes |
+|---|---|---|---|
+| [SpliceBERT-1024nt](https://huggingface.co/Taykhoom/SpliceBERT-1024nt) | 1024 nt | 72 vertebrates | Variable-length; general purpose |
+| [SpliceBERT-510nt](https://huggingface.co/Taykhoom/SpliceBERT-510nt) | 510 nt (fixed) | 72 vertebrates | Multi-species 510 nt |
+| **[SpliceBERT-human-510nt](https://huggingface.co/Taykhoom/SpliceBERT-human-510nt)** | 510 nt (fixed) | Human only | This model |
+
+## Usage
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/SpliceBERT-human-510nt", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/SpliceBERT-human-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.18.0`.
+This port uses [BERT-updated](https://huggingface.co/Taykhoom/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
+
+```bibtex
+@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](https://github.com/biomed-AI/SpliceBERT).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+The checkpoint weights are distributed under
+[CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) by the upstream
+[Zenodo record](https://doi.org/10.5281/zenodo.7995778). The original
+SpliceBERT source code is BSD 3-Clause licensed.

config.json

@@ -0,0 +1,27 @@
+{
+  "_name_or_path": "Taykhoom/SpliceBERT-human-510nt",
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "model_type": "bert_updated",
+  "auto_map": {
+    "AutoConfig": "Taykhoom/BERT-updated--configuration_bert_updated.BertUpdatedConfig",
+    "AutoModel": "Taykhoom/BERT-updated--modeling_bert.BertModel",
+    "AutoModelForMaskedLM": "Taykhoom/BERT-updated--modeling_bert.BertForMaskedLM"
+  },
+  "vocab_size": 10,
+  "hidden_size": 512,
+  "num_hidden_layers": 6,
+  "num_attention_heads": 16,
+  "intermediate_size": 2048,
+  "hidden_act": "gelu",
+  "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,
+  "pad_token_id": 0,
+  "model_max_length": 510,
+  "transformers_version": "4.57.6"
+}

configuration_bert_updated.py

@@ -0,0 +1,41 @@
+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",
+        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.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:6e5465d8e1c0e827d36b88b9fd9ad4f84397760adc90d924391c2442da2f7b30
+size 77826008

modeling_bert.py

@@ -0,0 +1,372 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+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,
+    ) -> 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 = F.softmax(scores, dim=-1)
+        probs = self.dropout(probs)
+        context = torch.matmul(probs, 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,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(hidden_states, key_padding_mask, output_attentions=True)
+
+        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)
+        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,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(hidden_states, key_padding_mask, output_attentions=True)
+
+        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)
+
+        orig_dtype = q.dtype
+        if orig_dtype not in (torch.float16, torch.bfloat16):
+            q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
+
+        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,
+                causal=False,
+            )
+            out = pad_input(out_u, indices, B, T)
+        else:
+            out = flash_attn_func(q, k, v, causal=False)
+
+        out = out.to(orig_dtype).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,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        self_out, attn_weights = self.self(hidden_states, key_padding_mask, output_attentions)
+        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)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.gelu(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,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attn_out, attn_weights = self.attention(hidden_states, key_padding_mask, output_attentions)
+        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)])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor],
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> Tuple:
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for layer in self.layer:
+            hidden_states, attn_weights = layer(hidden_states, key_padding_mask, output_attentions)
+            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: torch.LongTensor, token_type_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        B, T = input_ids.shape
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+        x = self.word_embeddings(input_ids)
+        x = x + self.position_embeddings(self.position_ids[:, :T])
+        x = x + self.token_type_embeddings(token_type_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)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.LayerNorm(F.gelu(self.dense(hidden_states)))
+
+
+class BertModel(PreTrainedModel):
+    config_class = BertUpdatedConfig
+    base_model_prefix = "bert"
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.pooler = BertPooler(config)
+        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: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = 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 attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        key_padding_mask = attention_mask.eq(0)
+        if not key_padding_mask.any():
+            key_padding_mask = None
+
+        x = self.embeddings(input_ids, token_type_ids)
+        last_hidden_state, all_hidden_states, all_attentions = self.encoder(
+            x, key_padding_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        pooled = self.pooler(last_hidden_state)
+
+        if not return_dict:
+            return tuple(v for v in [last_hidden_state, pooled, all_hidden_states, all_attentions] if v is not None)
+
+        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
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config)
+        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 forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = 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, attention_mask=attention_mask, token_type_ids=token_type_ids,
+            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,) + outputs[2:]
+            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_splicebert.py

@@ -0,0 +1,98 @@
+import json
+import os
+from transformers import PreTrainedTokenizer
+
+_DEFAULT_VOCAB = {
+    "[PAD]": 0,
+    "[UNK]": 1,
+    "[CLS]": 2,
+    "[SEP]": 3,
+    "[MASK]": 4,
+    "N": 5,
+    "A": 6,
+    "C": 7,
+    "G": 8,
+    "T": 9,
+}
+
+
+class SpliceBERTTokenizer(PreTrainedTokenizer):
+    """Single-nucleotide tokenizer for SpliceBERT.
+
+    Automatically converts U->T and adds [CLS]/[SEP] special tokens.
+    Raw sequences (not pre-spaced) are accepted.
+    """
+
+    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("U", "T").replace(" ", ""))
+
+    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 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 + cls + 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 += [1] + [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] + token_ids_0 + [0]
+        return [0] + token_ids_0 + [0, 0] + token_ids_1 + [0]

tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_splicebert.SpliceBERTTokenizer",
+      null
+    ]
+  },
+  "model_max_length": 510,
+  "tokenizer_class": "SpliceBERTTokenizer",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "eos_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]"
+}

vocab.json

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