Initial GENA-LM 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,156 @@
+---
+library_name: transformers
+tags:
+- biology
+- DNA
+- language-model
+license: mit
+---
+
+# GENA-LM-bert-base
+
+Minimal HuggingFace port of the **bert-base** variant of
+[GENA-LM](https://huggingface.co/AIRI-Institute/gena-lm-bert-base) -- a transformer masked language
+model for long human / multi-species DNA sequences, using byte-pair (BPE) tokenization.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 12 |
+| Attention heads | 12 |
+| Embedding dimension | 768 |
+| FFN hidden dimension | 3072 (GELU) |
+| Vocabulary size | 32000 |
+| Positional encoding | learned absolute |
+| Normalization | Pre-LayerNorm (eps=1e-12); final-layer LayerNorm: No |
+| Architecture | Pre-LayerNorm BERT (without a final-layer LayerNorm) |
+| Max sequence length | 512 BPE tokens (~4608 nucleotides) |
+
+**Vocabulary:** 32,000 BPE tokens trained on DNA, including `[CLS]`, `[SEP]`, `[PAD]`,
+`[UNK]`, and `[MASK]`.
+
+## Pretraining
+
+- **Objective:** Masked language modeling (15% masking, BigBird-style).
+- **Data:** Human T2T genome assembly (GCA_009914755.3).
+- **Pretraining iterations:** 500,000 (batch size 256, sequence length 512).
+- **Source checkpoint:** `AIRI-Institute/gena-lm-bert-base`
+
+## Parity Verification
+
+All 13 representation levels (embedding + 12 transformer blocks) and the
+masked-LM logits were verified to be bit-exact (max abs diff = 0.00) against the
+original `AIRI-Institute/gena-lm-bert-base` weights, for the `eager` backend. The added `sdpa` and
+`flash_attention_2` backends agree with `eager` up to the expected fused-kernel
+floating-point tolerance. Verified on GPU with PyTorch 2.7 / CUDA 12.9.
+
+## Related Models
+
+See the full [GENA-LM collection](https://huggingface.co/collections/Taykhoom/gena-lm-6a8cec0862e11d4f81d059ab).
+
+| Model | Parameters | Notes |
+|---|---|---|
+| **[GENA-LM-bert-base](https://huggingface.co/Taykhoom/GENA-LM-bert-base)** | 110M | 12L / 768d, 512 ctx (this model) |
+| [GENA-LM-t2t-bert-base](https://huggingface.co/Taykhoom/GENA-LM-t2t-bert-base) | 110M | 12L / 768d, 512 ctx |
+| [GENA-LM-t2t-multi-species-bert-base](https://huggingface.co/Taykhoom/GENA-LM-t2t-multi-species-bert-base) | 110M | 12L / 768d, 512 ctx |
+| [GENA-LM-t2t-lastln-base](https://huggingface.co/Taykhoom/GENA-LM-t2t-lastln-base) | 110M | 12L / 768d, 512 ctx |
+| [GENA-LM-t2t-bert-large](https://huggingface.co/Taykhoom/GENA-LM-t2t-bert-large) | 336M | 24L / 1024d, 512 ctx |
+| [GENA-LM-t2t-bigbird-base](https://huggingface.co/Taykhoom/GENA-LM-t2t-bigbird-base) | 110M | 12L / 768d, 4096 ctx |
+| [GENA-LM-t2t-sparse-bigbird-base](https://huggingface.co/Taykhoom/GENA-LM-t2t-sparse-bigbird-base) | 110M | 12L / 768d, 4096 ctx |
+| [GENA-LM-sparse-bigbird-base](https://huggingface.co/Taykhoom/GENA-LM-sparse-bigbird-base) | 110M | 12L / 768d, 4096 ctx |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
+model.eval()
+
+sequences = ["ACGTACGTACGTACGT", "TTACGGGCATACGACGT"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, dim) -- CLS token
+token_emb = out.last_hidden_state             # (batch, seq_len, dim)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer6_emb = out_all.hidden_states[6]
+```
+
+### MLM logits
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["ACGT[MASK]CGTACGT"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, seq_len, vocab_size)
+```
+
+### Faster attention backends
+
+```python
+# SDPA (PyTorch 2.0+) -- recommended for production
+model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True,
+                                  attn_implementation="sdpa")
+
+# Flash Attention 2 (requires flash-attn) -- fastest on long sequences
+import torch
+model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True,
+                                  attn_implementation="flash_attention_2",
+                                  dtype=torch.bfloat16)
+```
+
+### Fine-tuning
+
+Standard HuggingFace conventions. For sequence-level tasks, pool over non-padding
+positions or use the `[CLS]` token embedding as input to a prediction head.
+
+## Implementation Notes
+
+This is a minimal, self-contained reimplementation of the GENA-LM pre-LayerNorm BERT
+backbone. The only behavioural additions over the original (eager-only) code are the
+`sdpa` and `flash_attention_2` attention backends, selectable via `attn_implementation`;
+the `eager` backend reproduces the original outputs bit-for-bit. The original NSP head
+and pooler are not included, since this port targets embedding and masked-LM use.
+`AutoModel` returns the backbone without a pooler; use the `[CLS]` hidden state or
+masked mean pooling for sequence embeddings. The input embeddings and MLM decoder are tied.
+
+## Citation
+
+```bibtex
+@article{fishman2025_genalm,
+  title   = {{GENA-LM}: a family of open-source foundational {DNA} language models for long sequences},
+  author  = {Fishman, Veniamin and Kuratov, Yuri and Shmelev, Aleksei and Petrov, Maxim and Penzar, Dmitry and Shepelin, Denis and Chekanov, Nikolay and Kardymon, Olga and Burtsev, Mikhail},
+  journal = {Nucleic Acids Research},
+  volume  = {53},
+  number  = {2},
+  pages   = {gkae1310},
+  year    = {2025},
+  doi     = {10.1093/nar/gkae1310}
+}
+```
+
+## Credits
+
+Original model and code by Fishman, Kuratov, et al. (AIRI Institute).
+Source: [GitHub](https://github.com/AIRI-Institute/GENA_LM).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+MIT, following the original repository.

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "GenaLMForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "last_layer_norm": false,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_max_length": 512,
+  "model_type": "genalm",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 3,
+  "position_embedding_type": "absolute",
+  "pre_layer_norm": true,
+  "rotary_base": 10000,
+  "rotary_dim": 32,
+  "sparse_block_size": 0,
+  "transformers_version": "4.57.6",
+  "type_vocab_size": 2,
+  "vocab_size": 32000,
+  "auto_map": {
+    "AutoConfig": "configuration_genalm.GenaLMConfig",
+    "AutoModel": "modeling_genalm.GenaLMModel",
+    "AutoModelForMaskedLM": "modeling_genalm.GenaLMForMaskedLM"
+  },
+  "tie_word_embeddings": true
+}

configuration_genalm.py

@@ -0,0 +1,73 @@
+"""Configuration for the minimal GENA-LM (pre-LayerNorm BERT) port."""
+
+from transformers import PretrainedConfig
+
+
+class GenaLMConfig(PretrainedConfig):
+    """Configuration for the GENA-LM backbone.
+
+    GENA-LM is a pre-LayerNorm BERT (x = x + mha(ln(x)); x = x + ffn(ln(x))) with an
+    optional final LayerNorm. It supports learned absolute position embeddings and,
+    for some checkpoints, rotary position embeddings.
+    """
+
+    model_type = "genalm"
+
+    auto_map = {
+        "AutoConfig": "configuration_genalm.GenaLMConfig",
+        "AutoModel": "modeling_genalm.GenaLMModel",
+        "AutoModelForMaskedLM": "modeling_genalm.GenaLMForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        hidden_act: str = "gelu",
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 512,
+        type_vocab_size: int = 2,
+        initializer_range: float = 0.02,
+        layer_norm_eps: float = 1e-12,
+        pad_token_id: int = 3,
+        position_embedding_type: str = "absolute",
+        pre_layer_norm: bool = True,
+        last_layer_norm: bool = False,
+        rotary_base: int = 10000,
+        rotary_dim: int = 32,
+        sparse_block_size: int = 0,
+        model_max_length: int = 512,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **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
+        self.position_embedding_type = position_embedding_type
+        self.pre_layer_norm = pre_layer_norm
+        self.last_layer_norm = last_layer_norm
+        self.rotary_base = rotary_base
+        self.rotary_dim = rotary_dim
+        # BigBird block-sparse attention: 0 disables (dense). When > 0, attention is
+        # restricted to the (per-head) block layout stored in each layer's master_layout
+        # buffer, reproducing the original DeepSpeed block-sparse attention pattern.
+        self.sparse_block_size = sparse_block_size
+        self.model_max_length = model_max_length

model.safetensors

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

modeling_genalm.py

@@ -0,0 +1,635 @@
+"""Minimal GENA-LM backbone: pre-LayerNorm BERT with eager / sdpa / flash_attention_2.
+
+GENA-LM (AIRI-Institute) is a pre-LayerNorm BERT trained on human / multi-species DNA.
+The residual structure is::
+
+    a = pre_attention_ln(x)
+    x = x + attn_out_dense(self_attention(a))
+    b = post_attention_ln(x)
+    x = x + ffn_out_dense(act(ffn_in_dense(b)))
+
+and, for checkpoints with ``last_layer_norm=True``, a final LayerNorm is applied after
+the last block. Both learned absolute and rotary position embeddings are supported.
+
+This file adds ``sdpa`` and ``flash_attention_2`` dispatch on top of the original eager
+attention; the upstream model only ships eager attention.
+"""
+
+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
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput
+
+from .configuration_genalm import GenaLMConfig
+
+
+# ---------------------------------------------------------------------------
+# Rotary position embeddings (used by the rotary checkpoints)
+# ---------------------------------------------------------------------------
+
+class GenaLMRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, base: int = 10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        # persistent so the checkpoint's (fp16-rounded) inv_freq is loaded exactly rather
+        # than recomputed in fp32, which keeps rotary outputs bit-identical to the original.
+        self.register_buffer("inv_freq", inv_freq, persistent=True)
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def forward(self, seq_len: int, device, dtype) -> Tuple[torch.Tensor, torch.Tensor]:
+        if seq_len != self._seq_len_cached or self._cos_cached is None or self._cos_cached.device != device:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self._cos_cached = emb.cos()[None, None, :, :]
+            self._sin_cached = emb.sin()[None, None, :, :]
+        # A cache populated under torch.inference_mode() contains inference
+        # tensors, which cannot later participate in a grad-enabled forward.
+        # Clone in the caller's current mode so logits()/embed() can be mixed.
+        return self._cos_cached.to(dtype).clone(), self._sin_cached.to(dtype).clone()
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    cos = cos[:, :, : q.shape[2], :].to(q.dtype)
+    sin = sin[:, :, : q.shape[2], :].to(q.dtype)
+    return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)
+
+
+# ---------------------------------------------------------------------------
+# Attention variants
+# ---------------------------------------------------------------------------
+
+class GenaLMSelfAttention(nn.Module):
+    """Eager scaled dot-product self-attention (matches the original)."""
+
+    def __init__(self, config: GenaLMConfig):
+        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)
+
+        self.position_embedding_type = config.position_embedding_type
+        if self.position_embedding_type == "rotary":
+            self.rotary_dim = config.rotary_dim
+            self.rotary_emb = GenaLMRotaryEmbedding(self.rotary_dim, base=config.rotary_base)
+
+    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 _maybe_rotary(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.position_embedding_type != "rotary":
+            return q, k
+        cos, sin = self.rotary_emb(q.shape[2], q.device, q.dtype)
+        if self.rotary_dim < self.attention_head_size:
+            q_rot, q_pass = q[..., : self.rotary_dim], q[..., self.rotary_dim:]
+            k_rot, k_pass = k[..., : self.rotary_dim], k[..., self.rotary_dim:]
+            q_rot, k_rot = _apply_rotary(q_rot, k_rot, cos, sin)
+            q = torch.cat((q_rot, q_pass), dim=-1)
+            k = torch.cat((k_rot, k_pass), dim=-1)
+        else:
+            q, k = _apply_rotary(q, k, cos, sin)
+        return q, k
+
+    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))
+        q, k = self._maybe_rotary(q, k)
+
+        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.attention_head_size)
+        if key_padding_mask is not None:
+            scores = scores.masked_fill(key_padding_mask[:, None, None, :], float("-inf"))
+        probs = F.softmax(scores, dim=-1)
+        attn_weights = probs if output_attentions else None
+        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)
+        return context, attn_weights
+
+
+class GenaLMSdpaSelfAttention(GenaLMSelfAttention):
+    """SDPA self-attention via torch.nn.functional.scaled_dot_product_attention."""
+
+    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))
+        q, k = self._maybe_rotary(q, k)
+
+        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 GenaLMFlashSelfAttention(GenaLMSelfAttention):
+    """Flash Attention 2 self-attention."""
+
+    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))  # (B, H, T, D)
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+        q, k = self._maybe_rotary(q, k)
+
+        # flash_attn expects (B, T, H, D)
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.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  # True = valid token
+            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
+
+
+GENALM_SELF_ATTENTION_CLASSES = {
+    "eager": GenaLMSelfAttention,
+    "sdpa": GenaLMSdpaSelfAttention,
+    "flash_attention_2": GenaLMFlashSelfAttention,
+}
+
+
+class GenaLMBlockSparseSelfAttention(GenaLMSelfAttention):
+    """BigBird block-sparse self-attention, computed as masked dense attention.
+
+    Block-sparse attention is, by definition, dense attention restricted to a set of
+    allowed (query-block, key-block) pairs, with the softmax taken over the allowed keys
+    only. Each layer's allowed pattern is stored in ``master_layout`` (num_heads,
+    max_blocks, max_blocks), loaded from the original checkpoint; for a sequence of
+    ``nb = ceil(T / block)`` blocks we use ``master_layout[:, :nb, :nb]`` exactly as the
+    original DeepSpeed implementation does. This reproduces the original attention pattern
+    without requiring DeepSpeed's (triton-1.x-only) sparse kernels.
+
+    ``use_sdpa=True`` evaluates the masked attention with
+    ``F.scaled_dot_product_attention``; ``output_attentions`` always uses the explicit
+    (eager) path so the post-softmax probabilities can be returned.
+    """
+
+    def __init__(self, config: GenaLMConfig, use_sdpa: bool = False):
+        super().__init__(config)
+        self.block = config.sparse_block_size
+        self.use_sdpa = use_sdpa
+        max_blocks = config.max_position_embeddings // self.block
+        self.register_buffer(
+            "master_layout",
+            torch.zeros(self.num_attention_heads, max_blocks, max_blocks, dtype=torch.int64),
+            persistent=True,
+        )
+
+    def _block_mask(self, nb: int, device) -> torch.Tensor:
+        """(num_heads, nb*block, nb*block) bool, True where attention is allowed."""
+        layout = self.master_layout[:, :nb, :nb].to(device).bool()
+        return layout.repeat_interleave(self.block, dim=1).repeat_interleave(self.block, dim=2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        B, T, _ = hidden_states.shape
+        q = self._split_heads(self.query(hidden_states))  # (B, H, T, D)
+        k = self._split_heads(self.key(hidden_states))
+        v = self._split_heads(self.value(hidden_states))
+        q, k = self._maybe_rotary(q, k)
+
+        # pad sequence up to a multiple of the block size
+        pad = (-T) % self.block
+        Tp = T + pad
+        if pad:
+            q = F.pad(q, (0, 0, 0, pad))
+            k = F.pad(k, (0, 0, 0, pad))
+            v = F.pad(v, (0, 0, 0, pad))
+        nb = Tp // self.block
+
+        # key positions that must never be attended to: original padding + block padding
+        key_pad = torch.zeros(B, Tp, dtype=torch.bool, device=q.device)
+        if key_padding_mask is not None:
+            key_pad[:, :T] = key_padding_mask
+        if pad:
+            key_pad[:, T:] = True
+
+        allowed = self._block_mask(nb, q.device)  # (H, Tp, Tp)
+        disallowed = (~allowed)[None]  # (1, H, Tp, Tp)
+        key_pad_b = key_pad[:, None, None, :]  # (B, 1, 1, Tp)
+
+        if self.use_sdpa and not output_attentions:
+            attn_mask = torch.zeros(B, self.num_attention_heads, Tp, Tp, dtype=q.dtype, device=q.device)
+            attn_mask = attn_mask.masked_fill(disallowed, float("-inf"))
+            attn_mask = attn_mask.masked_fill(key_pad_b, float("-inf"))
+            context = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+            attn_weights = None
+        else:
+            scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.attention_head_size)
+            scores = scores.masked_fill(disallowed, float("-inf"))
+            scores = scores.masked_fill(key_pad_b, float("-inf"))
+            probs = F.softmax(scores, dim=-1)
+            attn_weights = probs[:, :, :T, :T] if output_attentions else None
+            probs = self.dropout(probs)
+            context = torch.matmul(probs, v)
+
+        context = context[:, :, :T, :].permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+        return context, attn_weights
+
+
+# ---------------------------------------------------------------------------
+# Layer components (pre-LayerNorm)
+# ---------------------------------------------------------------------------
+
+class GenaLMSelfOutput(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.dropout(self.dense(hidden_states))
+
+
+class GenaLMAttention(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        impl = getattr(config, "_attn_implementation", "eager")
+        if getattr(config, "sparse_block_size", 0):
+            if impl == "flash_attention_2":
+                raise ValueError(
+                    "flash_attention_2 is not supported by GENA-LM block-sparse checkpoints: "
+                    "Flash Attention cannot express their checkpoint-specific block mask. "
+                    "Use attn_implementation='eager' or 'sdpa'."
+                )
+            self.self = GenaLMBlockSparseSelfAttention(config, use_sdpa=(impl == "sdpa"))
+        else:
+            self.self = GENALM_SELF_ATTENTION_CLASSES[impl](config)
+        self.output = GenaLMSelfOutput(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), attn_weights
+
+
+class GenaLMIntermediate(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.intermediate_act_fn = ACT2FN[config.hidden_act] if isinstance(config.hidden_act, str) else config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.intermediate_act_fn(self.dense(hidden_states))
+
+
+class GenaLMOutput(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.dropout(self.dense(hidden_states))
+
+
+class GenaLMLayer(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.pre_attention_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.post_attention_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.attention = GenaLMAttention(config)
+        self.intermediate = GenaLMIntermediate(config)
+        self.output = GenaLMOutput(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(
+            self.pre_attention_ln(hidden_states), key_padding_mask, output_attentions
+        )
+        hidden_states = hidden_states + attn_out
+        ffn_out = self.output(self.intermediate(self.post_attention_ln(hidden_states)))
+        hidden_states = hidden_states + ffn_out
+        return hidden_states, attn_weights
+
+
+class GenaLMEncoder(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.layer = nn.ModuleList([GenaLMLayer(config) for _ in range(config.num_hidden_layers)])
+        self.last_layer_norm = config.last_layer_norm
+        if self.last_layer_norm:
+            self.last_layer_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    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 = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for layer in self.layer:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            hidden_states, attn_weights = layer(hidden_states, key_padding_mask, output_attentions)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
+
+        if self.last_layer_norm:
+            hidden_states = self.last_layer_ln(hidden_states)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        return hidden_states, all_hidden_states, all_attentions
+
+
+# ---------------------------------------------------------------------------
+# Embeddings, pooler, MLM head
+# ---------------------------------------------------------------------------
+
+class GenaLMEmbeddings(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.position_embedding_type = config.position_embedding_type
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        if self.position_embedding_type == "absolute":
+            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) + self.token_type_embeddings(token_type_ids)
+        if self.position_embedding_type == "absolute":
+            x = x + self.position_embeddings(self.position_ids[:, :T])
+        return self.dropout(self.LayerNorm(x))
+
+
+class GenaLMPooler(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        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 GenaLMPredictionHeadTransform(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.transform_act_fn = ACT2FN[config.hidden_act] if isinstance(config.hidden_act, str) else config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.LayerNorm(self.transform_act_fn(self.dense(hidden_states)))
+
+
+class GenaLMLMPredictionHead(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.transform = GenaLMPredictionHeadTransform(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.transform(hidden_states))
+
+    def _tie_weights(self):
+        self.decoder.bias = self.bias
+
+
+class GenaLMOnlyMLMHead(nn.Module):
+    def __init__(self, config: GenaLMConfig):
+        super().__init__()
+        self.predictions = GenaLMLMPredictionHead(config)
+
+    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+        return self.predictions(sequence_output)
+
+
+# ---------------------------------------------------------------------------
+# Top-level models
+# ---------------------------------------------------------------------------
+
+class GenaLMPreTrainedModel(PreTrainedModel):
+    config_class = GenaLMConfig
+    base_model_prefix = "bert"
+    supports_gradient_checkpointing = False
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+def _key_padding_mask(input_ids, attention_mask, pad_token_id):
+    if attention_mask is None:
+        mask = input_ids.ne(pad_token_id)
+    else:
+        mask = attention_mask.ne(0)
+    key_padding_mask = ~mask  # True = padding
+    if not key_padding_mask.any():
+        return None
+    return key_padding_mask
+
+
+class GenaLMModel(GenaLMPreTrainedModel):
+    def __init__(self, config: GenaLMConfig, add_pooling_layer: bool = False):
+        super().__init__(config)
+        self.embeddings = GenaLMEmbeddings(config)
+        self.encoder = GenaLMEncoder(config)
+        self.pooler = GenaLMPooler(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: 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
+
+        key_padding_mask = _key_padding_mask(input_ids, attention_mask, self.config.pad_token_id)
+
+        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 self.pooler is not None else None
+
+        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 GenaLMForMaskedLM(GenaLMPreTrainedModel):
+    _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]
+
+    def __init__(self, config: GenaLMConfig):
+        super().__init__(config)
+        self.bert = GenaLMModel(config, add_pooling_layer=False)
+        self.cls = GenaLMOnlyMLMHead(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_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(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.hidden_states, outputs.attentions)
+            output = tuple(v for v in output if v is not None)
+            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 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 512,
+  "unk_token": "[UNK]",
+  "sep_token": "[SEP]",
+  "pad_token": "[PAD]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]"
+}