model.py

"""A modern decoder-only transformer, from scratch, readable.

The same building blocks as Llama / DeepSeek — at mini scale:
  - RoPE     : tells attention WHERE each token sits (rotary positions)
  - RMSNorm  : keeps numbers stable between blocks (cheap normalization)
  - SwiGLU   : the "thinking" feed-forward block, gated for quality
  - weight tying: input embedding and output head share one matrix

Contract R2: every block commented in plain English.
"""

import math
from dataclasses import dataclass

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


@dataclass
class ModelConfig:
    vocab_size: int = 4096
    dim: int = 128          # width of every token's vector
    n_layers: int = 4       # how many attention+thinking floors
    n_heads: int = 4        # parallel attention "viewpoints"
    max_seq_len: int = 256  # longest context the model can see
    dropout: float = 0.0    # tiny data? keep 0; large runs may use 0.1


class RMSNorm(nn.Module):
    """Normalize a vector to a stable length, then let a learned gain rescale.

    Why: after many layers, numbers drift huge or tiny; training explodes.
    RMSNorm is LayerNorm's cheaper cousin (no mean subtraction) — the choice
    of Llama/DeepSeek.
    """

    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * rms * self.weight


def build_rope_cache(dim, max_seq_len, base=10000.0):
    """Precompute rotation angles for RoPE.

    Idea in one line: rotate each pair of vector coordinates by an angle that
    depends on the token's POSITION — attention then "feels" distance between
    words without any extra position embeddings.
    """
    half = dim // 2
    freqs = 1.0 / (base ** (torch.arange(0, half).float() / half))
    t = torch.arange(max_seq_len).float()
    angles = torch.outer(t, freqs)             # [seq, half]
    return torch.cos(angles), torch.sin(angles)


def apply_rope(x, cos, sin):
    """Rotate query/key pairs. x: [batch, heads, seq, head_dim]."""
    seq = x.shape[2]
    cos = cos[:seq].view(1, 1, seq, -1)
    sin = sin[:seq].view(1, 1, seq, -1)
    x1, x2 = x.chunk(2, dim=-1)                 # split coords into pairs
    return torch.cat([x1 * cos - x2 * sin,      # classic 2D rotation
                      x2 * cos + x1 * sin], dim=-1)


class Attention(nn.Module):
    """Causal self-attention: every token looks at all PREVIOUS tokens
    and decides whom to listen to. "Causal" = no peeking at the future,
    because the model's job is to predict it.
    """

    def __init__(self, cfg: ModelConfig):
        super().__init__()
        assert cfg.dim % cfg.n_heads == 0
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.dim // cfg.n_heads
        self.wq = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wk = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wv = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wo = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.dropout = cfg.dropout

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        # each token asks a Question, offers a Key, carries a Value
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        # fused attention with the causal (no-future) mask
        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.dropout if self.training else 0.0)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.wo(y)


class SwiGLU(nn.Module):
    """The "thinking" block. Gated: one path proposes, another path decides
    how much of it to let through. Empirically beats plain ReLU MLPs —
    used by Llama, PaLM, DeepSeek.
    """

    def __init__(self, cfg: ModelConfig):
        super().__init__()
        hidden = int(8 * cfg.dim / 3 / 64) * 64 or 2 * cfg.dim  # round to 64
        self.w_gate = nn.Linear(cfg.dim, hidden, bias=False)
        self.w_up = nn.Linear(cfg.dim, hidden, bias=False)
        self.w_down = nn.Linear(hidden, cfg.dim, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    """One floor of the tower: attention (talk to context), then SwiGLU
    (think it over). Residual "+x" paths let information skip floors —
    that is what makes deep towers trainable.
    """

    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.norm1 = RMSNorm(cfg.dim)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.dim)
        self.ffn = SwiGLU(cfg)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.ffn(self.norm2(x))
        return x


class TinyLLM(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm_out = RMSNorm(cfg.dim)
        self.head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        # weight tying: reading and writing tokens share one matrix —
        # fewer parameters, better tiny-model quality
        self.head.weight = self.tok_emb.weight
        cos, sin = build_rope_cache(cfg.dim // cfg.n_heads, cfg.max_seq_len)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(self._init)

    @staticmethod
    def _init(module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        x = self.tok_emb(idx)                       # tokens -> vectors
        for block in self.blocks:                   # up the tower
            x = block(x, self.rope_cos, self.rope_sin)
        x = self.norm_out(x)
        logits = self.head(x)                       # vectors -> vocab scores
        loss = None
        if targets is not None:
            # cross-entropy: how surprised was the model by the true next token
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50):
        """Write text: predict, sample, append, repeat."""
        self.eval()
        for _ in range(max_new_tokens):
            ctx = idx[:, -self.cfg.max_seq_len:]
            logits, _ = self(ctx)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx

    def num_params(self):
        return sum(p.numel() for p in self.parameters())


if __name__ == "__main__":
    # self-check: forward + backward on random tokens, and a generation call
    cfg = ModelConfig(vocab_size=4096, dim=128, n_layers=4, n_heads=4)
    model = TinyLLM(cfg)
    print(f"model built: {model.num_params():,} parameters")
    x = torch.randint(0, cfg.vocab_size, (2, 32))
    y = torch.randint(0, cfg.vocab_size, (2, 32))
    _, loss = model(x, y)
    loss.backward()
    print(f"forward+backward OK, random-guess loss = {loss.item():.3f} "
          f"(expected ~{math.log(cfg.vocab_size):.3f} = ln(vocab))")
    out = model.generate(x[:1, :4], max_new_tokens=8)
    print(f"generate OK: {out.shape[1]} tokens")