Lab: train complete transformer models

Keep embeddings, linear projections, normalization, and tensor operations. Implement attention, masks, residual branches, encoder, decoder, and the training loop yourself. Start with width 32, four heads, and one or two layers. Add end tokens and variable-length padding. Your solution must pass future-token independence, padding invariance, finite-gradient, and free-running evaluation checks before you increase its size.

#!/usr/bin/env python3
"""Inspectable transformer parts implemented without ``nn.Transformer``.

Mask convention
---------------
All boolean masks in this file use ``True`` for an allowed token or attention
connection.  This is intentionally stated here because some PyTorch APIs use
the opposite convention.  Attention raises ``ValueError`` instead of silently
running softmax when a query has no allowed key.

Important shapes use these names:

* ``B``: batch size
* ``T``: sequence length
* ``D``: model width (``d_model``)
* ``H``: number of attention heads
* ``Dh``: width of one head (``D // H``)
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Sequence

import torch
from torch import nn
from torch.nn import functional as F


def _check_valid_mask(mask: torch.Tensor, shape: tuple[int, int], name: str) -> None:
    if mask.dtype != torch.bool:
        raise TypeError(f"{name} must have dtype torch.bool (True means valid)")
    if tuple(mask.shape) != shape:
        raise ValueError(f"{name} must have shape {shape}, got {tuple(mask.shape)}")


def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
    """Rotate every adjacent pair ``(a, b)`` to ``(-b, a)``."""
    if x.shape[-1] % 2:
        raise ValueError("RoPE requires an even final dimension")
    pairs = x.reshape(*x.shape[:-1], -1, 2)
    first, second = pairs.unbind(dim=-1)
    return torch.stack((-second, first), dim=-1).flatten(start_dim=-2)


def rope_sin_cos(
    length: int,
    width: int,
    *,
    offset: int = 0,
    base: float = 10_000.0,
    device: torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return RoPE sine/cosine tables with shape ``(length, width)``."""
    if width % 2:
        raise ValueError("RoPE width must be even")
    if length < 0 or offset < 0:
        raise ValueError("RoPE length and offset must be non-negative")

    # Compute angles in float32 for stable CPU behavior, then match x's dtype.
    positions = torch.arange(offset, offset + length, device=device, dtype=torch.float32)
    pair_indices = torch.arange(0, width, 2, device=device, dtype=torch.float32)
    inverse_frequencies = base ** (-pair_indices / width)
    pair_angles = positions[:, None] * inverse_frequencies[None, :]
    angles = torch.repeat_interleave(pair_angles, repeats=2, dim=-1)
    return angles.sin().to(dtype=dtype), angles.cos().to(dtype=dtype)


def apply_rope(x: torch.Tensor, *, offset: int = 0, base: float = 10_000.0) -> torch.Tensor:
    """Apply rotary position embedding to ``(..., T, Dh)`` with an offset."""
    length, width = x.shape[-2:]
    sin, cos = rope_sin_cos(
        length,
        width,
        offset=offset,
        base=base,
        device=x.device,
        dtype=x.dtype,
    )
    # Add leading singleton dimensions so (T, Dh) broadcasts to (B, H, T, Dh).
    while sin.ndim < x.ndim:
        sin = sin.unsqueeze(0)
        cos = cos.unsqueeze(0)
    return x * cos + rotate_pairs(x) * sin


class RotaryEmbedding(nn.Module):
    """Small module wrapper around :func:`apply_rope`."""

    def __init__(self, base: float = 10_000.0) -> None:
        super().__init__()
        self.base = base

    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        return apply_rope(x, offset=offset, base=self.base)


@dataclass(frozen=True)
class KVCache:
    """Keys/values for one decoder layer, both shaped ``(B, H, T, Dh)``."""

    key: torch.Tensor
    value: torch.Tensor
    key_padding_mask: torch.Tensor  # (B, T), True for a real token

    def __post_init__(self) -> None:
        if self.key.ndim != 4 or self.value.ndim != 4:
            raise ValueError("cached key and value must have shape (B, H, T, Dh)")
        if self.key.shape != self.value.shape:
            raise ValueError("cached key and value shapes must match")
        expected = (self.key.shape[0], self.key.shape[2])
        _check_valid_mask(self.key_padding_mask, expected, "cached key_padding_mask")

    @property
    def length(self) -> int:
        return self.key.shape[2]


class ManualLayerNorm(nn.Module):
    """Layer normalization written from mean and variance for inspection."""

    def __init__(self, width: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.bias = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        normalized = (x - mean) * torch.rsqrt(variance + self.eps)
        return normalized * self.weight + self.bias


class ManualMultiHeadAttention(nn.Module):
    """Multi-head attention with explicit score, scale, mask, and softmax.

    ``query`` has shape ``(B, Tq, D)``. ``key_value`` has shape ``(B, Tk, D)``
    and defaults to ``query`` for self-attention.  Normally the return value is
    a tensor ``(B, Tq, D)``.  With ``use_cache=True`` the return value is
    ``(output, new_cache)``.
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        *,
        dropout: float = 0.0,
        use_rope: bool = False,
        bias: bool = True,
    ) -> None:
        super().__init__()
        if d_model % num_heads:
            raise ValueError("d_model must be divisible by num_heads")
        head_dim = d_model // num_heads
        if use_rope and head_dim % 2:
            raise ValueError("RoPE requires an even head dimension")
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = head_dim**-0.5
        self.use_rope = use_rope
        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
        self.k_proj = nn.Linear(d_model, d_model, bias=bias)
        self.v_proj = nn.Linear(d_model, d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
        self.attention_dropout = nn.Dropout(dropout)
        self.output_dropout = nn.Dropout(dropout)
        self.rope = RotaryEmbedding()

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, T, D)`` to ``(B, H, T, Dh)``."""
        batch, length, width = x.shape
        if width != self.d_model:
            raise ValueError(f"last dimension must be d_model={self.d_model}")
        return x.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)

    def merge_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, H, T, Dh)`` back to ``(B, T, D)``."""
        batch, heads, length, head_dim = x.shape
        if heads != self.num_heads or head_dim != self.head_dim:
            raise ValueError("head dimensions do not match this attention module")
        return x.transpose(1, 2).contiguous().reshape(batch, length, self.d_model)

    def _expand_attention_mask(
        self,
        mask: torch.Tensor,
        batch: int,
        query_length: int,
        key_length: int,
    ) -> torch.Tensor:
        if mask.dtype != torch.bool:
            raise TypeError("attention_mask must be bool (True means allowed)")
        if mask.ndim == 2:
            if tuple(mask.shape) != (query_length, key_length):
                raise ValueError("2D attention_mask must have shape (Tq, Tk)")
            return mask[None, None, :, :]
        if mask.ndim == 3:
            if tuple(mask.shape) != (batch, query_length, key_length):
                raise ValueError("3D attention_mask must have shape (B, Tq, Tk)")
            return mask[:, None, :, :]
        if mask.ndim == 4 and mask.shape[0] in (1, batch) and mask.shape[1] in (
            1,
            self.num_heads,
        ):
            if tuple(mask.shape[-2:]) != (query_length, key_length):
                raise ValueError("4D attention_mask has incorrect sequence dimensions")
            return mask
        raise ValueError("attention_mask must have 2, 3, or 4 dimensions")

    def forward(
        self,
        query: torch.Tensor,
        key_value: torch.Tensor | None = None,
        *,
        attention_mask: torch.Tensor | None = None,
        key_padding_mask: torch.Tensor | None = None,
        query_padding_mask: torch.Tensor | None = None,
        causal: bool = False,
        cache: KVCache | None = None,
        use_cache: bool = False,
        query_position_offset: int | None = None,
        key_position_offset: int = 0,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        if query.ndim != 3 or query.shape[-1] != self.d_model:
            raise ValueError("query must have shape (B, Tq, d_model)")
        if query.shape[1] == 0:
            raise ValueError("attention needs at least one query token")
        if cache is not None and not use_cache:
            raise ValueError("cache input requires use_cache=True so the updated cache is returned")
        if use_cache:
            expected_offset = cache.length if cache is not None else 0
            if key_position_offset != 0 or query_position_offset not in (None, expected_offset):
                raise ValueError("cached positions must start at zero and continue from cache length")
        if cache is not None and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and self.training and (
            self.attention_dropout.p > 0.0 or self.output_dropout.p > 0.0
        ):
            raise ValueError("KV caching with dropout is inference-only; call eval() first")

        is_self_attention = key_value is None or key_value is query
        source = query if key_value is None else key_value
        if source.ndim != 3 or source.shape[0] != query.shape[0] or source.shape[-1] != self.d_model:
            raise ValueError("key_value must have shape (B, Tk, d_model)")

        batch, query_length, _ = query.shape
        current_key_length = source.shape[1]
        if key_padding_mask is None:
            current_valid = torch.ones(
                (batch, current_key_length), dtype=torch.bool, device=query.device
            )
        else:
            _check_valid_mask(
                key_padding_mask, (batch, current_key_length), "key_padding_mask"
            )
            current_valid = key_padding_mask
        if query_padding_mask is not None:
            _check_valid_mask(
                query_padding_mask, (batch, query_length), "query_padding_mask"
            )

        q = self.split_heads(self.q_proj(query))
        new_k = self.split_heads(self.k_proj(source))
        new_v = self.split_heads(self.v_proj(source))

        inferred_offset = cache.length if cache is not None else (
            key_position_offset if is_self_attention else 0
        )
        q_offset = inferred_offset if query_position_offset is None else query_position_offset
        if q_offset < 0 or key_position_offset < 0:
            raise ValueError("position offsets must be non-negative")
        if self.use_rope:
            q = self.rope(q, offset=q_offset)
            new_k_offset = inferred_offset if cache is not None else key_position_offset
            new_k = self.rope(new_k, offset=new_k_offset)

        if cache is not None:
            if cache.key.shape[:2] != (batch, self.num_heads):
                raise ValueError("cache batch/head dimensions do not match")
            if cache.key.shape[-1] != self.head_dim:
                raise ValueError("cache head width does not match")
            k = torch.cat((cache.key, new_k), dim=2)
            v = torch.cat((cache.value, new_v), dim=2)
            valid_keys = torch.cat((cache.key_padding_mask, current_valid), dim=1)
        else:
            k, v, valid_keys = new_k, new_v, current_valid

        key_length = k.shape[2]
        allowed = valid_keys[:, None, None, :]  # (B, 1, 1, Tk)
        if causal:
            query_positions = torch.arange(
                q_offset, q_offset + query_length, device=query.device
            )
            first_key_position = 0 if cache is not None else key_position_offset
            key_positions = torch.arange(
                first_key_position, first_key_position + key_length, device=query.device
            )
            causal_mask = key_positions[None, :] <= query_positions[:, None]
            allowed = allowed & causal_mask[None, None, :, :]
        if attention_mask is not None:
            expanded = self._expand_attention_mask(
                attention_mask, batch, query_length, key_length
            )
            allowed = allowed & expanded

        # Expand before checking so each batch/head/query row is checked.
        allowed = allowed.expand(batch, self.num_heads, query_length, key_length)
        if not bool(allowed.any(dim=-1).all()):
            raise ValueError("attention mask contains an all-masked query row")

        # (B,H,Tq,Dh) @ (B,H,Dh,Tk) -> one score per query/key pair.
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores, dim=-1)
        weights = self.attention_dropout(weights)
        mixed_values = torch.matmul(weights, v)  # (B, H, Tq, Dh)
        output = self.out_proj(self.merge_heads(mixed_values))
        output = self.output_dropout(output)
        if query_padding_mask is not None:
            output = output * query_padding_mask.unsqueeze(-1).to(output.dtype)

        if use_cache:
            return output, KVCache(k, v, valid_keys)
        return output


class FeedForward(nn.Module):
    """Position-wise ``D -> hidden -> D`` network used by each block."""

    def __init__(self, d_model: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.input_projection = nn.Linear(d_model, hidden_dim)
        self.output_projection = nn.Linear(hidden_dim, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.output_projection(self.dropout(F.gelu(self.input_projection(x)))))


class EncoderBlock(nn.Module):
    """Pre-norm encoder block: non-causal self-attention, then FFN."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        attention_input = self.attention_norm(x)
        x = x + self.self_attention(
            attention_input,
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * valid_tokens.unsqueeze(-1).to(x.dtype)


class DecoderBlock(nn.Module):
    """Pre-norm decoder block with causal self-attention and cross-attention."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.self_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.cross_norm = ManualLayerNorm(d_model)
        self.cross_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        normalized = self.self_norm(x)
        x = x + self.self_attention(
            normalized,
            key_padding_mask=target_valid_tokens,
            query_padding_mask=target_valid_tokens,
            causal=True,
        )
        normalized = self.cross_norm(x)
        x = x + self.cross_attention(
            normalized,
            memory,
            key_padding_mask=memory_valid_tokens,
            query_padding_mask=target_valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * target_valid_tokens.unsqueeze(-1).to(x.dtype)


class DecoderOnlyBlock(nn.Module):
    """Pre-norm causal block that can append to a per-layer KV cache."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        valid_tokens: torch.Tensor,
        *,
        cache: KVCache | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        attention_result = self.self_attention(
            self.attention_norm(x),
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
            causal=True,
            cache=cache,
            use_cache=use_cache,
        )
        if use_cache:
            attention_output, new_cache = attention_result
        else:
            attention_output = attention_result
        x = x + attention_output
        x = x + self.feed_forward(self.ffn_norm(x))
        x = x * valid_tokens.unsqueeze(-1).to(x.dtype)
        if use_cache:
            return x, new_cache
        return x


class TransformerEncoder(nn.Module):
    """Stack of non-causal encoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [EncoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, valid_tokens)
        return self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)


class TransformerDecoder(nn.Module):
    """Stack of causal decoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [DecoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        for layer in self.layers:
            x = layer(
                x,
                memory,
                target_valid_tokens=target_valid_tokens,
                memory_valid_tokens=memory_valid_tokens,
            )
        return self.final_norm(x) * target_valid_tokens.unsqueeze(-1).to(x.dtype)


class Seq2SeqTransformer(nn.Module):
    """Complete encoder-decoder transformer for token IDs."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.source_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.target_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.encoder = TransformerEncoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.decoder = TransformerDecoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.output_projection = nn.Linear(d_model, vocab_size, bias=False)

    def encode(
        self, source_ids: torch.Tensor, source_valid: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if source_valid is None:
            source_valid = source_ids.ne(self.padding_id)
        _check_valid_mask(source_valid, tuple(source_ids.shape), "source_valid")
        if not bool(source_valid.any(dim=1).all()):
            raise ValueError("each source sequence needs at least one valid token")
        embedded = self.source_embedding(source_ids) * math.sqrt(self.d_model)
        return self.encoder(embedded, source_valid), source_valid

    def decode(
        self,
        target_ids: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid: torch.Tensor | None = None,
        memory_valid: torch.Tensor,
    ) -> torch.Tensor:
        if target_valid is None:
            target_valid = target_ids.ne(self.padding_id)
        embedded = self.target_embedding(target_ids) * math.sqrt(self.d_model)
        hidden = self.decoder(
            embedded,
            memory,
            target_valid_tokens=target_valid,
            memory_valid_tokens=memory_valid,
        )
        return self.output_projection(hidden)

    def forward(
        self,
        source_ids: torch.Tensor,
        target_input_ids: torch.Tensor,
        *,
        source_valid: torch.Tensor | None = None,
        target_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        memory, source_valid = self.encode(source_ids, source_valid)
        return self.decode(
            target_input_ids,
            memory,
            target_valid=target_valid,
            memory_valid=source_valid,
        )

    @torch.no_grad()
    def generate(
        self,
        source_ids: torch.Tensor,
        *,
        bos_id: int,
        eos_id: int,
        max_new_tokens: int,
        source_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Greedy decode and return generated IDs, including BOS and possible EOS."""
        self.eval()
        memory, source_valid = self.encode(source_ids, source_valid)
        generated = torch.full(
            (source_ids.shape[0], 1), bos_id, dtype=torch.long, device=source_ids.device
        )
        finished = torch.zeros(source_ids.shape[0], dtype=torch.bool, device=source_ids.device)
        for _ in range(max_new_tokens):
            logits = self.decode(generated, memory, memory_valid=source_valid)
            next_ids = logits[:, -1].argmax(dim=-1)
            next_ids = torch.where(finished, torch.full_like(next_ids, eos_id), next_ids)
            generated = torch.cat((generated, next_ids[:, None]), dim=1)
            finished |= next_ids.eq(eos_id)
            if bool(finished.all()):
                break
        return generated


class DecoderOnlyTransformer(nn.Module):
    """Decoder-only language model with one :class:`KVCache` per layer."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.layers = nn.ModuleList(
            [DecoderOnlyBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        token_ids: torch.Tensor,
        *,
        valid_tokens: torch.Tensor | None = None,
        caches: Sequence[KVCache | None] | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, list[KVCache]]:
        if token_ids.ndim != 2:
            raise ValueError("token_ids must have shape (B, T)")
        if valid_tokens is None:
            valid_tokens = token_ids.ne(self.padding_id)
        _check_valid_mask(valid_tokens, tuple(token_ids.shape), "valid_tokens")
        if not bool(valid_tokens.any(dim=1).all()):
            raise ValueError("each input sequence needs at least one valid token")
        if caches is None:
            caches = [None] * len(self.layers)
        if len(caches) != len(self.layers):
            raise ValueError("provide exactly one cache entry per layer")
        initialized = [cache is not None for cache in caches]
        if any(initialized) and not all(initialized):
            raise ValueError("layer caches must be either all initialized or all empty")
        lengths = {cache.length for cache in caches if cache is not None}
        if len(lengths) > 1:
            raise ValueError("all layer caches must have the same sequence length")

        x = self.token_embedding(token_ids) * math.sqrt(self.d_model)
        new_caches: list[KVCache] = []
        for layer, cache in zip(self.layers, caches):
            result = layer(x, valid_tokens, cache=cache, use_cache=use_cache)
            if use_cache:
                x, new_cache = result
                new_caches.append(new_cache)
            else:
                x = result
        x = self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)
        # Tied output weights avoid a separate, hidden language-model projection.
        logits = F.linear(x, self.token_embedding.weight)
        if use_cache:
            return logits, new_caches
        return logits

    @torch.no_grad()
    def generate_cached(self, prefix_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        """Prefill once, then decode one token at a time using the KV caches."""
        self.eval()
        generated = prefix_ids
        logits, caches = self(prefix_ids, use_cache=True)
        for _ in range(max_new_tokens):
            next_id = logits[:, -1].argmax(dim=-1, keepdim=True)
            generated = torch.cat((generated, next_id), dim=1)
            logits, caches = self(next_id, caches=caches, use_cache=True)
        return generated


__all__ = [
    "DecoderBlock",
    "DecoderOnlyBlock",
    "DecoderOnlyTransformer",
    "EncoderBlock",
    "FeedForward",
    "KVCache",
    "ManualLayerNorm",
    "ManualMultiHeadAttention",
    "RotaryEmbedding",
    "Seq2SeqTransformer",
    "TransformerDecoder",
    "TransformerEncoder",
    "apply_rope",
    "rope_sin_cos",
    "rotate_pairs",
]

#!/usr/bin/env python3
"""Train small manual transformers on deterministic generated CPU tasks.

Examples:

    python examples/pytorch/build_transformer.py seq2seq --task reverse --quick
    python examples/pytorch/build_transformer.py decoder-only --steps 100 --quick

The first task maps a source sequence to either a copy or a reversal.  The
second task predicts a repeating ascending token pattern.  Both use separate
held-out data and write a checkpoint only when ``--checkpoint`` is supplied.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import torch
from torch import nn

from transformer_parts import DecoderOnlyTransformer, Seq2SeqTransformer


PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
FIRST_DATA_ID = 3


@dataclass(frozen=True)
class SequenceBatch:
    """Padded IDs; ``targets`` use PAD_ID where loss must be ignored."""

    inputs: torch.Tensor
    decoder_inputs: torch.Tensor
    targets: torch.Tensor


def make_seq2seq_dataset(
    count: int,
    *,
    seed: int,
    task: str = "reverse",
    variable_length: bool = True,
    min_length: int = 3,
    max_length: int = 7,
    vocab_size: int = 16,
) -> SequenceBatch:
    """Create padded copy/reverse examples with independent deterministic RNG."""
    if task not in {"copy", "reverse"}:
        raise ValueError("task must be 'copy' or 'reverse'")
    if not (1 <= min_length <= max_length):
        raise ValueError("expected 1 <= min_length <= max_length")
    if vocab_size <= FIRST_DATA_ID:
        raise ValueError("vocab_size needs at least one ordinary data token")

    generator = torch.Generator(device="cpu").manual_seed(seed)
    source = torch.full((count, max_length), PAD_ID, dtype=torch.long)
    decoder_input = torch.full((count, max_length + 1), PAD_ID, dtype=torch.long)
    target = torch.full((count, max_length + 1), PAD_ID, dtype=torch.long)
    for row in range(count):
        if variable_length:
            length = int(
                torch.randint(min_length, max_length + 1, (), generator=generator).item()
            )
        else:
            length = max_length
        values = torch.randint(
            FIRST_DATA_ID, vocab_size, (length,), generator=generator, dtype=torch.long
        )
        answer = values if task == "copy" else values.flip(0)
        source[row, :length] = values
        decoder_input[row, 0] = BOS_ID
        decoder_input[row, 1 : length + 1] = answer
        target[row, :length] = answer
        target[row, length] = EOS_ID
    return SequenceBatch(source, decoder_input, target)


def make_pattern_dataset(
    count: int,
    *,
    seed: int,
    pattern_length: int = 8,
    vocab_size: int = 16,
) -> SequenceBatch:
    """Create cycles from independently sampled starts in a finite token set."""
    if pattern_length < 2:
        raise ValueError("pattern_length must be at least 2")
    data_token_count = vocab_size - FIRST_DATA_ID
    if data_token_count < 2:
        raise ValueError("vocab_size needs at least two ordinary data tokens")
    generator = torch.Generator(device="cpu").manual_seed(seed)
    starts = torch.randint(0, data_token_count, (count,), generator=generator)
    offsets = torch.arange(pattern_length)
    values = (starts[:, None] + offsets[None, :]) % data_token_count + FIRST_DATA_ID
    inputs = torch.cat(
        (torch.full((count, 1), BOS_ID, dtype=torch.long), values), dim=1
    )
    targets = torch.cat(
        (values, torch.full((count, 1), EOS_ID, dtype=torch.long)), dim=1
    )
    # SequenceBatch keeps one layout for both tasks. Decoder-only code uses
    # decoder_inputs and ignores inputs.
    return SequenceBatch(values, inputs, targets)


def masked_cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Average cross-entropy over non-padding targets only."""
    if logits.shape[:-1] != targets.shape:
        raise ValueError("logit and target sequence shapes do not match")
    if not bool(targets.ne(PAD_ID).any()):
        raise ValueError("cannot compute a loss for an all-padding target")
    return nn.functional.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        targets.reshape(-1),
        ignore_index=PAD_ID,
    )


def token_and_sequence_accuracy(
    logits: torch.Tensor, targets: torch.Tensor
) -> tuple[float, float]:
    """Return masked token accuracy and whole-sequence accuracy."""
    predicted = logits.argmax(dim=-1)
    valid = targets.ne(PAD_ID)
    token_accuracy = (predicted.eq(targets) & valid).sum().item() / valid.sum().item()
    correct_or_padding = predicted.eq(targets) | ~valid
    sequence_accuracy = correct_or_padding.all(dim=1).float().mean().item()
    return token_accuracy, sequence_accuracy


def most_common_baseline(train_targets: torch.Tensor, heldout_targets: torch.Tensor) -> float:
    """Accuracy when every non-padding position gets one frequent train token."""
    valid_train = train_targets[train_targets.ne(PAD_ID)]
    most_common = int(torch.bincount(valid_train).argmax().item())
    valid_heldout = heldout_targets.ne(PAD_ID)
    return (
        heldout_targets.eq(most_common).logical_and(valid_heldout).sum().item()
        / valid_heldout.sum().item()
    )


def minibatch_indices(
    count: int, batch_size: int, steps: int, *, seed: int
) -> list[torch.Tensor]:
    generator = torch.Generator(device="cpu").manual_seed(seed)
    return [torch.randint(count, (batch_size,), generator=generator) for _ in range(steps)]


def _generated_matches(generated: torch.Tensor, targets: torch.Tensor) -> float:
    """Compare greedy sequences after removing BOS and trailing padding."""
    matches = 0
    for generated_row, target_row in zip(generated.tolist(), targets.tolist()):
        generated_answer = generated_row[1:]
        expected = [token for token in target_row if token != PAD_ID]
        if EOS_ID in generated_answer:
            generated_answer = generated_answer[: generated_answer.index(EOS_ID) + 1]
        if generated_answer == expected:
            matches += 1
    return matches / targets.shape[0]


def _save_and_verify(
    model: nn.Module,
    checkpoint_path: Path,
    config: dict[str, int | float],
    example_ids: torch.Tensor,
    *,
    decoder_only: bool,
    source_ids: torch.Tensor | None = None,
) -> None:
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model_state_dict": model.state_dict(), "config": config}, checkpoint_path)
    model_class = DecoderOnlyTransformer if decoder_only else Seq2SeqTransformer
    restored = model_class(**config)
    payload = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    restored.load_state_dict(payload["model_state_dict"])
    restored.eval()
    model.eval()
    with torch.no_grad():
        if decoder_only:
            original = model(example_ids)
            reloaded = restored(example_ids)
        else:
            if source_ids is None:
                raise ValueError("source_ids are required for seq2seq verification")
            original = model(source_ids, example_ids)
            reloaded = restored(source_ids, example_ids)
    if not torch.equal(original, reloaded):
        raise AssertionError("checkpoint reload changed model output")


def run_seq2seq(
    *,
    steps: int,
    task: str,
    variable_length: bool,
    quick: bool,
    checkpoint_path: Path | None = None,
) -> dict[str, float]:
    """Train and evaluate an encoder-decoder transformer."""
    torch.manual_seed(2026)
    vocab_size = 16
    # A broad generated training set discourages memorizing a few sequences.
    train_count, heldout_count = (4_096, 128) if quick else (16_384, 512)
    batch_size = 32 if quick else 64
    config: dict[str, int | float] = {
        "vocab_size": vocab_size,
        "d_model": 32 if quick else 64,
        "num_heads": 4,
        "hidden_dim": 64 if quick else 128,
        "num_layers": 1 if quick else 2,
        "dropout": 0.0,
        "padding_id": PAD_ID,
    }
    train = make_seq2seq_dataset(
        train_count,
        seed=10,
        task=task,
        variable_length=variable_length,
        vocab_size=vocab_size,
    )
    heldout = make_seq2seq_dataset(
        heldout_count,
        seed=20,
        task=task,
        variable_length=variable_length,
        vocab_size=vocab_size,
    )
    model = Seq2SeqTransformer(**config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    losses: list[float] = []
    model.train()
    for indices in minibatch_indices(train_count, batch_size, steps, seed=30):
        optimizer.zero_grad(set_to_none=True)
        logits = model(train.inputs[indices], train.decoder_inputs[indices])
        loss = masked_cross_entropy(logits, train.targets[indices])
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        losses.append(loss.item())

    model.eval()
    with torch.no_grad():
        heldout_logits = model(heldout.inputs, heldout.decoder_inputs)
        heldout_loss = masked_cross_entropy(heldout_logits, heldout.targets).item()
        token_accuracy, teacher_forced_exact = token_and_sequence_accuracy(
            heldout_logits, heldout.targets
        )
        evaluation_count = min(64, heldout_count)
        generated = model.generate(
            heldout.inputs[:evaluation_count],
            bos_id=BOS_ID,
            eos_id=EOS_ID,
            max_new_tokens=heldout.targets.shape[1],
        )
        greedy_exact = _generated_matches(generated, heldout.targets[:evaluation_count])
    baseline = most_common_baseline(train.targets, heldout.targets)

    print(f"task: seq2seq {task}, variable_length={variable_length}")
    print(f"train source shape: {tuple(train.inputs.shape)}")
    print(f"held-out source shape: {tuple(heldout.inputs.shape)}")
    print(f"first/final training loss: {losses[0]:.4f} / {losses[-1]:.4f}")
    print(f"held-out loss: {heldout_loss:.4f}")
    print(f"most-common-token baseline accuracy: {baseline:.3f}")
    print(f"held-out token accuracy (teacher forced): {token_accuracy:.3f}")
    print(f"held-out exact accuracy (teacher forced): {teacher_forced_exact:.3f}")
    print(f"held-out exact accuracy (greedy decode): {greedy_exact:.3f}")

    if checkpoint_path is not None:
        _save_and_verify(
            model,
            checkpoint_path,
            config,
            heldout.decoder_inputs[:4],
            decoder_only=False,
            source_ids=heldout.inputs[:4],
        )
        print(f"checkpoint saved and verified: {checkpoint_path}")

    return {
        "initial_loss": losses[0],
        "final_loss": losses[-1],
        "heldout_loss": heldout_loss,
        "baseline_accuracy": baseline,
        "token_accuracy": token_accuracy,
        "teacher_forced_exact_accuracy": teacher_forced_exact,
        "greedy_exact_accuracy": greedy_exact,
    }


def run_decoder_only(
    *,
    steps: int,
    quick: bool,
    checkpoint_path: Path | None = None,
) -> dict[str, float]:
    """Train and evaluate a decoder-only transformer on a cyclic pattern."""
    torch.manual_seed(2027)
    vocab_size = 16
    train_count, heldout_count = (256, 96) if quick else (1_024, 256)
    batch_size = 32 if quick else 64
    config: dict[str, int | float] = {
        "vocab_size": vocab_size,
        "d_model": 32 if quick else 64,
        "num_heads": 4,
        "hidden_dim": 64 if quick else 128,
        "num_layers": 1 if quick else 2,
        "dropout": 0.0,
        "padding_id": PAD_ID,
    }
    train = make_pattern_dataset(train_count, seed=40, vocab_size=vocab_size)
    heldout = make_pattern_dataset(heldout_count, seed=50, vocab_size=vocab_size)
    model = DecoderOnlyTransformer(**config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    losses: list[float] = []
    model.train()
    for indices in minibatch_indices(train_count, batch_size, steps, seed=60):
        optimizer.zero_grad(set_to_none=True)
        logits = model(train.decoder_inputs[indices])
        loss = masked_cross_entropy(logits, train.targets[indices])
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        losses.append(loss.item())

    model.eval()
    with torch.no_grad():
        heldout_logits = model(heldout.decoder_inputs)
        heldout_loss = masked_cross_entropy(heldout_logits, heldout.targets).item()
        token_accuracy, exact_accuracy = token_and_sequence_accuracy(
            heldout_logits, heldout.targets
        )
        # Prefix BOS + first data token. The cache handles all later positions.
        evaluation_count = min(64, heldout_count)
        prefix = heldout.decoder_inputs[:evaluation_count, :2]
        generated = model.generate_cached(prefix, heldout.targets.shape[1] - 1)
        greedy_exact = _generated_matches(generated, heldout.targets[:evaluation_count])
    baseline = most_common_baseline(train.targets, heldout.targets)

    print("task: decoder-only ascending cyclic pattern")
    print(
        "dataset note: train and held-out rows are sampled independently, but "
        f"there are only {vocab_size - FIRST_DATA_ID} starts, so patterns repeat "
        "across both sets"
    )
    print(
        "generation note: the first target is random, so cached greedy decoding "
        "is given the first data token"
    )
    print(f"train input shape: {tuple(train.decoder_inputs.shape)}")
    print(f"held-out input shape: {tuple(heldout.decoder_inputs.shape)}")
    print(f"first/final training loss: {losses[0]:.4f} / {losses[-1]:.4f}")
    print(f"held-out loss: {heldout_loss:.4f}")
    print(f"most-common-token baseline accuracy: {baseline:.3f}")
    print(f"held-out token accuracy: {token_accuracy:.3f}")
    print(f"held-out exact accuracy (teacher forced): {exact_accuracy:.3f}")
    print(
        "held-out exact accuracy (cached greedy decode, given first data token): "
        f"{greedy_exact:.3f}"
    )

    if checkpoint_path is not None:
        _save_and_verify(
            model,
            checkpoint_path,
            config,
            heldout.decoder_inputs[:4],
            decoder_only=True,
        )
        print(f"checkpoint saved and verified: {checkpoint_path}")

    return {
        "initial_loss": losses[0],
        "final_loss": losses[-1],
        "heldout_loss": heldout_loss,
        "baseline_accuracy": baseline,
        "token_accuracy": token_accuracy,
        "teacher_forced_exact_accuracy": exact_accuracy,
        "greedy_exact_accuracy": greedy_exact,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="build", required=True)

    seq2seq = subparsers.add_parser("seq2seq", help="train an encoder-decoder model")
    seq2seq.add_argument("--task", choices=("copy", "reverse"), default="reverse")
    seq2seq.add_argument(
        "--fixed-length",
        action="store_true",
        help="use length 7 only (the default includes lengths 3 through 7)",
    )
    seq2seq.add_argument("--steps", type=int)
    seq2seq.add_argument("--quick", action="store_true")
    seq2seq.add_argument("--checkpoint", type=Path)

    decoder = subparsers.add_parser("decoder-only", help="train a causal language model")
    decoder.add_argument("--steps", type=int)
    decoder.add_argument("--quick", action="store_true")
    decoder.add_argument("--checkpoint", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.set_num_threads(1)
    steps = args.steps if args.steps is not None else (40 if args.quick else 500)
    if steps < 1:
        raise ValueError("--steps must be at least 1")
    if args.build == "seq2seq":
        run_seq2seq(
            steps=steps,
            task=args.task,
            variable_length=not args.fixed_length,
            quick=args.quick,
            checkpoint_path=args.checkpoint,
        )
    else:
        run_decoder_only(
            steps=steps,
            quick=args.quick,
            checkpoint_path=args.checkpoint,
        )


if __name__ == "__main__":
    main()