mlm.py

"""MLM head for the hybrid MLM/CLM objective.

Ported from temp/gpt-bert-main/pretraining/model.py MaskClassifier, but
configurable to match the trunk's customizability:
- norm_type: "layernorm" or "rmsnorm" (reuses the existing field).
- ffn_activation: "gelu", "swiglu", "relu", or "identity" (reuses the
  existing field). The head's nonlinearity follows this choice.
- use_bias_ffn: whether the linear layers in the head use bias.
- weight tying: the final linear is tied to the word embedding, matching
  gpt-bert's MaskClassifier and our lm_head.

The head is controlled by model config (mlm_head_enabled) so checkpoint
loading and HF export remain strict and honest.
"""

import torch
import torch.nn as nn

from .transformer_components import RMSNorm, FeedForward


def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
    if norm_type == "layernorm":
        return nn.LayerNorm(hidden_dim)
    if norm_type == "rmsnorm":
        return RMSNorm(hidden_dim)
    raise ValueError(f"Unsupported norm_type: {norm_type}")


class MaskClassifier(nn.Module):
    """MLM prediction head with configurable norm and activation.

    Architecture (matching gpt-bert's MaskClassifier):
        norm -> FeedForward (activation, bias configurable) -> norm -> dropout -> linear (tied)

    The final linear is tied to the word embedding weight, matching
    gpt-bert's MaskClassifier and our lm_head.

    Args:
        hidden_dim: Model hidden dimension.
        vocab_size: Vocabulary size for the output projection.
        norm_type: "layernorm" or "rmsnorm".
        ffn_activation: "gelu", "swiglu", "relu", or "identity".
        use_bias_ffn: Whether the intermediate linear layers use bias.
        dropout: Dropout probability.
        word_embedding: The word embedding weight to tie the final linear to.
    """

    def __init__(
        self,
        hidden_dim: int,
        vocab_size: int,
        norm_type: str,
        ffn_activation: str,
        use_bias_ffn: bool,
        dropout: float,
        word_embedding: nn.Parameter,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.norm_type = norm_type
        self.ffn_activation = ffn_activation
        self.use_bias_ffn = use_bias_ffn
        self.dropout_p = dropout

        self.norm1 = _create_norm(norm_type, hidden_dim)
        self.feed_forward = FeedForward(
            input_dim=hidden_dim,
            hidden_dim=hidden_dim,
            dropout=dropout,
            activation=ffn_activation,
            use_bias=use_bias_ffn,
        )
        self.norm2 = _create_norm(norm_type, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.linear_out = nn.Linear(hidden_dim, vocab_size, bias=False)
        self.linear_out.weight = word_embedding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply the MLM head to contextualized embeddings.

        Args:
            x: Contextualized embeddings [batch, seq_len, hidden_dim] or
               flattened [num_masked, hidden_dim].

        Returns:
            Logits [batch, seq_len, vocab_size] or [num_masked, vocab_size].
        """
        x = self.norm1(x)
        x = self.feed_forward(x)
        x = self.norm2(x)
        x = self.dropout(x)
        x = self.linear_out(x)
        return x