presentation/app/model.py

"""From-scratch pre-LayerNorm BERT-style encoder for the BITS Phishing Detector.

Plain PyTorch only (torch / math / dataclasses / typing) so this file can be copied
verbatim to app/model.py for the Hugging Face Space. No pretrained weights, no
``transformers``.

Architecture (docs/research/research_model-math.md sec. 1 and 5):

* token + learned absolute position embeddings -> LayerNorm -> dropout (no segment
  embeddings: the classifier sees a single "subject\\n\\nbody" segment).
* ``n_layers`` pre-LN blocks::

      x = x + Dropout(OutProj(SDPA(QKV(LN1(x)))))
      x = x + Dropout(FFN_down(GELU(FFN_up(LN2(x)))))

  Pre-LN is chosen over BERT's post-LN because a 24-layer stack trained from
  scratch at batch 256 is far more tolerant of the learning rate and warmup under
  pre-LN (Xiong et al. 2020), and we have no pretrained weights to fall back on.
* a final LayerNorm (pre-LN needs it: the residual stream is otherwise un-normalised).
* attention through ``torch.nn.functional.scaled_dot_product_attention`` so the
  fused Flash / memory-efficient kernels run and the L x L score matrix is never
  stored for backward. No mask -> ``attn_mask=None`` (Flash-eligible, the packed
  pretraining path); a ``BoolTensor[B, L]`` key-padding mask (True = keep) ->
  ``[B, 1, 1, L]`` (memory-efficient backend, the fine-tune / inference path).

Default config (24L / d1024 / 16 heads / FFN 4096 / V 32768 / L 512) has exactly
337,476,608 parameters with the tied MLM head and 337,443,842 as the classifier.
"""

import math
from dataclasses import asdict, dataclass, fields
from typing import Any, Mapping

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.checkpoint import checkpoint


@dataclass
class EncoderConfig:
    vocab_size: int = 32768
    max_seq_len: int = 512
    d_model: int = 1024
    n_layers: int = 24
    n_heads: int = 16
    d_ff: int = 4096
    dropout: float = 0.1
    attention_dropout: float = 0.1
    # 1e-5 rather than BERT's 1e-12. Under autocast LayerNorm runs in fp32 so both
    # behave identically in training; 1e-5 is the safer value if the deployed model
    # is ever run in pure bf16 (8-bit mantissa: var + 1e-12 == var), and it is the
    # value every modern pre-LN encoder (RoBERTa-style, GPT-2) uses.
    layer_norm_eps: float = 1e-5
    pad_token_id: int = 0
    tie_mlm_head: bool = True
    grad_checkpointing: bool = False
    num_labels: int = 2

    def __post_init__(self) -> None:
        if self.d_model % self.n_heads:
            raise ValueError(f"d_model={self.d_model} not divisible by n_heads={self.n_heads}")
        if not 0 <= self.pad_token_id < self.vocab_size:
            raise ValueError("pad_token_id outside the vocabulary")

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> "EncoderConfig":
        """Build from the config.yaml ``model`` section.

        Keys that describe choices this file implements only one way (``norm``,
        ``activation``) are validated rather than silently ignored; any other
        unknown key is dropped.
        """
        norm = d.get("norm", "pre_ln")
        if norm != "pre_ln":
            raise ValueError(f"only norm=pre_ln is implemented, got {norm!r}")
        act = d.get("activation", "gelu")
        if act != "gelu":
            raise ValueError(f"only activation=gelu is implemented, got {act!r}")
        known = {f.name for f in fields(cls)}
        return cls(**{k: v for k, v in d.items() if k in known})

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


# --------------------------------------------------------------------------- init
def _init_module(module: nn.Module) -> None:
    """BERT init: N(0, 0.02) weights, zero biases, identity LayerNorm."""
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.padding_idx is not None:
            with torch.no_grad():
                module.weight[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)


def _zero_pad_row_grad(pad_token_id: int):
    """Post-accumulate hook keeping the [PAD] embedding row exactly zero through training.

    ``nn.Embedding(padding_idx)`` routes no gradient into that row, but with the tied MLM
    decoder the output matmul does (observed grad norm ~0.06 after one backward), so
    the row would drift from zero on the first optimizer step. Zeroing ``.grad[pad]``
    after accumulation is enough: AdamW then has zero moments for the row and decoupled
    weight decay leaves a zero weight at zero.
    """
    def hook(param: Tensor) -> None:
        if param.grad is not None:
            param.grad[pad_token_id].zero_()
    return hook


def _expand_key_padding_mask(attention_mask: Tensor | None) -> Tensor | None:
    """``BoolTensor[B, L]`` (True = keep) -> ``[B, 1, 1, L]`` broadcast over heads and queries."""
    if attention_mask is None:
        return None
    if attention_mask.dtype is not torch.bool:
        attention_mask = attention_mask.bool()
    return attention_mask[:, None, None, :]


# ------------------------------------------------------------------------ modules
class Embeddings(nn.Module):
    """Token + learned absolute position embeddings -> LayerNorm -> dropout.

    The [PAD] row of the token table is zero at init and stays zero while training
    (see ``_zero_pad_row_grad``). Tensor hooks are not carried by ``state_dict`` or
    ``copy.deepcopy``: a deep-copied module must not be trained, a module rebuilt from
    its config and ``load_state_dict`` is fine.
    """

    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_token_id)
        self.pos = nn.Embedding(cfg.max_seq_len, cfg.d_model)
        self.norm = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)
        self.dropout = nn.Dropout(cfg.dropout)
        self.register_buffer("position_ids", torch.arange(cfg.max_seq_len), persistent=False)
        self.tok.weight.register_post_accumulate_grad_hook(_zero_pad_row_grad(cfg.pad_token_id))

    def forward(self, input_ids: Tensor) -> Tensor:
        L = input_ids.shape[1]
        if L > self.pos.num_embeddings:
            raise ValueError(f"sequence length {L} > max_seq_len {self.pos.num_embeddings}")
        x = self.tok(input_ids) + self.pos(self.position_ids[:L])
        return self.dropout(self.norm(x))


class EncoderLayer(nn.Module):
    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.head_dim
        self.attention_dropout = cfg.attention_dropout
        self.ln1 = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.ln2 = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)
        self.ffn_up = nn.Linear(cfg.d_model, cfg.d_ff)
        self.ffn_down = nn.Linear(cfg.d_ff, cfg.d_model)
        self.dropout = nn.Dropout(cfg.dropout)

    def _qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        B, L, _ = x.shape
        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        return q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)  # [B, h, L, dh]

    def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor:
        q, k, v = self._qkv(self.ln1(x))
        o = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
        )
        o = o.transpose(1, 2).reshape(x.shape)
        x = x + self.dropout(self.out_proj(o))
        x = x + self.dropout(self.ffn_down(F.gelu(self.ffn_up(self.ln2(x)))))
        return x

    @torch.no_grad()
    def attention_probs(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor:
        """Explicit softmax(QK^T / sqrt(dh)) for inspection. Off the hot path.

        fp32, no dropout: this is the quantity the app shows as a signal, not the
        stochastic one used in training. Returns ``[B, h, L, L]``.
        """
        q, k, _ = self._qkv(self.ln1(x))
        scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(self.head_dim)
        if attn_mask is not None:
            scores = scores.masked_fill(~attn_mask, float("-inf"))
        return scores.softmax(dim=-1)


class Encoder(nn.Module):
    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.config = cfg
        self.embeddings = Embeddings(cfg)
        self.layers = nn.ModuleList(EncoderLayer(cfg) for _ in range(cfg.n_layers))
        self.final_norm = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)
        self.apply(_init_module)
        self._scale_residual_projections()

    def _scale_residual_projections(self) -> None:
        """GPT-2 style 1/sqrt(2 * n_layers) on the two projections that write into the
        residual stream (OutProj, FFN_down).

        Each of the 2 * n_layers residual branches adds an N(0, 0.02^2)-scaled
        contribution; without the rescale the residual variance grows linearly with
        depth, so the final-norm input and the early-training gradient magnitude
        would be ~sqrt(48) ~ 7x larger for 24 layers than for one. Keeping the stream
        O(1) at init is what lets a from-scratch 24-layer pre-LN model take a 1.5e-4
        peak LR with a short warmup instead of diverging in the first few hundred steps.
        """
        scale = 1.0 / math.sqrt(2 * self.config.n_layers)
        with torch.no_grad():
            for layer in self.layers:
                layer.out_proj.weight.mul_(scale)
                layer.ffn_down.weight.mul_(scale)

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        return_last_attn: bool = False,
    ) -> Tensor | tuple[Tensor, Tensor]:
        mask = _expand_key_padding_mask(attention_mask)
        x = self.embeddings(input_ids)
        use_ckpt = self.config.grad_checkpointing and self.training and torch.is_grad_enabled()
        last_attn: Tensor | None = None
        n = len(self.layers)
        for i, layer in enumerate(self.layers):
            if return_last_attn and i == n - 1:
                last_attn = layer.attention_probs(x, mask)
            if use_ckpt:
                x = checkpoint(layer, x, mask, use_reentrant=False)
            else:
                x = layer(x, mask)
        x = self.final_norm(x)
        if return_last_attn:
            assert last_attn is not None
            return x, last_attn
        return x


class MLMHead(nn.Module):
    """Linear(d, d) -> GELU -> LayerNorm -> decoder Linear(d, V) tied to the token embedding.

    Tying shares the Parameter object, so ``state_dict()`` carries two keys over one
    storage (``encoder.embeddings.tok.weight`` and ``head.decoder.weight``) and
    ``safetensors.torch.save_file`` refuses it ("Some tensors share memory"). Write
    checkpoints from ``EncoderForMLM.state_dict_for_safetensors()`` (or with
    ``safetensors.torch.save_model``, which drops the same alias) and read them back
    with ``EncoderForMLM.load_tied_state_dict``.
    """

    def __init__(self, cfg: EncoderConfig, token_embedding: nn.Embedding | None) -> None:
        super().__init__()
        self.dense = nn.Linear(cfg.d_model, cfg.d_model)
        self.norm = nn.LayerNorm(cfg.d_model, eps=cfg.layer_norm_eps)
        self.decoder = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        # Own bias so the tied matrix is not forced to also fit unigram frequencies.
        self.bias = nn.Parameter(torch.zeros(cfg.vocab_size))
        self.apply(_init_module)
        if token_embedding is not None:
            # Same Parameter object: one set of weights, one gradient accumulator.
            # Tying saves 33.5M parameters and (BERT/RoBERTa) regularises the
            # output distribution toward the input geometry.
            self.decoder.weight = token_embedding.weight

    def forward(self, hidden: Tensor) -> Tensor:
        h = self.norm(F.gelu(self.dense(hidden)))
        return self.decoder(h) + self.bias


class EncoderForMLM(nn.Module):
    """Encoder + MLM head for pretraining. Checkpoints: see ``state_dict_for_safetensors``."""

    TIED_KEY = "head.decoder.weight"
    EMBEDDING_KEY = "encoder.embeddings.tok.weight"

    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.config = cfg
        self.encoder = Encoder(cfg)
        self.head = MLMHead(cfg, self.encoder.embeddings.tok if cfg.tie_mlm_head else None)

    def state_dict_for_safetensors(self) -> dict[str, Tensor]:
        """``state_dict()`` minus the tied decoder's alias key.

        With ``tie_mlm_head`` the decoder weight IS the token embedding, and safetensors
        rejects two keys over one storage, so the pretraining trainer must save this
        dict (``safetensors.torch.save_file(model.state_dict_for_safetensors(), path)``)
        rather than the raw ``state_dict()``. ``from_pretrained_encoder`` reads only the
        ``encoder.*`` keys and accepts either form.
        """
        sd = self.state_dict()
        if self.config.tie_mlm_head:
            sd.pop(self.TIED_KEY)
        return sd

    def load_tied_state_dict(self, state_dict: Mapping[str, Tensor], strict: bool = True) -> Any:
        """``load_state_dict`` that accepts the dedup'd dict from ``state_dict_for_safetensors``
        / ``save_model``: the dropped alias is re-pointed at the embedding tensor before
        loading, so strict mode still reports every genuine mismatch."""
        sd = dict(state_dict)
        if self.config.tie_mlm_head and self.TIED_KEY not in sd and self.EMBEDDING_KEY in sd:
            sd[self.TIED_KEY] = sd[self.EMBEDDING_KEY]
        return self.load_state_dict(sd, strict=strict)

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        masked_positions: Tensor | None = None,
    ) -> Tensor:
        """Logits over the vocabulary.

        With ``masked_positions`` (``BoolTensor[B, L]``) the hidden states are
        gathered BEFORE the head, so logits are ``[N_masked, V]`` and the 85% of
        positions that carry no loss never materialise a V-wide row (saves ~136 MiB
        per 512-token sample at peak; research_model-math.md sec. 2).
        Without it, ``[B, L, V]``.
        """
        hidden = self.encoder(input_ids, attention_mask)
        if masked_positions is not None:
            hidden = hidden[masked_positions]
        return self.head(hidden)


class ClassificationHead(nn.Module):
    """CLS -> dropout -> Linear(d, d) -> tanh -> dropout -> Linear(d, num_labels)."""

    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.dropout = nn.Dropout(cfg.dropout)
        self.dense = nn.Linear(cfg.d_model, cfg.d_model)
        self.out = nn.Linear(cfg.d_model, cfg.num_labels)
        self.apply(_init_module)

    def forward(self, cls_hidden: Tensor) -> Tensor:
        h = torch.tanh(self.dense(self.dropout(cls_hidden)))
        return self.out(self.dropout(h))


class EncoderForClassification(nn.Module):
    def __init__(self, cfg: EncoderConfig) -> None:
        super().__init__()
        self.config = cfg
        self.encoder = Encoder(cfg)
        self.classifier = ClassificationHead(cfg)

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        return_last_attn: bool = False,
    ) -> Tensor | tuple[Tensor, Tensor]:
        out = self.encoder(input_ids, attention_mask, return_last_attn=return_last_attn)
        if return_last_attn:
            hidden, last_attn = out
            return self.classifier(hidden[:, 0]), last_attn
        return self.classifier(out[:, 0])

    @classmethod
    def from_pretrained_encoder(
        cls,
        mlm: "EncoderForMLM | Mapping[str, Tensor]",
        cfg: EncoderConfig | None = None,
        num_labels: int | None = None,
    ) -> "EncoderForClassification":
        """Classifier whose encoder is initialised from a pretrained ``EncoderForMLM``
        (module or its ``state_dict``). The MLM head is dropped; the classification
        head keeps its fresh init. Strict key matching: any missing/unexpected encoder
        key raises, because a silently half-loaded encoder would train but be garbage.
        """
        if isinstance(mlm, EncoderForMLM):
            cfg = cfg or mlm.config
            state = mlm.state_dict()
        else:
            if cfg is None:
                raise ValueError("cfg is required when loading from a state_dict")
            state = mlm
        if num_labels is not None and num_labels != cfg.num_labels:
            cfg = EncoderConfig(**{**cfg.to_dict(), "num_labels": num_labels})
        model = cls(cfg)
        load_encoder_weights(model.encoder, state, prefix="encoder.", strict=True)
        return model


def load_encoder_weights(
    encoder: Encoder,
    state_dict: Mapping[str, Tensor],
    prefix: str = "encoder.",
    strict: bool = True,
) -> tuple[list[str], list[str]]:
    """Load ``Encoder`` weights from a wrapper's state_dict (keys ``encoder.*``).

    Returns ``(missing_keys, unexpected_keys)`` relative to the encoder. Keys outside
    ``prefix`` (``head.*``, ``classifier.*``) are ignored, not reported. With
    ``strict`` any mismatch raises after reporting the exact keys.
    """
    sub = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)}
    result = encoder.load_state_dict(sub, strict=False)
    missing, unexpected = list(result.missing_keys), list(result.unexpected_keys)
    if strict and (missing or unexpected):
        raise RuntimeError(
            f"encoder weights do not match: missing={missing[:8]}{'...' if len(missing) > 8 else ''} "
            f"unexpected={unexpected[:8]}{'...' if len(unexpected) > 8 else ''}"
        )
    return missing, unexpected


def count_parameters(module: nn.Module) -> dict[str, Any]:
    """Unique-parameter counts (a tied weight is counted once, under the encoder's
    embedding where ``named_parameters`` first sees it).

    Breakdown keys: ``embeddings`` (token + position + embedding LN),
    ``encoder_layers``, ``final_norm``, ``head`` (MLM head or classifier).
    """
    total = sum(p.numel() for p in module.parameters())
    trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
    breakdown = {"embeddings": 0, "encoder_layers": 0, "final_norm": 0, "head": 0}
    for name, p in module.named_parameters():
        parts = name.split(".")
        if "embeddings" in parts:
            breakdown["embeddings"] += p.numel()
        elif "layers" in parts:
            breakdown["encoder_layers"] += p.numel()
        elif "final_norm" in parts:
            breakdown["final_norm"] += p.numel()
        else:
            breakdown["head"] += p.numel()
    return {"total": total, "trainable": trainable, **breakdown}


def _format_count_table(title: str, counts: Mapping[str, int]) -> str:
    rows = [f"{title}"]
    for k in ("embeddings", "encoder_layers", "final_norm", "head", "total", "trainable"):
        rows.append(f"  {k:<15} {counts[k]:>15,}")
    return "\n".join(rows)


if __name__ == "__main__":
    import sys
    from pathlib import Path

    import yaml  # only here: the module body must stay torch-only for app/model.py

    cfg_path = Path(__file__).resolve().parent.parent / "config.yaml"
    with open(cfg_path) as fh:
        model_cfg = EncoderConfig.from_dict(yaml.safe_load(fh)["model"])
    print(f"config: {model_cfg.to_dict()}")

    mlm = EncoderForMLM(model_cfg)
    mlm_counts = count_parameters(mlm)
    print(_format_count_table("EncoderForMLM (tied head)", mlm_counts))
    assert mlm.head.decoder.weight is mlm.encoder.embeddings.tok.weight, "MLM head is not tied"

    cls = EncoderForClassification.from_pretrained_encoder(mlm)
    cls_counts = count_parameters(cls)
    print(_format_count_table("EncoderForClassification", cls_counts))

    is_default = (
        model_cfg.n_layers, model_cfg.d_model, model_cfg.n_heads, model_cfg.d_ff,
        model_cfg.vocab_size, model_cfg.max_seq_len, model_cfg.tie_mlm_head, model_cfg.num_labels,
    ) == (24, 1024, 16, 4096, 32768, 512, True, 2)
    if is_default:
        # research_model-math.md sec. 1: per layer 12,596,224; encoder 302,311,424;
        # embeddings 34,080,768; tied MLM head 1,084,416; classifier head 1,051,650.
        assert mlm_counts["total"] == 337_476_608, mlm_counts["total"]
        assert cls_counts["total"] == 337_443_842, cls_counts["total"]
        assert mlm_counts["encoder_layers"] == 302_309_376, mlm_counts["encoder_layers"]
        assert mlm_counts["embeddings"] == 34_080_768, mlm_counts["embeddings"]
        assert mlm_counts["head"] == 1_084_416, mlm_counts["head"]
        assert cls_counts["head"] == 1_051_650, cls_counts["head"]
        print("parameter counts match research_model-math.md sec. 1")
    else:
        print("non-default config: skipping the 337,476,608 / 337,443,842 assertion")

    # Checkpoint format: the tied alias is dropped for safetensors and restored on load.
    sd = mlm.state_dict_for_safetensors()
    assert EncoderForMLM.TIED_KEY not in sd and len(sd) == len(mlm.state_dict()) - 1
    mlm.load_tied_state_dict(sd, strict=True)
    EncoderForClassification.from_pretrained_encoder(sd, model_cfg)
    print(f"safetensors-safe state_dict: {len(sd)} tensors, strict round-trip OK")

    # The [PAD] embedding row must receive no gradient despite the tied decoder.
    ids = torch.randint(5, model_cfg.vocab_size, (2, 16))
    mlm.train()
    mlm(ids).float().mean().backward()
    pad_grad = mlm.encoder.embeddings.tok.weight.grad[model_cfg.pad_token_id]
    assert float(pad_grad.abs().max()) == 0.0, "pad row received gradient"
    assert float(mlm.encoder.embeddings.tok.weight.grad[5:].abs().max()) > 0.0
    mlm.zero_grad(set_to_none=True)
    print("pad-row gradient is zero under the tied decoder")

    # Round-trip check of the shared contract on a tiny CPU batch.
    mlm.eval()
    mask = torch.ones(2, 16, dtype=torch.bool)
    mask[1, 10:] = False
    with torch.inference_mode():
        full = mlm(ids, mask)
        pos = torch.zeros(2, 16, dtype=torch.bool)
        pos[0, 3] = pos[1, 7] = True
        gathered = mlm(ids, mask, masked_positions=pos)
        logits, attn = cls(ids, mask, return_last_attn=True)
    assert full.shape == (2, 16, model_cfg.vocab_size) and gathered.shape == (2, model_cfg.vocab_size)
    assert logits.shape == (2, model_cfg.num_labels) and attn.shape == (2, model_cfg.n_heads, 16, 16)
    assert torch.allclose(attn[1, :, :, 10:], torch.zeros(()), atol=0), "padded keys received attention"
    print("forward shapes OK:", tuple(full.shape), tuple(gathered.shape), tuple(logits.shape), tuple(attn.shape))
    sys.exit(0)