src/mini_transformer/transformer.py

"""Implementation of the encoder-decoder transformer used across the project."""

from __future__ import annotations

import math
from typing import cast

import torch
import torch.nn as nn
from torch import Tensor

from . import modules
from .configs import ModelCfg
from .utils import broadcast_padding_mask, create_causal_mask, sample_from_logits

__all__ = ["BasicEncoderDecoderTransformer"]


class BasicEncoderDecoderTransformer(nn.Module):
    """Full encoder-decoder model with tied input/output embeddings."""

    def __init__(self, cfg: ModelCfg) -> None:
        super().__init__()

        if not isinstance(cfg, ModelCfg):
            raise TypeError(f"cfg must be a ModelCfg, got {type(cfg)}")

        if not isinstance(cfg.vocab_size, int):
            raise TypeError(f"vocab_size must be an int, got {type(cfg.vocab_size)}")
        if not cfg.vocab_size > 0:
            raise ValueError(f"vocab_size must be strictly greater than 0, got {cfg.vocab_size}")

        if not isinstance(cfg.d_model, int):
            raise TypeError(f"d_model must be an int, got {type(cfg.d_model)}")
        if not cfg.d_model > 0:
            raise ValueError(f"d_model must be strictly greater than 0, got {cfg.d_model}")

        if not isinstance(cfg.max_seq_len, int):
            raise TypeError(f"max_seq_len must be an int, got {type(cfg.max_seq_len)}")
        if not cfg.max_seq_len > 0:
            raise ValueError(f"max_seq_len must be greater than 0, got {cfg.max_seq_len}")
        if not 0 <= cfg.dropout_rate < 1:
            raise ValueError(f"dropout_rate must be in [0,1[, got {cfg.dropout_rate}")

        if not isinstance(cfg.pad_id, int):
            raise TypeError(f"pad_id must be an int, got {type(cfg.pad_id)}")
        if not (0 <= cfg.pad_id < cfg.vocab_size):
            raise ValueError(f"pad_id must be in [0, {cfg.vocab_size - 1}], got {cfg.pad_id}")
        if not isinstance(cfg.bos_id, int):
            raise TypeError(f"bos_id must be an int, got {type(cfg.bos_id)}")
        if not (0 <= cfg.bos_id < cfg.vocab_size):
            raise ValueError(f"bos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.bos_id}")
        if not isinstance(cfg.eos_id, int):
            raise TypeError(f"eos_id must be an int, got {type(cfg.eos_id)}")
        if not (0 <= cfg.eos_id < cfg.vocab_size):
            raise ValueError(f"eos_id must be in [0, {cfg.vocab_size - 1}], got {cfg.eos_id}")

        self.cfg = cfg

        style_raw = getattr(cfg, "layer_norm_style", None)
        if style_raw is None:
            style = "post" if cfg.num_layers <= 4 else "pre"
        elif not isinstance(style_raw, str):
            raise TypeError(
                f"layer_norm_style must be a string when provided, got {type(style_raw)}"
            )
        else:
            style = style_raw.lower()
        if style not in {"pre", "post"}:
            raise ValueError(
                "layer_norm_style must be either 'pre' or 'post'"
                f" (case-insensitive); got {style_raw!r}"
            )
        self.layer_norm_style = style
        cfg.layer_norm_style = style

        self.embed = modules.InputEmbedding(
            cfg.vocab_size, cfg.d_model, cfg.max_seq_len, cfg.pad_id, cfg.dropout_rate
        )
        self.encoder = modules.TransformerEncoder(
            cfg.d_model,
            cfg.num_heads,
            cfg.d_ff,
            cfg.num_layers,
            cfg.dropout_rate,
            layer_norm_style=style,
        )
        self.decoder = modules.TransformerDecoder(
            cfg.d_model,
            cfg.num_heads,
            cfg.d_ff,
            cfg.num_layers,
            cfg.dropout_rate,
            layer_norm_style=style,
        )

        self.lm_head = modules.LMHead(cfg.d_model, cfg.vocab_size)

        # tie weights
        self.lm_head.fc.weight = self.embed.token_embed.weight

        self.max_seq_len = cfg.max_seq_len
        self.pad_id = cfg.pad_id
        self.bos_id = cfg.bos_id
        self.eos_id = cfg.eos_id

        self.gradient_checkpointing = False

    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled
        self.encoder.use_ckpt = enabled
        self.decoder.use_ckpt = enabled

    def forward(
        self,
        src_ids: Tensor,
        tgt_ids: Tensor,
        src_padding_mask: Tensor,
        tgt_padding_mask: Tensor,
    ) -> Tensor:
        """Compute logits given input/output token ids and boolean padding masks."""

        if not isinstance(src_ids, Tensor):
            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
        if src_ids.dim() != 2:
            raise ValueError(
                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
            )
        if src_ids.dtype != torch.long:
            raise TypeError(f"src_ids must be torch.long (int64), got {src_ids.dtype}")

        if not isinstance(tgt_ids, Tensor):
            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
        if tgt_ids.dim() != 2:
            raise ValueError(
                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
            )
        if tgt_ids.dtype != torch.long:
            raise TypeError(f"tgt_ids must be torch.long (int64), got {tgt_ids.dtype}")

        if (
            not isinstance(src_padding_mask, Tensor)
            or src_padding_mask.dtype != torch.bool
            or src_padding_mask.dim() != 2
        ):
            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, S)")

        if (
            not isinstance(tgt_padding_mask, Tensor)
            or tgt_padding_mask.dtype != torch.bool
            or tgt_padding_mask.dim() != 2
        ):
            raise TypeError("tgt_padding_mask must be a boolean tensor shaped (B, S)")
        src_padding_mask = broadcast_padding_mask(src_padding_mask, self.cfg.num_heads)
        tgt_padding_mask = broadcast_padding_mask(tgt_padding_mask, self.cfg.num_heads)
        hidden_states = self.encode(src_ids, src_padding_mask)
        dec_hidden_states = self.decode(hidden_states, tgt_ids, src_padding_mask, tgt_padding_mask)
        logits = self.lm_head(dec_hidden_states)
        return logits  # (B, T_tgt, V)

    def encode(self, src_ids: Tensor, src_padding_mask: Tensor) -> Tensor:
        """Encode source tokens into memory representations."""

        if not isinstance(src_ids, torch.Tensor):
            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
        if src_ids.dim() != 2:
            raise ValueError(
                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
            )
        if src_ids.dtype != torch.long:
            raise TypeError("src_ids must be torch.long (int64)")

        if (
            not isinstance(src_padding_mask, Tensor)
            or src_padding_mask.dtype != torch.bool
            or src_padding_mask.dim() != 4
        ):
            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, H, 1, S)")

        x = self.embed(src_ids)  # (B, Sx, D)
        hidden_states = self.encoder(x, src_padding_mask)
        return hidden_states  # (B, Sx, D)

    def decode(
        self,
        hidden_states: Tensor,
        tgt_ids: Tensor,
        src_padding_mask: Tensor,
        tgt_padding_mask: Tensor,
    ) -> Tensor:
        if not isinstance(hidden_states, torch.Tensor):
            raise TypeError(f"hidden_states must be a torch.Tensor, got {type(hidden_states)}")
        if hidden_states.dim() != 3:
            raise ValueError(
                f"hidden_states must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(hidden_states.shape)}"
            )

        if not isinstance(tgt_ids, torch.Tensor):
            raise TypeError(f"tgt_ids must be a torch.Tensor, got {type(tgt_ids)}")
        if tgt_ids.dim() != 2:
            raise ValueError(
                f"tgt_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(tgt_ids.shape)}"
            )
        if tgt_ids.dtype != torch.long:
            raise TypeError("tgt_ids must be torch.long (int64)")

        if (
            not isinstance(src_padding_mask, Tensor)
            or src_padding_mask.dtype != torch.bool
            or src_padding_mask.dim() != 4
        ):
            raise TypeError("src_padding_mask must be a boolean tensor shaped (B, H, 1, S)")

        if (
            not isinstance(tgt_padding_mask, Tensor)
            or tgt_padding_mask.dtype != torch.bool
            or tgt_padding_mask.dim() != 4
        ):
            raise TypeError("tgt_padding_mask must be a boolean tensor shaped (B, H, 1, S)")

        tgt_causal_mask = create_causal_mask(tgt_ids, self.cfg.num_heads)

        y = self.embed(tgt_ids)  # (B, Sy, D)
        out = self.decoder(hidden_states, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask)
        return out  # (B, Sy, D)

    def generate(
        self,
        src_ids: Tensor,
        src_padding_mask: Tensor,
        max_new_tokens: int = 20,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        do_sample: bool = False,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0,
        no_repeat_ngram: int | None = None,
        min_steps_before_eos: int = 0,
        *,
        seed: int | None = None,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        if not isinstance(src_ids, torch.Tensor):
            raise TypeError(f"src_ids must be a torch.Tensor, got {type(src_ids)}")
        if src_ids.dim() != 2:
            raise ValueError(
                f"src_ids must be a 2D torch.Tensor of shape (B, S), got shape {tuple(src_ids.shape)}"
            )
        if src_ids.dtype != torch.long:
            raise TypeError("src_ids must be torch.long (int64)")

        if not isinstance(src_padding_mask, torch.Tensor):
            raise TypeError(
                f"src_padding_mask must be a torch.Tensor, got {type(src_padding_mask)}"
            )

        if not isinstance(max_new_tokens, int):
            raise TypeError(f"max_new_tokens must be an int, got {type(max_new_tokens)}")
        if max_new_tokens < 0:
            raise ValueError(
                f"max_new_tokens must be greater than or equal to 0, got {max_new_tokens}"
            )
        if not isinstance(temperature, float):
            raise TypeError(f"temperature must be a float, got {type(temperature)}")
        if not (temperature > 0.0):
            raise ValueError(f"temperature must be strictly greater than 0, got {temperature}")
        if top_k is not None and not isinstance(top_k, int):
            raise TypeError(f"top_k must be an int or None, got {type(top_k)}")
        if top_k is not None and top_k < 0:
            raise ValueError(f"top_k must be >= 0, got {top_k}")
        if top_p is not None and not isinstance(top_p, int | float):
            raise TypeError(f"top_p must be a number or None, got {type(top_p)}")
        if top_p is not None:
            top_p = float(top_p)
            if not (0.0 < top_p <= 1.0):
                raise ValueError(f"top_p must be in (0, 1], got {top_p}")
        if not isinstance(do_sample, bool):
            raise TypeError(f"do_sample must be a bool, got {type(do_sample)}")
        if not isinstance(presence_penalty, int | float) or not math.isfinite(
            float(presence_penalty)
        ):
            raise ValueError(
                f"presence_penalty must be a finite number >= 0, got {presence_penalty!r}."
            )
        presence_penalty = float(presence_penalty)
        if presence_penalty < 0.0:
            raise ValueError(f"presence_penalty must be >= 0, got {presence_penalty}.")
        if not isinstance(frequency_penalty, int | float) or not math.isfinite(
            float(frequency_penalty)
        ):
            raise ValueError(
                f"frequency_penalty must be a finite number >= 0, got {frequency_penalty!r}."
            )
        frequency_penalty = float(frequency_penalty)
        if frequency_penalty < 0.0:
            raise ValueError(f"frequency_penalty must be >= 0, got {frequency_penalty}.")
        if no_repeat_ngram is not None:
            if not isinstance(no_repeat_ngram, int):
                raise ValueError(f"no_repeat_ngram must be int or None, got {no_repeat_ngram!r}.")
            if no_repeat_ngram < 2:
                no_repeat_ngram = None  # size <2 is meaningless; treat as disabled
        if not isinstance(min_steps_before_eos, int) or min_steps_before_eos < 0:
            raise ValueError(
                f"min_steps_before_eos must be int >= 0, got {min_steps_before_eos!r}."
            )
        if seed is not None and not isinstance(seed, int):
            raise TypeError(f"seed must be int or None, got {type(seed)}")
        if generator is not None and not isinstance(generator, torch.Generator):
            raise TypeError(f"generator must be a torch.Generator or None, got {type(generator)}")

        self.eval()
        with torch.no_grad():
            batch_size, _ = src_ids.shape
            device = src_ids.device
            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
            allowed_new = max(0, self.max_seq_len - tgt_ids.size(1))
            max_new_tokens = min(max_new_tokens, allowed_new)
            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
            src_padding_mask = broadcast_padding_mask(src_padding_mask, self.cfg.num_heads)
            hidden_states = self.encode(src_ids, src_padding_mask)

            def _make_generator() -> torch.Generator:
                if device.type == "cpu":
                    return torch.Generator()
                return torch.Generator(device=device)

            rng = generator
            if seed is not None:
                rng = rng if rng is not None else _make_generator()
                rng.manual_seed(seed)
            elif rng is None and do_sample:
                rng = _make_generator()
            for step in range(max_new_tokens):
                # decode on current tgt_ids; take last-step logits
                tgt_padding_mask = tgt_ids == self.pad_id
                tgt_padding_mask = broadcast_padding_mask(tgt_padding_mask, self.cfg.num_heads)
                step_hidden = self.decode(
                    hidden_states, tgt_ids, src_padding_mask, tgt_padding_mask
                )  # (B, T, D)
                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
                disallowed = [self.pad_id] + ([self.eos_id] if step < min_steps_before_eos else [])
                next_ids = cast(
                    Tensor,
                    sample_from_logits(
                        logits,
                        temperature=temperature,
                        top_k=top_k,
                        top_p=top_p,
                        do_sample=do_sample,
                        disallowed_tokens=disallowed,
                        repetition_ctx=tgt_ids,
                        presence_penalty=presence_penalty,
                        frequency_penalty=frequency_penalty,
                        no_repeat_ngram_size=no_repeat_ngram,
                        rng=rng,
                    ),
                )  # (B,)
                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
                finished |= next_ids == self.eos_id
                if finished.all():
                    return tgt_ids
            return tgt_ids

    def debug_generate(
        self,
        src_ids: Tensor,
        src_padding_mask: Tensor | None,
        tokenizer,
        max_new_tokens: int = 20,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        do_sample: bool = False,
        *,
        seed: int | None = None,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        """
        Like generate(), but prints debug info at each step: input tokens, ids, output ids, tokens, logits, etc.
        """
        self.eval()
        with torch.no_grad():
            batch_size, _ = src_ids.shape
            device = src_ids.device
            tgt_ids = torch.full((batch_size, 1), self.bos_id, device=device, dtype=torch.long)
            allowed_new = max(0, self.max_seq_len - tgt_ids.size(1))
            max_new_tokens = min(max_new_tokens, allowed_new)
            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
            if src_padding_mask is not None:
                if (
                    src_padding_mask.dtype != torch.bool
                    or src_padding_mask.dim() != 2
                    or src_padding_mask.size(0) != batch_size
                ):
                    raise TypeError(
                        "src_padding_mask must be boolean tensor shaped (B, S) when provided"
                    )
                src_mask_2d = src_padding_mask
            else:
                src_mask_2d = torch.ones_like(src_ids, dtype=torch.bool, device=device)
            src_mask_4d = broadcast_padding_mask(src_mask_2d, self.cfg.num_heads)
            hidden_states = self.encode(src_ids, src_mask_4d)

            def _make_generator() -> torch.Generator:
                if device.type == "cpu":
                    return torch.Generator()
                return torch.Generator(device=device)

            rng = generator
            if seed is not None:
                rng = rng if rng is not None else _make_generator()
                rng.manual_seed(seed)
            elif rng is None and do_sample:
                rng = _make_generator()
            # print("[DEBUG] Input ids:", src_ids)
            print(
                "[DEBUG] Input tokens:", tokenizer.batch_decode(src_ids, skip_special_tokens=False)
            )
            for step in range(max_new_tokens):
                print(f"[DEBUG] Step {step + 1}")
                # print("  Output ids so far:", tgt_ids)
                print(
                    "  Output tokens so far:",
                    tokenizer.batch_decode(tgt_ids, skip_special_tokens=False),
                )
                tgt_padding_mask_2d = tgt_ids != self.pad_id
                tgt_padding_mask_4d = broadcast_padding_mask(
                    tgt_padding_mask_2d, self.cfg.num_heads
                )
                step_hidden = self.decode(
                    hidden_states, tgt_ids, src_mask_4d, tgt_padding_mask_4d
                )  # (B, T, D)
                logits = self.lm_head(step_hidden)[:, -1, :]  # (B, V)
                next_ids = cast(
                    Tensor,
                    sample_from_logits(
                        logits,
                        temperature=temperature,
                        top_k=top_k,
                        top_p=top_p,
                        do_sample=do_sample,
                        disallowed_tokens=[self.pad_id],
                        rng=rng,
                    ),
                )  # (B,)
                next_ids = torch.where(finished, torch.full_like(next_ids, self.eos_id), next_ids)
                tgt_ids = torch.cat([tgt_ids, next_ids.unsqueeze(1)], dim=-1)
                finished |= next_ids == self.eos_id

                # print("  Next token ids:", next_ids)
                print(
                    "  Next tokens:",
                    tokenizer.batch_decode(next_ids.unsqueeze(1), skip_special_tokens=False),
                )
                # print("  Logits (first 5):", logits[0, :5].cpu().numpy() if logits.shape[0] > 0 else None)
                if finished.all():
                    print("[DEBUG] All sequences finished.")
                    return tgt_ids
            print("[DEBUG] Max steps reached.")
            return tgt_ids