"""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, }