src/mini_transformer/utils.py

from __future__ import annotations

import math
import random
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor


def _require_matplotlib() -> Any:
    try:
        import matplotlib.pyplot as plt  # type: ignore
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise ImportError("Install mini-transformer[viz] to enable plotting utilities.") from exc
    return plt


# -----------------------------------------------------------------------------
# Global seeding utilities
# -----------------------------------------------------------------------------


def set_global_seed(seed: int, *, deterministic: bool = True) -> None:
    """Seed Python, NumPy, and PyTorch (CPU and CUDA) for reproducible runs."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def make_worker_init_fn(seed: int) -> Callable[[int], None]:
    """Return a DataLoader worker init function that derives unique seeds."""

    def _init_fn(worker_id: int) -> None:
        worker_seed = seed + worker_id
        random.seed(worker_seed)
        np.random.seed(worker_seed)
        torch.manual_seed(worker_seed)

    return _init_fn


# -----------------------------------------------------------------------------
# Attention building blocks
# -----------------------------------------------------------------------------


def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
    """Split the last dimension into ``(num_heads, d_head)`` and permute to (B, H, S, d_head)."""

    if not isinstance(x, torch.Tensor):
        raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
    if not isinstance(num_heads, int):
        raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
    if x.ndim != 3:
        raise ValueError(
            f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
        )
    if num_heads <= 0:
        raise ValueError(f"num_heads must be > 0; got {num_heads}")

    batch_size, seq_length, d_model = x.shape
    if d_model % num_heads != 0:
        raise ValueError(
            f"d_model ({d_model}) must be divisible by num_heads ({num_heads});"
            f" got remainder {d_model % num_heads}"
        )

    d_head = d_model // num_heads
    return x.reshape(batch_size, seq_length, num_heads, d_head).permute(0, 2, 1, 3)


def join_heads(x: torch.Tensor) -> torch.Tensor:
    """Merge ``(num_heads, d_head)`` back into the model dimension."""

    if not isinstance(x, torch.Tensor):
        raise TypeError(f"Expected torch.Tensor, got {type(x)}")
    if x.ndim != 4:
        raise ValueError(
            f"Expected 4D torch.Tensor (batch, num_heads, seq_len, d_head), got {tuple(x.shape)}"
        )

    batch_size, num_heads, seq_length, d_head = x.shape
    return x.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, num_heads * d_head)


def calculate_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    mask: torch.Tensor | None,
    *,
    dropout_p: float = 0.0,
    return_probs: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Scaled dot-product attention with optional dropout and probability return.

    Args:
        query, key, value: ``[B, H, S, Dh]`` tensors on the same device/dtype.
        mask: Optional boolean tensor broadcastable to ``[B, H, Sq, Sk]`` where
            ``True`` entries are masked. ``None`` disables masking.
        dropout_p: Dropout probability applied to the attention weights. Callers
            should pass ``0.0`` when not training.
        return_probs: When ``True`` returns ``(context, probs)``; otherwise only
            the context tensor is returned.
    """

    if not (query.dim() == key.dim() == value.dim() == 4):
        raise ValueError(
            "query, key, value must be 4D tensors shaped (B, H, S, Dh);"
            f" got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
        )
    if query.device != key.device or query.device != value.device:
        raise RuntimeError("query, key, value must be on the same device")

    if mask is not None:
        if mask.dtype != torch.bool:
            raise TypeError("mask must be boolean when provided")
        if mask.device != query.device:
            mask = mask.to(query.device)
        target_shape = (query.size(0), query.size(1), query.size(2), key.size(2))
        if mask.shape != target_shape:
            try:
                mask = mask.expand(target_shape)
            except RuntimeError as exc:
                raise ValueError(
                    f"mask with shape {tuple(mask.shape)} not broadcastable to {target_shape}"
                ) from exc

    p = float(dropout_p)
    if p < 0 or p >= 1:
        raise ValueError(f"dropout_p must be in [0, 1), got {dropout_p}")

    return _attention_with_probs(query, key, value, mask, p, return_probs)


def _attention_with_probs(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    mask: torch.Tensor | None,
    dropout_p: float,
    return_probs: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Manual attention path supporting optional probability return."""

    head_dim = query.size(-1)
    work_dtype = torch.float32 if query.dtype in (torch.float16, torch.bfloat16) else query.dtype

    q = query.to(work_dtype)
    k = key.to(work_dtype)
    scores = torch.matmul(q, k.transpose(-2, -1))
    scores.mul_(1.0 / math.sqrt(head_dim))

    full_mask_rows = None
    if mask is not None:
        mask = mask.to(scores.device)
        fill_value = torch.finfo(scores.dtype).min
        scores = scores.masked_fill(mask, fill_value)
        full_mask_rows = mask.all(dim=-1, keepdim=True)
        if full_mask_rows.any():
            scores = scores.masked_fill(full_mask_rows, 0.0)

    row_max = scores.max(dim=-1, keepdim=True).values
    row_max = torch.where(torch.isfinite(row_max), row_max, torch.zeros_like(row_max))
    logits = scores - row_max
    probs = torch.softmax(logits, dim=-1)
    if full_mask_rows is not None:
        probs = torch.where(full_mask_rows, torch.zeros_like(probs), probs)

    if dropout_p > 0.0:
        probs = F.dropout(probs, p=dropout_p, training=True)

    context = torch.matmul(probs.to(value.dtype), value)
    if return_probs:
        probs_out = probs.to(work_dtype)
        return context, probs_out
    return context


# -----------------------------------------------------------------------------
# Positional encodings
# -----------------------------------------------------------------------------


def sinusoidal_positional_encoding(S: int, D: int) -> torch.Tensor:
    """Return the classic sinusoidal positional encoding table (shape ``[S, D]``)."""

    if S <= 0 or D <= 0:
        raise ValueError(f"S and D must be > 0, got S={S}, D={D}")

    positions = torch.arange(S, dtype=torch.float32).unsqueeze(1)
    div_terms = torch.exp(torch.arange(0, D, 2, dtype=torch.float32) * -(math.log(10000.0) / D))

    pe = torch.zeros(S, D, dtype=torch.float32)
    pe[:, 0::2] = torch.sin(positions * div_terms)
    pe[:, 1::2] = torch.cos(positions * div_terms[: D // 2])
    return pe


# -----------------------------------------------------------------------------
# Sampling utilities
# -----------------------------------------------------------------------------


def _ensure_2d_logits(logits: torch.Tensor) -> torch.Tensor:
    """Reshape logits so that the last dimension is vocabulary sized and the rest flatten."""

    if logits.dim() == 2:
        return logits
    if logits.dim() >= 3:
        V = logits.size(-1)
        return logits.reshape(-1, V)
    raise ValueError(f"sample_from_logits: logits must be at least 2D, got {tuple(logits.shape)}")


def _apply_allow_deny_mask(
    logits: torch.Tensor,
    *,
    allowed_tokens: Iterable[int] | None,
    disallowed_tokens: Iterable[int] | None,
    filter_value: float,
) -> torch.Tensor:
    if allowed_tokens is not None:
        mask = torch.zeros_like(logits, dtype=torch.bool)
        idx = torch.tensor(list(allowed_tokens), device=logits.device)
        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
        if idx.numel() > 0:
            mask.index_fill_(-1, idx, True)
        logits = torch.where(mask, logits, torch.full_like(logits, filter_value))
    if disallowed_tokens is not None:
        idx = torch.tensor(list(disallowed_tokens), device=logits.device)
        idx = idx[(idx >= 0) & (idx < logits.size(-1))]
        if idx.numel() > 0:
            logits.index_fill_(-1, idx, filter_value)
    return logits


def _top_k_filtering(
    logits: torch.Tensor,
    top_k: int | None,
    *,
    min_tokens_to_keep: int,
    filter_value: float,
) -> torch.Tensor:
    if top_k is None or top_k <= 0:
        return logits
    k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
    values, _ = torch.topk(logits, k, dim=-1)
    threshold = values[..., -1, None]
    return torch.where(logits < threshold, torch.full_like(logits, filter_value), logits)


def _top_p_filtering(
    logits_scaled: torch.Tensor,
    probs: torch.Tensor,
    top_p: float,
    *,
    min_tokens_to_keep: int,
    filter_value: float,
) -> torch.Tensor:
    if top_p is None or not (0.0 < top_p < 1.0):
        return logits_scaled

    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)
    to_remove = cumulative > top_p
    to_remove[..., :min_tokens_to_keep] = False
    scatter_mask = torch.zeros_like(to_remove, dtype=torch.bool).scatter(-1, sorted_idx, to_remove)
    return torch.where(scatter_mask, torch.full_like(logits_scaled, filter_value), logits_scaled)


def sample_from_logits(
    logits: torch.Tensor,
    *,
    do_sample: bool = False,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    min_tokens_to_keep: int = 1,
    allowed_tokens: Iterable[int] | None = None,
    disallowed_tokens: Iterable[int] | None = None,
    repetition_ctx: torch.Tensor | None = None,
    presence_penalty: float = 0.0,
    frequency_penalty: float = 0.0,
    no_repeat_ngram_size: int | None = None,
    filter_value: float = -float("inf"),
    rng: torch.Generator | None = None,
    return_probs: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Warp logits (temperature, top-k, nucleus, penalties) and produce next-token ids."""

    if temperature <= 0:
        raise ValueError("temperature must be > 0")

    logits2d = _ensure_2d_logits(logits).to(dtype=torch.float32)
    batch, vocab = logits2d.shape

    logits2d = _apply_allow_deny_mask(
        logits2d,
        allowed_tokens=allowed_tokens,
        disallowed_tokens=disallowed_tokens,
        filter_value=filter_value,
    )

    if repetition_ctx is not None and (presence_penalty > 0.0 or frequency_penalty > 0.0):
        if repetition_ctx.dim() != 2 or repetition_ctx.size(0) != batch:
            raise ValueError(
                f"repetition_ctx must be [B, T]; got {tuple(repetition_ctx.shape)} with B={batch}"
            )
        counts = torch.zeros((batch, vocab), device=logits2d.device, dtype=torch.float32)
        ctx = repetition_ctx
        valid = (ctx >= 0) & (ctx < vocab)
        if valid.any():
            ids = ctx.masked_select(valid)
            bidx = (
                torch.arange(batch, device=logits2d.device)
                .unsqueeze(1)
                .expand_as(ctx)
                .masked_select(valid)
            )
            counts.index_put_(
                (bidx, ids), torch.ones_like(ids, dtype=torch.float32), accumulate=True
            )
        if presence_penalty > 0.0:
            logits2d = logits2d - presence_penalty * (counts > 0).to(logits2d.dtype)
        if frequency_penalty > 0.0:
            logits2d = logits2d - frequency_penalty * counts

    if (
        no_repeat_ngram_size is not None
        and no_repeat_ngram_size >= 2
        and repetition_ctx is not None
    ):
        ctx_cpu = repetition_ctx.detach().to("cpu")
        for b, seq in enumerate(ctx_cpu.tolist()):
            tokens = [t for t in seq if 0 <= t < vocab]
            if len(tokens) < no_repeat_ngram_size:
                continue
            history: dict[tuple[int, ...], set[int]] = {}
            n = int(no_repeat_ngram_size)
            for i in range(len(tokens) - (n - 1)):
                key = tuple(tokens[i : i + n - 1])
                nxt = tokens[i + n - 1]
                history.setdefault(key, set()).add(nxt)
            key = tuple(tokens[-(n - 1) :])
            banned = history.get(key)
            if banned:
                idx = torch.tensor(list(banned), device=logits2d.device, dtype=torch.long)
                logits2d[b].index_fill_(0, idx, filter_value)

    logits_scaled = logits2d / float(temperature)
    logits_scaled = _top_k_filtering(
        logits_scaled, top_k=top_k, min_tokens_to_keep=min_tokens_to_keep, filter_value=filter_value
    )
    probs = F.softmax(logits_scaled, dim=-1)
    if top_p is not None and 0.0 < top_p < 1.0:
        logits_scaled = _top_p_filtering(
            logits_scaled,
            probs,
            top_p=top_p,
            min_tokens_to_keep=min_tokens_to_keep,
            filter_value=filter_value,
        )
        probs = F.softmax(logits_scaled, dim=-1)

    if do_sample:
        next_ids = torch.multinomial(probs, num_samples=1, replacement=True, generator=rng).squeeze(
            -1
        )
    else:
        next_ids = torch.argmax(probs, dim=-1)

    next_ids = next_ids.to(dtype=torch.long, device=logits.device)
    if return_probs:
        return next_ids, probs.to(device=logits.device, dtype=probs.dtype)
    return next_ids


# -----------------------------------------------------------------------------
# Attention mask helpers
# -----------------------------------------------------------------------------


def create_causal_mask(x: torch.Tensor, num_heads: int) -> torch.Tensor:
    """Return a boolean causal mask ``[B, num_heads, S, S]`` for decoder self-attention."""

    if x.ndim != 2:
        raise ValueError(f"Expected input of shape (B, S), got {x.shape}")
    if num_heads <= 0:
        raise ValueError(f"num_heads must be > 0, got {num_heads}")

    batch, seq_len = x.shape
    if seq_len <= 0:
        raise ValueError(f"Sequence length must be > 0, got {seq_len}")

    base = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device), diagonal=1)
    return base.view(1, 1, seq_len, seq_len).expand(batch, num_heads, seq_len, seq_len).contiguous()


def create_qk_padding_mask(
    query_attention_mask: torch.Tensor, key_attention_mask: torch.Tensor
) -> torch.Tensor:
    """Combine query and key padding masks (boolean) into a ``[B, H, Sq, Sk]`` mask."""

    Bq, Hq, _, Sq = query_attention_mask.shape
    Bk, Hk, _, Sk = key_attention_mask.shape
    if Bq != Bk or Hq != Hk:
        raise ValueError(
            f"Padding mask batch/head mismatch: query {query_attention_mask.shape} vs key {key_attention_mask.shape}"
        )

    q_mask = query_attention_mask.to(torch.bool).view(Bq, Hq, Sq, 1)
    k_mask = key_attention_mask.to(torch.bool).view(Bk, Hk, 1, Sk)
    return (q_mask | k_mask).expand(Bq, Hq, Sq, Sk)


def broadcast_padding_mask(mask: torch.Tensor, num_heads: int) -> torch.Tensor:
    """Expand a ``[B, S]`` padding mask to ``[B, H, 1, S]``."""

    if not isinstance(mask, torch.Tensor):
        raise TypeError(f"mask must be a torch.Tensor, got {type(mask)}")
    if mask.dim() != 2:
        raise ValueError(f"mask must be [B, S], got shape {tuple(mask.shape)}")
    if not isinstance(num_heads, int) or num_heads <= 0:
        raise ValueError(f"num_heads must be a positive int, got {num_heads}")

    batch, seq_len = mask.shape
    return mask.unsqueeze(1).unsqueeze(2).expand(batch, num_heads, 1, seq_len)


def combine_masks(
    causal_mask: torch.Tensor | None, padding_mask: torch.Tensor | None
) -> torch.Tensor | None:
    """Combine causal and padding masks (boolean OR) while handling ``None`` values."""

    if padding_mask is None:
        return causal_mask.to(torch.bool) if causal_mask is not None else None
    if causal_mask is None:
        return padding_mask.to(torch.bool)

    m1 = causal_mask.to(torch.bool)
    m2 = padding_mask.to(torch.bool).to(device=m1.device)
    try:
        return m1 | m2
    except RuntimeError as exc:
        raise ValueError(
            f"Masks not broadcastable: causal {tuple(m1.shape)} vs padding {tuple(m2.shape)}"
        ) from exc


# -----------------------------------------------------------------------------
# Attention introspection / plotting
# -----------------------------------------------------------------------------


def extract_all_attention_maps(
    model,
    src_ids: torch.Tensor,
    tgt_ids: torch.Tensor,
    src_padding_2d: torch.Tensor,
    tgt_padding_2d: torch.Tensor,
):
    """Collect attention probability tensors for encoder only/decoder self/cross attention."""

    was_training = model.training
    model.eval()
    device = src_ids.device
    num_heads = model.cfg.num_heads

    src_pad_b = broadcast_padding_mask(src_padding_2d.to(device), num_heads)
    tgt_pad_b = broadcast_padding_mask(tgt_padding_2d.to(device), num_heads)
    tgt_causal = create_causal_mask(tgt_ids.to(device), num_heads)

    x = model.embed(src_ids.to(device))
    y = model.embed(tgt_ids.to(device))

    enc_layers = model.encoder.layers
    dec_layers = model.decoder.layers

    enc_self_maps: list[Tensor] = []
    dec_self_maps: list[Tensor] = []
    dec_cross_maps: list[Tensor] = []

    cur = x
    for layer in enc_layers:
        mha = layer.attention_layer
        attn_input = layer.norm1(cur) if getattr(layer, "pre_norm", False) else cur
        q = split_heads(mha.query_linear(attn_input), mha.num_heads)
        k = split_heads(mha.key_linear(attn_input), mha.num_heads)
        v = split_heads(mha.value_linear(attn_input), mha.num_heads)
        mask = create_qk_padding_mask(src_pad_b, src_pad_b)
        _, probs = calculate_attention(q, k, v, mask, dropout_p=0.0, return_probs=True)
        enc_self_maps.append(probs)
        cur = layer(cur, src_pad_b)

    mem_full = cur

    y_in = y
    for layer in dec_layers:
        pre_norm = getattr(layer, "pre_norm", False)

        self_mha = layer.self_attention_layer
        self_attn_input = layer.norm1(y_in) if pre_norm else y_in
        q = split_heads(self_mha.query_linear(self_attn_input), self_mha.num_heads)
        k = split_heads(self_mha.key_linear(self_attn_input), self_mha.num_heads)
        v = split_heads(self_mha.value_linear(self_attn_input), self_mha.num_heads)
        pad = create_qk_padding_mask(tgt_pad_b, tgt_pad_b)
        combined_mask = combine_masks(tgt_causal, pad)
        self_ctx, self_probs = calculate_attention(
            q, k, v, combined_mask, dropout_p=0.0, return_probs=True
        )
        dec_self_maps.append(self_probs)

        self_ctx_merged = join_heads(self_ctx)
        self_projected = self_mha.output_linear(self_ctx_merged)
        self_residual = y_in + layer.dropout1(self_projected)

        if pre_norm:
            cross_query_input = layer.norm2(self_residual)
        else:
            cross_query_input = layer.norm1(self_residual)

        cross_mha = layer.cross_attention_layer
        cq = split_heads(cross_mha.query_linear(cross_query_input), cross_mha.num_heads)
        ck = split_heads(cross_mha.key_linear(mem_full), cross_mha.num_heads)
        cv = split_heads(cross_mha.value_linear(mem_full), cross_mha.num_heads)
        cmask = create_qk_padding_mask(tgt_pad_b, src_pad_b)
        _, cross_probs = calculate_attention(cq, ck, cv, cmask, dropout_p=0.0, return_probs=True)
        dec_cross_maps.append(cross_probs)

        y_in = layer(mem_full, y_in, src_pad_b, tgt_pad_b, tgt_causal)

    maps = {"enc_self": enc_self_maps, "dec_self": dec_self_maps, "dec_cross": dec_cross_maps}
    model.train(was_training)
    return maps


def attach_tokens(maps_dict, tokenizer, src_ids, tgt_ids):
    """Attach decoded token strings to a maps dictionary produced by ``extract_all_attention_maps``."""

    src_tok = [
        tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in src_ids.tolist()
    ]
    tgt_tok = [
        tokenizer.convert_ids_to_tokens(row, skip_special_tokens=False) for row in tgt_ids.tolist()
    ]
    maps_dict["src_tokens"] = src_tok
    maps_dict["tgt_tokens"] = tgt_tok
    return maps_dict


def _imshow(ax, array2d, title="", vmin=0.0, vmax=1.0):
    image = ax.imshow(array2d, aspect="auto", vmin=vmin, vmax=vmax)
    ax.set_title(title, fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
    return image


def plot_layer_heads_grid(
    maps_dict: dict,
    layer_idx: int,
    batch: int = 0,
    heads: list[int] | None = None,
    figsize_per_cell: tuple[float, float] = (2.2, 2.2),
    show_colorbar: bool = True,
    vmin: float = 0.0,
    vmax: float = 1.0,
    *,
    show: bool = True,
    save_path: str | Path | None = None,
):
    """Plot a single layer as a 3-row (enc self / dec self / dec cross) grid."""

    plt = _require_matplotlib()

    enc_layers = maps_dict["enc_self"]
    dec_self_layers = maps_dict["dec_self"]
    dec_cross_layers = maps_dict["dec_cross"]

    enc_idx = min(layer_idx, len(enc_layers) - 1)
    dec_idx = min(layer_idx, len(dec_self_layers) - 1)

    enc = enc_layers[enc_idx][batch]
    dec_self = dec_self_layers[dec_idx][batch]
    dec_cross = dec_cross_layers[dec_idx][batch]

    num_heads = enc.size(0)
    heads = list(range(num_heads)) if heads is None else heads

    rows = 3
    cols = len(heads)
    fig_width = max(1, int(round(figsize_per_cell[0] * cols)))
    fig_height = max(1, int(round(figsize_per_cell[1] * rows)))
    fig, axes = plt.subplots(
        rows, cols, figsize=(fig_width, fig_height), squeeze=False, constrained_layout=True
    )

    images = []
    for column, head in enumerate(heads):
        images.append(
            _imshow(axes[0, column], enc[head].cpu().float().numpy(), f"Enc h{head}", vmin, vmax)
        )
        _imshow(
            axes[1, column], dec_self[head].cpu().float().numpy(), f"Dec self h{head}", vmin, vmax
        )
        _imshow(
            axes[2, column], dec_cross[head].cpu().float().numpy(), f"Dec cross h{head}", vmin, vmax
        )

    if show_colorbar and images:
        fig.colorbar(images[0], ax=axes, fraction=0.02, pad=0.01)
    fig.suptitle(f"Layer {layer_idx}", fontsize=12)

    if save_path is not None:
        save_path = Path(save_path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(save_path, dpi=150, bbox_inches="tight")

    if show:
        plt.show()
    else:
        plt.close(fig)
    return fig


def plot_all_layers_all_heads(
    maps_dict: dict,
    batch: int = 0,
    max_layers: int | None = None,
    heads: list[int] | None = None,
    figsize_per_cell: tuple[float, float] = (2.0, 2.0),
    vmin: float = 0.0,
    vmax: float = 1.0,
    save_pdf_path: str | None = None,
    *,
    show: bool = True,
):
    """Render attention grids for every layer (optionally saving a multi-page PDF)."""

    plt = _require_matplotlib()
    from matplotlib.backends.backend_pdf import PdfPages  # type: ignore

    enc_layers = len(maps_dict["enc_self"])
    dec_layers = len(maps_dict["dec_self"])
    total_layers = max(enc_layers, dec_layers)
    if max_layers is not None:
        total_layers = min(total_layers, max_layers)

    def _render_layers(record_page):
        for layer in range(total_layers):
            fig = plot_layer_heads_grid(
                maps_dict,
                layer_idx=layer,
                batch=batch,
                heads=heads,
                figsize_per_cell=figsize_per_cell,
                show_colorbar=True,
                vmin=vmin,
                vmax=vmax,
                show=show,
            )
            if record_page is not None:
                record_page(fig)
                plt.close(fig)

    if save_pdf_path:
        with PdfPages(save_pdf_path) as pdf:
            _render_layers(pdf.savefig)
        print(f"[saved] multi-page attention viewer -> {save_pdf_path}")
    else:
        _render_layers(None)


def _format_token_sequence(tokens: list[str]) -> str:
    return " | ".join(tokens)


def _resolve_attention_layers(available: int, requested: Sequence[int] | None) -> list[int]:
    if requested is None:
        return list(range(available))
    return sorted({int(idx) for idx in requested if 0 <= int(idx) < available})


def _resolve_attention_heads(available: int, requested: Sequence[int] | None) -> list[int]:
    if requested is None:
        return list(range(available))
    return sorted({int(idx) for idx in requested if 0 <= int(idx) < available})


# -----------------------------------------------------------------------------
# Debug helper
# -----------------------------------------------------------------------------


def debug_transformer_forward(
    model,
    tokenizer,
    src_ids: torch.Tensor,
    tgt_ids: torch.Tensor,
    *,
    logits: torch.Tensor | None = None,
    pad_id: int,
    device: str | torch.device = "cpu",
    batch_index: int = 0,
    sample_index: int = 0,
    show_attention: bool = False,
    save_attention: bool = False,
    average_heads: bool = False,
    attention_types: Sequence[str] = ("enc_self", "dec_self", "dec_cross"),
    attention_layers: Sequence[int] | None = None,
    attention_heads: Sequence[int] | None = None,
    attention_figsize: tuple[float, float] = (4.0, 4.0),
    skip_special_tokens: bool = False,
    log_fn: Callable[[str], None] | None = print,
    return_maps: bool = False,
    save_dir: str | Path | None = None,
    run_dir: str | Path | None = None,
) -> dict[str, Any]:
    """Run a forward pass, compute diagnostics, and optionally render attention heatmaps."""

    if log_fn is None:

        def log_fn(_: str) -> None:  # type: ignore[redefinition]
            return

    device = torch.device(device)

    if src_ids.dim() != 2 or tgt_ids.dim() != 2:
        raise ValueError("src_ids and tgt_ids must be rank-2 tensors")

    batch_size = src_ids.size(0)
    if not (0 <= sample_index < batch_size):
        raise IndexError(f"sample_index {sample_index} out of range for batch size {batch_size}")

    src_ids = src_ids.to(device)
    tgt_ids = tgt_ids.to(device)

    if tgt_ids.size(1) < 2:
        raise ValueError("tgt_ids must contain at least BOS and one target token")

    decoder_in = tgt_ids[:, :-1]
    labels = tgt_ids[:, 1:]

    src_pad_mask = src_ids.eq(pad_id)
    tgt_pad_mask = decoder_in.eq(pad_id)
    tgt_pad_full = tgt_ids.eq(pad_id)

    was_training = model.training
    attention_requested = show_attention or save_attention or return_maps

    if logits is None:
        if not callable(model):
            raise TypeError("model must be callable")

        model.eval()
        with torch.no_grad():
            logits = model(src_ids, decoder_in, src_pad_mask, tgt_pad_mask)
        model.train(was_training)
    else:
        logits = logits.to(device)

    if logits.dim() != 3 or logits.size(0) != batch_size:
        raise ValueError("logits must be of shape [batch, seq_len, vocab]")

    vocab_size = logits.size(-1)
    mask = labels.ne(pad_id)

    with torch.no_grad():
        log_probs = F.log_softmax(logits, dim=-1)
        nll = F.nll_loss(
            log_probs.reshape(-1, vocab_size),
            labels.reshape(-1),
            reduction="sum",
            ignore_index=pad_id,
        )

    tokens_total = int(mask.sum().item())
    avg_nll = (nll / max(tokens_total, 1)).item()
    perplexity = math.exp(avg_nll) if tokens_total > 0 else float("nan")

    pred_ids = logits.argmax(dim=-1)
    correct_mask = pred_ids.eq(labels) & mask
    correct_tokens = int(correct_mask.sum().item())
    token_accuracy = correct_tokens / max(tokens_total, 1)

    sample_labels = labels[sample_index]
    sample_preds = pred_ids[sample_index]
    sample_mask = mask[sample_index]

    sample_tokens_total = int(sample_mask.sum().item())
    sample_correct_tokens = int((sample_preds.eq(sample_labels) & sample_mask).sum().item())
    sample_token_accuracy = sample_correct_tokens / max(sample_tokens_total, 1)
    sample_exact_match = bool(torch.equal(sample_preds[sample_mask], sample_labels[sample_mask]))

    sample_decoder_in = decoder_in[sample_index]
    sample_pred_sequence = torch.cat([sample_decoder_in[:1], sample_preds], dim=0)

    with torch.no_grad():
        sample_log_probs = F.log_softmax(logits[sample_index], dim=-1)
        sample_nll = F.nll_loss(
            sample_log_probs,
            sample_labels,
            reduction="sum",
            ignore_index=pad_id,
        )

    sample_avg_nll = (sample_nll / max(sample_tokens_total, 1)).item()
    sample_perplexity = math.exp(sample_avg_nll) if sample_tokens_total > 0 else float("nan")

    sample_src_ids = src_ids[sample_index]
    sample_tgt_ids = tgt_ids[sample_index]

    sample_src_tokens = tokenizer.convert_ids_to_tokens(
        sample_src_ids.tolist(),
        skip_special_tokens=skip_special_tokens,
    )
    sample_tgt_tokens = tokenizer.convert_ids_to_tokens(
        sample_tgt_ids.tolist(),
        skip_special_tokens=skip_special_tokens,
    )
    sample_pred_tokens = tokenizer.convert_ids_to_tokens(
        sample_pred_sequence.tolist(),
        skip_special_tokens=skip_special_tokens,
    )

    sample_src_text = tokenizer.decode(sample_src_ids.tolist(), skip_special_tokens=True)
    sample_tgt_text = tokenizer.decode(sample_tgt_ids.tolist(), skip_special_tokens=True)
    sample_pred_text = tokenizer.decode(sample_pred_sequence.tolist(), skip_special_tokens=True)

    log_fn("\n--- DEBUG TRANSFORMER FORWARD ---")
    log_fn(f"Batch index : {batch_index}")
    log_fn(f"Sample index: {sample_index}")
    log_fn(f"Batch token accuracy: {token_accuracy * 100:.2f}% ({correct_tokens}/{tokens_total})")
    log_fn(
        f"Batch NLL / ppl : {avg_nll:.4f} / {perplexity:.4f}"
        if tokens_total > 0
        else "Batch NLL / ppl : n/a"
    )
    log_fn(
        f"Sample token accuracy: {sample_token_accuracy * 100:.2f}% ({sample_correct_tokens}/{sample_tokens_total})"
    )
    log_fn(
        f"Sample NLL / ppl : {sample_avg_nll:.4f} / {sample_perplexity:.4f}"
        if sample_tokens_total > 0
        else "Sample NLL / ppl : n/a"
    )
    log_fn(f"Sample exact match : {sample_exact_match}")

    log_fn("-- Source sequence --")
    log_fn(f"IDs   : {sample_src_ids.tolist()}")
    log_fn(f"Tokens: {_format_token_sequence(sample_src_tokens)}")
    log_fn(f"Text  : {sample_src_text}")

    log_fn("-- Target sequence --")
    log_fn(f"IDs   : {sample_tgt_ids.tolist()}")
    log_fn(f"Tokens: {_format_token_sequence(sample_tgt_tokens)}")
    log_fn(f"Text  : {sample_tgt_text}")

    log_fn("-- Predicted sequence --")
    log_fn(f"IDs   : {sample_pred_sequence.tolist()}")
    log_fn(f"Tokens: {_format_token_sequence(sample_pred_tokens)}")
    log_fn(f"Text  : {sample_pred_text}")

    result: dict[str, Any] = {
        "batch_index": batch_index,
        "token_accuracy": token_accuracy,
        "avg_negative_log_likelihood": avg_nll,
        "perplexity": perplexity,
        "tokens_total": tokens_total,
        "correct_tokens": correct_tokens,
        "sample": {
            "index": sample_index,
            "token_accuracy": sample_token_accuracy,
            "exact_match": sample_exact_match,
            "avg_negative_log_likelihood": sample_avg_nll,
            "perplexity": sample_perplexity,
            "source_ids": sample_src_ids.tolist(),
            "target_ids": sample_tgt_ids.tolist(),
            "predicted_ids": sample_pred_sequence.tolist(),
            "source_tokens": sample_src_tokens,
            "target_tokens": sample_tgt_tokens,
            "predicted_tokens": sample_pred_tokens,
            "source_text": sample_src_text,
            "target_text": sample_tgt_text,
            "predicted_text": sample_pred_text,
        },
    }

    if attention_requested:
        plt = _require_matplotlib()
        if not hasattr(model, "embed"):
            raise AttributeError("model must expose an embed() method for debugging")
        with torch.no_grad():
            attention_maps = extract_all_attention_maps(
                model,
                src_ids,
                tgt_ids,
                src_pad_mask,
                tgt_pad_full,
            )

        target_types = [t for t in attention_types if t in attention_maps]
        if not target_types:
            target_types = [k for k in ("enc_self", "dec_self", "dec_cross") if k in attention_maps]

        target_dir: Path | None = None
        if save_attention:
            if save_dir is not None:
                target_dir = Path(save_dir)
            elif run_dir is not None:
                target_dir = Path(run_dir) / "debug" / "attention"
            else:
                log_fn(
                    "[warn] save_attention requested but no save_dir/run_dir provided; skipping save."
                )
        if target_dir is not None:
            target_dir.mkdir(parents=True, exist_ok=True)

        sample_label = f"batch{batch_index}_sample{sample_index}"
        figure_records: list[dict[str, Any]] = []

        for att_type in target_types:
            layer_maps = attention_maps.get(att_type, [])
            if not layer_maps:
                continue

            layers_to_plot = _resolve_attention_layers(len(layer_maps), attention_layers)
            for layer_idx in layers_to_plot:
                layer_tensor = layer_maps[layer_idx][sample_index]
                head_candidates = _resolve_attention_heads(layer_tensor.size(0), attention_heads)
                if not head_candidates:
                    continue

                if average_heads:
                    averaged = layer_tensor[head_candidates].mean(dim=0).cpu().float().numpy()
                    fig, ax = plt.subplots(figsize=attention_figsize)
                    im = ax.imshow(averaged, aspect="auto", vmin=0.0, vmax=1.0)
                    ax.set_title(f"{att_type} L{layer_idx} (avg heads)")
                    ax.set_xlabel("Key index")
                    ax.set_ylabel("Query index")
                    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

                    save_path = None
                    if target_dir is not None:
                        save_path = target_dir / f"{sample_label}_{att_type}_L{layer_idx}_avg.png"
                        fig.savefig(save_path, dpi=150, bbox_inches="tight")
                    if show_attention:
                        plt.show()
                    else:
                        plt.close(fig)

                    figure_records.append(
                        {
                            "type": att_type,
                            "layer": layer_idx,
                            "heads": "average",
                            "path": str(save_path) if save_path else None,
                        }
                    )
                else:
                    cols = len(head_candidates)
                    fig, axes = plt.subplots(
                        1,
                        cols,
                        figsize=(attention_figsize[0] * cols, attention_figsize[1]),
                        squeeze=False,
                        constrained_layout=True,
                    )
                    for col, head_idx in enumerate(head_candidates):
                        ax = axes[0, col]
                        head_map = layer_tensor[head_idx].cpu().float().numpy()
                        im = ax.imshow(head_map, aspect="auto", vmin=0.0, vmax=1.0)
                        ax.set_title(f"{att_type} L{layer_idx} H{head_idx}")
                        ax.set_xlabel("Key index")
                        if col == 0:
                            ax.set_ylabel("Query index")
                    fig.colorbar(im, ax=axes.ravel().tolist(), fraction=0.046, pad=0.04)

                    save_path = None
                    if target_dir is not None:
                        head_tag = "-".join(str(h) for h in head_candidates)
                        save_path = (
                            target_dir / f"{sample_label}_{att_type}_L{layer_idx}_H{head_tag}.png"
                        )
                        fig.savefig(save_path, dpi=150, bbox_inches="tight")
                    if show_attention:
                        plt.show()
                    else:
                        plt.close(fig)

                    figure_records.append(
                        {
                            "type": att_type,
                            "layer": layer_idx,
                            "heads": head_candidates,
                            "path": str(save_path) if save_path else None,
                        }
                    )

        if target_dir is not None:
            raw_path = target_dir / f"{sample_label}_attention.pt"
            torch.save(
                {k: [v.cpu() for v in tensors] for k, tensors in attention_maps.items()}, raw_path
            )
            result.setdefault("attention", {})["raw_path"] = str(raw_path)

            summary_path = target_dir / f"{sample_label}_attention_all_layers.pdf"
            plot_all_layers_all_heads(
                attention_maps,
                batch=sample_index,
                figsize_per_cell=attention_figsize,
                save_pdf_path=str(summary_path),
                show=False,
            )
            result.setdefault("attention", {})["summary_path"] = str(summary_path)

        if figure_records:
            result.setdefault("attention", {})["figures"] = figure_records
        if return_maps:
            result["attention_maps"] = attention_maps

    log_fn("--- END DEBUG ---\n")
    return result


# -----------------------------------------------------------------------------
# Config compatibility checker
# -----------------------------------------------------------------------------


def check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg):
    """Ensure tokenizer and model configuration agree on core vocabulary settings."""

    if model_cfg.vocab_size != tokenizer_cfg.vocab_size:
        raise ValueError(
            f"Vocab size mismatch: model={model_cfg.vocab_size}, tokenizer={tokenizer_cfg.vocab_size}"
        )
    if model_cfg.max_seq_len != tokenizer_cfg.max_seq_len:
        raise ValueError(
            f"Max sequence length mismatch: model={model_cfg.max_seq_len}, tokenizer={tokenizer_cfg.max_seq_len}"
        )
    if model_cfg.pad_id != tokenizer_cfg.pad_id:
        raise ValueError(
            f"pad_id mismatch: model={model_cfg.pad_id}, tokenizer={tokenizer_cfg.pad_id}"
        )
    if model_cfg.bos_id != tokenizer_cfg.bos_id:
        raise ValueError(
            f"bos_id mismatch: model={model_cfg.bos_id}, tokenizer={tokenizer_cfg.bos_id}"
        )
    if model_cfg.eos_id != tokenizer_cfg.eos_id:
        raise ValueError(
            f"eos_id mismatch: model={model_cfg.eos_id}, tokenizer={tokenizer_cfg.eos_id}"
        )