model_modernbert.py

"""ModernBERT-based encoder for genomic MLM, built on HuggingFace.

We construct ``ModernBertForMaskedLM`` from scratch with our 9-token DNA
vocabulary and configurable max_position. Supports both ``sdpa`` and
``flash_attention_2`` attention implementations.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import torch
import torch.nn as nn
import torch.nn.functional as F


# ─── Patch HF ModernBertConfig for global_attn_every_n_layers=1 ───────────
# transformers 5.7.0 has a bug: when all layers are "full_attention", the
# original convert_rope_params_to_dict unconditionally injects a
# "sliding_attention" entry into rope_parameters, then standardize_rope_params
# takes its Case-1 (single-rope) path and mutates rope_parameters with
# top-level "rope_type"/"rope_theta" keys — the strict-dataclass validator
# then rejects the result because those keys aren't in the Literal subset.
# Fix: only add "sliding_attention" rope params when layer_types actually has
# at least one sliding_attention layer. Drop it otherwise so standardize takes
# the Case-2 (per-layer-type) path with only "full_attention".
def _patch_modernbert_rope_params():
    from transformers import ModernBertConfig
    if getattr(ModernBertConfig, "_nucengram_patched", False):
        return
    def patched(self, **kwargs):
        rope_scaling = kwargs.pop("rope_scaling", None)
        default_rope_params = {
            "sliding_attention": {"rope_type": "default"},
            "full_attention":    {"rope_type": "default"},
        }
        self.rope_parameters = (self.rope_parameters
                                if self.rope_parameters is not None
                                else default_rope_params)
        if rope_scaling is not None:
            if "full_attention" in self.rope_parameters:
                self.rope_parameters["full_attention"].update(rope_scaling)
            if "sliding_attention" in self.rope_parameters:
                self.rope_parameters["sliding_attention"].update(rope_scaling)

        needs_sliding = (self.layer_types is not None
                         and "sliding_attention" in set(self.layer_types))

        if self.rope_parameters.get("full_attention") is None:
            self.rope_parameters["full_attention"] = {"rope_type": "default"}
        self.rope_parameters["full_attention"].setdefault(
            "rope_theta", kwargs.pop("global_rope_theta", self.default_theta["global"]))

        if needs_sliding:
            if self.rope_parameters.get("sliding_attention") is None:
                self.rope_parameters["sliding_attention"] = {"rope_type": "default"}
            self.rope_parameters["sliding_attention"].setdefault(
                "rope_theta", kwargs.pop("local_rope_theta", self.default_theta["local"]))
        else:
            # No sliding-attention layers → drop the entry so standardize_rope_params
            # takes Case-2 (per-layer-type) with only "full_attention" key.
            self.rope_parameters.pop("sliding_attention", None)
            kwargs.pop("local_rope_theta", None)

        self.standardize_rope_params()
        return kwargs
    ModernBertConfig.convert_rope_params_to_dict = patched
    ModernBertConfig._nucengram_patched = True

_patch_modernbert_rope_params()

from .tokenizer import VOCAB_SIZE, PAD_ID, BOS_ID, EOS_ID, MASK_ID


AttentionImpl = Literal["sdpa", "flash_attention_2", "eager"]


@dataclass
class ModernBertGenomicConfig:
    vocab_size: int = VOCAB_SIZE
    hidden_size: int = 384
    intermediate_size: int = 1024
    num_hidden_layers: int = 8
    num_attention_heads: int = 6
    max_position_embeddings: int = 8192
    local_attention: int = 128            # span for local attention layers
    global_attn_every_n_layers: int = 3   # =3 alternating GLOBAL+2×LOCAL; =1 every layer GLOBAL
    rope_theta_global: float = 160000.0   # bigger theta for long context
    rope_theta_local: float = 10000.0
    attn_implementation: AttentionImpl = "sdpa"
    pad_token_id: int = PAD_ID
    bos_token_id: int = BOS_ID
    eos_token_id: int = EOS_ID
    cls_token_id: int = BOS_ID
    sep_token_id: int = EOS_ID
    mask_token_id: int = MASK_ID
    norm_eps: float = 1e-5
    embedding_dropout: float = 0.0
    attention_dropout: float = 0.0
    mlp_dropout: float = 0.0
    initializer_range: float = 0.02
    tie_word_embeddings: bool = True
    sparse_prediction: bool = False
    deterministic_flash_attn: bool = False
    bf16: bool = True


def build_modernbert(cfg: ModernBertGenomicConfig):
    """Construct a ModernBertForMaskedLM with the given config."""
    from transformers import ModernBertConfig, ModernBertForMaskedLM

    # When global_attn_every_n_layers=1 every layer is full_attention — no
    # sliding rope params needed. The HF patch above keeps strict validation happy.
    if cfg.global_attn_every_n_layers == 1:
        rope_parameters = {"full_attention": {"rope_theta": cfg.rope_theta_global}}
    else:
        rope_parameters = {
            "full_attention":     {"rope_theta": cfg.rope_theta_global},
            "sliding_attention":  {"rope_theta": cfg.rope_theta_local},
        }

    hf_cfg = ModernBertConfig(
        vocab_size=cfg.vocab_size,
        hidden_size=cfg.hidden_size,
        intermediate_size=cfg.intermediate_size,
        num_hidden_layers=cfg.num_hidden_layers,
        num_attention_heads=cfg.num_attention_heads,
        hidden_activation="gelu",
        max_position_embeddings=cfg.max_position_embeddings,
        norm_eps=cfg.norm_eps,
        norm_bias=False,
        pad_token_id=cfg.pad_token_id,
        bos_token_id=cfg.bos_token_id,
        eos_token_id=cfg.eos_token_id,
        cls_token_id=cfg.cls_token_id,
        sep_token_id=cfg.sep_token_id,
        attention_bias=False,
        attention_dropout=cfg.attention_dropout,
        local_attention=cfg.local_attention,
        global_attn_every_n_layers=cfg.global_attn_every_n_layers,
        rope_parameters=rope_parameters,
        embedding_dropout=cfg.embedding_dropout,
        mlp_bias=False,
        mlp_dropout=cfg.mlp_dropout,
        decoder_bias=True,
        classifier_pooling="mean",
        deterministic_flash_attn=cfg.deterministic_flash_attn,
        sparse_prediction=cfg.sparse_prediction,
        tie_word_embeddings=cfg.tie_word_embeddings,
        initializer_range=cfg.initializer_range,
    )

    # Use FA2 only if we asked for it AND the install supports it.
    impl = cfg.attn_implementation
    if impl == "flash_attention_2":
        try:
            import flash_attn  # noqa: F401
        except ImportError as e:  # pragma: no cover
            raise RuntimeError(
                "attn_implementation='flash_attention_2' requested but "
                "flash-attn is not importable. Install flash-attn or use 'sdpa'."
            ) from e

    model = ModernBertForMaskedLM(hf_cfg)
    # HuggingFace stores the chosen attention impl on the config as well as on
    # individual modules; we set it explicitly on the loaded model so it takes
    # effect for our from-scratch instantiation.
    model.config._attn_implementation = impl
    if hasattr(model, "set_attn_implementation"):
        model.set_attn_implementation(impl)

    return model, hf_cfg


# ---------------------------------------------------------------------------
# common adapter so the training loop can swap models
# ---------------------------------------------------------------------------
class ModernBertWrapper(nn.Module):
    """Thin adapter that gives a ModernBertForMaskedLM the same forward
    signature as our in-house ``NucEngramModel.forward`` (returns dict with
    ``loss``, ``logits``, ``hidden``).
    """

    def __init__(self, cfg: ModernBertGenomicConfig):
        super().__init__()
        self.cfg = cfg
        self.model, self.hf_cfg = build_modernbert(cfg)

    @property
    def attn_implementation(self) -> str:
        impl = getattr(self.model.config, "_attn_implementation", None)
        if impl:
            return impl
        return getattr(self.model.config, "attn_implementation", "unknown")

    def num_params(self, only_trainable: bool = True) -> dict:
        backbone = sum(p.numel() for p in self.parameters()
                       if p.requires_grad or not only_trainable)
        return {"backbone": backbone, "memory": 0, "total": backbone}

    def forward(self,
                input_ids: torch.Tensor,
                labels: torch.Tensor | None = None,
                attention_mask: torch.Tensor | None = None) -> dict:
        # ModernBert wants attention_mask in {0, 1} (1 = keep).
        if attention_mask is not None:
            attention_mask = attention_mask.to(torch.long)
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            output_hidden_states=False,
        )
        return {
            "loss": out.loss,
            "logits": out.logits,
            "hidden": None,
        }