model_modernbert_engram.py

"""ModernBERT-based encoder for genomic MLM, augmented with Engram injection.

This wrapper adds Engram (conditional N-gram memory) deltas to specific
ModernBERT encoder layers via forward pre-hooks. It does NOT modify HF's
ModernBERT code, the ``engram`` module, or the ``model_modernbert`` module —
all three are reused unchanged.

Backbone parameter names match ``ModernBertWrapper`` exactly so a checkpoint
trained as ``model_type='modernbert'`` can be warm-started into this class
with ``strict=False``: only ``engram_blocks.*`` keys are missing.
"""
from __future__ import annotations

from dataclasses import replace
from typing import Sequence

import torch
import torch.nn as nn

from .engram import Engram, EngramConfig
from .model_modernbert import (
    ModernBertGenomicConfig,
    build_modernbert,
)


def _validate_layer_inject_ids(ids: Sequence[int], num_layers: int) -> list[int]:
    out: list[int] = []
    for idx, i in enumerate(ids):
        if isinstance(i, bool) or not isinstance(i, int):
            raise ValueError(
                f"layer_inject_ids[{idx}] = {i!r} is not an int (bools are rejected too)"
            )
        if i < 0 or i >= num_layers:
            raise ValueError(
                f"layer_inject_ids[{idx}] = {i} out of range [0, {num_layers})"
            )
        out.append(i)
    if len(set(out)) != len(out):
        raise ValueError(f"layer_inject_ids contains duplicates: {list(ids)}")
    return out


def _engram_key(i: int) -> str:
    return f"layer_{i:04d}"


class ModernBertEngramWrapper(nn.Module):
    """ModernBERT encoder for MLM with Engram modules injected at chosen layers.

    The wrapper:
      * Constructs a stock ``ModernBertForMaskedLM`` via :func:`build_modernbert`
        and stores it at ``self.model`` (identical attribute layout to
        :class:`ModernBertWrapper`, so state_dict keys are exactly the same
        for every backbone tensor).
      * Adds ``self.engram_blocks: nn.ModuleDict`` keyed by ``layer_{i:04d}``,
        one :class:`Engram` per index in ``engram_cfg.layer_inject_ids``.
      * Registers a ``forward_pre_hook(with_kwargs=True)`` on each target
        encoder layer that adds ``engram_blocks[layer_{i:04d}](hidden, input_ids)``
        to the layer's input ``hidden_states``.

    ``input_ids`` is stashed on ``self._current_input_ids`` at the start of
    each ``forward`` and never cleared inside the same call, so layers
    recomputed during backward (e.g. under ``torch.utils.checkpoint``) still
    see the right ids.
    """

    def __init__(self, cfg: ModernBertGenomicConfig, engram_cfg: EngramConfig):
        super().__init__()
        ids = _validate_layer_inject_ids(engram_cfg.layer_inject_ids, cfg.num_hidden_layers)
        engram_cfg = replace(
            engram_cfg,
            hidden_size=cfg.hidden_size,
            vocab_size=cfg.vocab_size,
            pad_id=cfg.pad_token_id,
        )
        self.cfg = cfg
        self.engram_cfg = engram_cfg
        # IDENTICAL outer attribute layout to ModernBertWrapper:
        self.model, self.hf_cfg = build_modernbert(cfg)
        # Sibling engram modules — NOT inside self.model, so backbone keys are unchanged.
        self.engram_blocks = nn.ModuleDict(
            {_engram_key(i): Engram(engram_cfg, layer_id=i) for i in ids}
        )
        for i in ids:
            self.model.model.layers[i].register_forward_pre_hook(
                self._make_hook(i), with_kwargs=True
            )
        self._current_input_ids: torch.Tensor | None = None

    @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 = 0
        memory = 0
        for name, p in self.named_parameters():
            if only_trainable and not p.requires_grad:
                continue
            n = p.numel()
            if "engram" in name and "embedding.emb" in name:
                memory += n
            else:
                backbone += n
        return {"backbone": backbone, "memory": memory, "total": backbone + memory}

    def _make_hook(self, layer_id: int):
        key = _engram_key(layer_id)
        # Cache to avoid one dict lookup per call.
        block_ref = self.engram_blocks[key]
        wrapper_self = self

        def hook(module, args, kwargs):
            ids = wrapper_self._current_input_ids
            if ids is None:
                return None  # cold-call safety (e.g. tracing); do nothing
            if args:
                hidden = args[0]
                delta = block_ref(hidden, input_ids=ids)
                new_args = (hidden + delta, *args[1:])
                return new_args, kwargs
            hidden = kwargs["hidden_states"]
            delta = block_ref(hidden, input_ids=ids)
            kwargs = {**kwargs, "hidden_states": hidden + delta}
            return args, kwargs

        return hook

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ) -> dict:
        # Overwrite-only on every forward; do NOT clear before/after model() so
        # that any backward-time recomputed forward (e.g. gradient checkpointing)
        # still reads the same ids that the original forward used.
        self._current_input_ids = input_ids
        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}