model.py

"""
GPT Language Model — Built Entirely from Scratch
=================================================
This file contains the COMPLETE transformer architecture.
Nothing is imported from `transformers` — every layer is hand-written
so you understand exactly what's happening inside a GPT.

Architecture overview (what happens when text goes through the model):

    Input Text
        ↓
    Token Embedding  (words → vectors)
        +
    Position Embedding  (position → vectors)
        ↓
    N × Transformer Blocks, each containing:
        ├── Layer Norm
        ├── Multi-Head Self-Attention  (tokens attend to each other)
        ├── Residual Connection
        ├── Layer Norm
        ├── Feed-Forward Network  (MLP — process each token independently)
        └── Residual Connection
        ↓
    Final Layer Norm
        ↓
    Linear Projection → Vocabulary Logits
        ↓
    Softmax → Next Token Probabilities

Note on positions: the diagram shows "position embedding" as a separate additive step; in **this**
implementation, absolute position is injected inside attention via **RoPE** on Q/K instead,
so `tok_emb` is only token IDs → vectors (no learned positional embedding table).
"""

# --- Imports ---
# `math`: scalar math (sqrt for init scaling, layer count in residual init).
import math
# `torch`: tensors, autograd, device placement (CPU/GPU).
import torch
# `nn.Module`, `nn.Linear`, etc.: building blocks with registered parameters.
import torch.nn as nn
# `F`: stateless ops (softmax, cross_entropy, silu, scaled_dot_product_attention).
import torch.nn.functional as F
# `@dataclass`: auto-generates __init__ so config fields read like a clean struct.
from dataclasses import dataclass


@dataclass
class GPTConfig:
    """
    All hyperparameters for the model in one place.

    Typical shapes implied by these numbers:
    - vocab_size V: token IDs live in [0, V-1]; lm_head maps embd → V logits per position.
    - n_embd C: every token becomes a length-C vector everywhere inside the trunk.
    - n_head H: attention splits C into H heads, each with head_dim = C/H (must divide evenly).
    - n_layer: depth — how many TransformerBlocks are stacked.
    - block_size T_max: max context length this model was laid out for (RoPE cache, masks, crop in generate).
    """
    vocab_size: int = 32000
    n_embd: int = 768
    n_head: int = 12
    n_layer: int = 12
    block_size: int = 512
    dropout: float = 0.1
    bias: bool = False


class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization (RMSNorm).

    Why RMSNorm instead of LayerNorm?
    - RMSNorm is simpler and faster (no mean subtraction, no bias)
    - Used in modern models like Llama, Mistral, etc.
    - Achieves similar or better training stability

    Math: RMSNorm(x) = x / RMS(x) * γ
    where RMS(x) = sqrt(mean(x²) + ε)

    Tensor view:
    - forward(x) with x.shape = (..., dim) keeps all leading axes; only the last axis is normalized.
    - Output matches x.shape exactly; γ broadcasts on the last axis.
    """

    def __init__(self, dim: int, eps: float = 1e-6):
        # `super().__init__` registers this submodule so parameters move with `.to(device)`.
        super().__init__()
        # Tiny constant so we never divide by zero if a feature is all zeros.
        self.eps = eps
        # Learnable scale γ (one per feature). Starts at 1 so initially behavior is pure RMS rescaling.
        # Shape of weight: (dim,) — broadcasts against last axis of any input tensor.
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Normalize the last dimension of `x` to RMS 1, then apply γ.

        Args:
            x: tensor with last dim = `dim`, e.g. (B, T, dim).

        Returns:
            Same shape as x.
        """
        # Input x: arbitrary leading dims ..., final dim must equal `dim` (the normalized axis).
        # Example: x shape (batch=B, seq=T, dim=C) — we normalize each (C,) vector independently.
        # mean(x*x, dim=-1, keepdim=True): average squared value along the feature axis only.
        # Result shape matches x but last dim is 1, e.g. (B, T, 1) — broadcasts cleanly on divide/multiply.
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        # Divide: rescales each vector to RMS ≈ 1 before γ; γ then learns per-feature gain.
        # WHY no centering (subtract mean)? Empirically stable; cheaper than full LayerNorm; used in Llama-style stacks.
        return x / rms * self.weight


class RotaryPositionalEmbedding(nn.Module):
    """
    Rotary Positional Embedding (RoPE).

    Why RoPE instead of learned positional embeddings?
    - Encodes RELATIVE position, not absolute
    - Generalizes better to sequence lengths not seen during training
    - Used in Llama, Mistral, Qwen — the modern standard
    - No learnable parameters (purely mathematical)

    How it works:
    - Pairs up embedding dimensions: (d0, d1), (d2, d3), ...
    - Rotates each pair by an angle proportional to position
    - Nearby tokens get small rotations, far tokens get large rotations
    - The dot product between two rotated vectors naturally encodes distance

    forward(x, seq_len) returns:
    - cos, sin each shaped for broadcasting to q,k in attention, typically (1, 1, seq_len, head_dim).
    """

    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        # `dim` here is head_dim (per-head width), always even in practice — RoPE pairs dimensions.
        # Indices 0,2,4,... up to dim-2: i / dim maps to exponents in [0, 1). Geometric series of wavelengths.
        # base=10000 is the GPT-2/LLama-style default: very low frequencies for early pairs, higher for late pairs.
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        # Buffer, not Parameter: fixed trig table; saved in state_dict but not updated by optimizers.
        # Shape: (dim/2,) — one inverse frequency per 2D rotation plane.
        self.register_buffer("inv_freq", inv_freq)

        # Integer positions 0..max_seq_len-1 as floats for sin/cos.
        t = torch.arange(max_seq_len).float()
        # Outer product: each position t gets a vector of angles t * inv_freq.
        # Shape: (max_seq_len, dim/2) — cache all positions up to max context.
        freqs = torch.outer(t, inv_freq)
        self.register_buffer("cos_cached", freqs.cos())
        self.register_buffer("sin_cached", freqs.sin())

    def forward(self, x: torch.Tensor, seq_len: int):
        """
        Args:
            x: unused in this minimal implementation (some codebases use dtype/device from x).
            seq_len: current sequence length T (≤ cached max).

        Returns:
            cos, sin broadcast tensors for apply_rotary_emb, shape (1, 1, T, head_dim) each.
        """
        # x is only used in some RoPE APIs; here we slice caches to current T = seq_len.
        # cos_cached is (seq_len, dim/2) — need to repeat to match full head_dim
        # Broadcast-friendly for attention: leading 1s for batch and head dims.
        cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0)  # (1, 1, T, dim/2)
        sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)
        # Repeat along last dim: (1, 1, T, dim/2) → (1, 1, T, dim)
        # We duplicate cos/sin because rotate_half() splits the head_dim into two halves;
        # this matches standard implementations that pack angles for both halves of each pair.
        cos = torch.cat([cos, cos], dim=-1)
        sin = torch.cat([sin, sin], dim=-1)
        return cos, sin


def apply_rotary_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """
    Apply rotary embeddings to query and key tensors.

    Takes pairs of dimensions and rotates them:
    [x0, x1] → [x0*cos - x1*sin, x0*sin + x1*cos]

    Inputs / shapes (typical):
    - q, k: (B, H, T, head_dim) after projection and head split — RoPE only touches last axis.
    - cos, sin: broadcastable to q,k — here (1, 1, T, head_dim) so same angles for all batch/heads.
    Outputs: same shapes as q and k, with phase encoding of absolute position folded into q,k.

    WHY rotate q and k (not v)? Attention score is q·k; rotating both by the same angle
    preserves relative geometry we care about (dot products encode relative position via θ_i - θ_j).
    """
    def rotate_half(x):
        # Split last dimension into two halves and swap with sign flip — equivalent form of 2D rotations
        # packed as one matmul-free op. Shape: x (..., D) → same shape out.
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    # (q * cos) + (rotate_half(q) * sin) is the fused implementation of applying rotations per pair.
    q_rotated = (q * cos) + (rotate_half(q) * sin)
    k_rotated = (k * cos) + (rotate_half(k) * sin)
    return q_rotated, k_rotated


class CausalSelfAttention(nn.Module):
    """
    Multi-Head Causal Self-Attention.

    This is the CORE mechanism of a transformer. Here's what it does:

    1. Each token creates three vectors: Query (Q), Key (K), Value (V)
       - Q = "what am I looking for?"
       - K = "what do I contain?"
       - V = "what information do I provide?"

    2. Attention scores = Q @ K^T / sqrt(d_k)
       - Each token's query is compared to all keys
       - Divided by sqrt(d_k) to keep gradients stable

    3. CAUSAL MASK: token at position i can only attend to positions ≤ i
       - This is what makes GPT "autoregressive" (left-to-right)

    4. Output = softmax(scores) @ V
       - Weighted combination of values

    Multi-head: we run this N times in parallel with different learned projections,
    then concatenate. This lets the model attend to different types of relationships.

    Shape recipe (forward):
    - Input x: (B, T, C)
    - After QKV + split + reshape: Q,K,V each (B, H, T, d) with C = H·d
    - SDPA output: (B, H, T, d) → merge heads → (B, T, C) → out_proj → (B, T, C)
    """

    def __init__(self, config: GPTConfig):
        super().__init__()
        # Each head must own an equal slice of the model width C so concat_heads restores C.
        assert config.n_embd % config.n_head == 0, \
            f"n_embd ({config.n_embd}) must be divisible by n_head ({config.n_head})"

        self.n_head = config.n_head
        # head_dim d_k = C / H; this is the per-head key/query width used in attention scores.
        self.head_dim = config.n_embd // config.n_head
        self.n_embd = config.n_embd

        # Single linear layer projects to Q, K, V all at once (more efficient)
        # Weight shape: (3*C, C) if we think y=xW^T in PyTorch Linear — maps each C-vector to 3C outputs.
        # One big matmul on hardware beats three smaller ones (better memory bandwidth).
        self.qkv_proj = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # `out_proj`: mixes information across heads after attention (learned linear recombination).
        # Shape: maps (B,T,C) → (B,T,C) — still "per token", but now heads are fused.
        self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)

        # Dropout on attention probs (implemented inside SDPA when dropout_p>0) — regularizes which keys we listen to.
        self.attn_dropout = nn.Dropout(config.dropout)
        # Dropout on the residual stream after projecting heads — classic transformer regularization.
        self.resid_dropout = nn.Dropout(config.dropout)

        # RoPE table sized to longest sequence in config and to **per-head** dimension (not full C).
        self.rope = RotaryPositionalEmbedding(self.head_dim, config.block_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, C) hidden states **before** this attention layer (already on residual stream).

        Returns:
            (B, T, C) attention output **before** the residual add (dropout applied); block adds residual.
        """
        B, T, C = x.size()  # batch, sequence length, embedding dim
        # Shape: x is (batch=B, seq=T, embd=C) → each token is a C-dim vector on the residual stream.

        # Project to Q, K, V in one shot and reshape for multi-head
        qkv = self.qkv_proj(x)
        # qkv shape: (B, T, 3*C) — contiguous chunks along last dim are Q, K, V each of width C.
        q, k, v = qkv.split(self.n_embd, dim=2)
        # Now q, k, v each: Shape: (B, T, C)

        # Reshape: (B, T, n_embd) → (B, n_head, T, head_dim)
        # view: unpack the C axis as (H, d) with C = H*d; transpose puts heads next to batch for batched matmuls.
        # After transpose — q, k, v: Shape: (B, H, T, d) where d = head_dim.
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # Apply rotary positional embeddings
        # cos/sin: Shape: (1, 1, T, d) — angles depend on token index t, same across batch/heads in usual setup.
        cos, sin = self.rope(q, T)
        q, k = apply_rotary_emb(q, k, cos, sin)
        # q, k now encode absolute position; v unchanged — what we pass through attention as "content".

        # Compute attention using PyTorch's efficient implementation
        # This automatically applies causal mask and is memory-efficient
        # Conceptually: scores = (q @ k^T) / sqrt(d_k) — divide by sqrt(d_k) keeps dot-product magnitude
        # ~O(1) as d_k grows, so softmax doesn't saturate (sharp peaks) and gradients stay healthier.
        # `is_causal=True` masks future keys (typically -inf in logits) so softmax puts ~0 mass on j > i.
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=self.attn_dropout.p if self.training else 0.0,
            is_causal=True,
        )
        # y: Shape: (B, H, T, d) — per head, per token, d-dim attended mix of values.

        # Reassemble all head outputs
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        # After transpose+view — y: Shape: (B, T, C) — concatenated heads back on the last axis.

        # Final projection
        # resid_dropout(out_proj(y)) — Shape: still (B, T, C); this gets added back in the block's residual.
        return self.resid_dropout(self.out_proj(y))


class FeedForward(nn.Module):
    """
    Feed-Forward Network (MLP) with SwiGLU activation.

    After attention lets tokens "talk" to each other, the FFN processes
    each token independently. Think of it as:
    - Attention = inter-token reasoning (relationships between words)
    - FFN = intra-token reasoning (processing within each word's representation)

    SwiGLU (used in Llama, Mistral, modern GPTs):
    - Better than plain ReLU or GELU
    - Uses a gating mechanism: output = Swish(xW₁) ⊙ (xW₃)
    - The ⊙ is element-wise multiplication (gating)
    - Expansion factor is typically 8/3 × n_embd (≈ 2.67×)
    """

    def __init__(self, config: GPTConfig):
        super().__init__()
        # Llama-2 style rounded hidden: ~2.67× wider inner layer than embedding dim for expressivity.
        hidden_dim = int(8 / 3 * config.n_embd)
        # Round up to nearest multiple of 64 for hardware efficiency
        # Tensor cores like fixed multiples of 8/16/32/64 — avoids slow "odd" matmul shapes on GPU.
        hidden_dim = 64 * ((hidden_dim + 63) // 64)

        # gate projection — controls how much of the "value path" passes through (Swish nonlinearity).
        self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=config.bias)  # gate projection
        # down projection — maps widened representation back to residual width C for the next sublayer.
        self.w2 = nn.Linear(hidden_dim, config.n_embd, bias=config.bias)  # down projection
        # up projection — linearly lifts x into the wide space before gating.
        self.w3 = nn.Linear(config.n_embd, hidden_dim, bias=config.bias)  # up projection
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, C) residual-stream activations.

        Returns:
            (B, T, C) FFN output to be added back in the residual.
        """
        # SwiGLU: Swish(x @ W1) * (x @ W3), then project down
        # x: Shape: (B, T, C)
        # self.w1(x), self.w3(x): Shape: (B, T, hidden_dim) — parallel up/gate streams in wide space.
        # silu (Swish): smooth gated ReLU-like; * is elementwise multiply — "GLU" gate.
        # self.w2(...): Shape: (B, T, C) again — FFN output matches residual width for the skip connection.
        # dropout: randomly zero FFN outputs during training to reduce co-adaptation of neurons.
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class TransformerBlock(nn.Module):
    """
    A single Transformer block.

    Structure (Pre-Norm architecture, used in all modern LLMs):

        x ─────────┐
        ↓           │
        RMSNorm     │
        ↓           │
        Attention   │
        ↓           │
        + ←─────────┘  (residual connection)
        │
        ├─────────┐
        ↓         │
        RMSNorm   │
        ↓         │
        FFN       │
        ↓         │
        + ←───────┘    (residual connection)

    Residual connections are critical:
    - They let gradients flow directly back through the network
    - Without them, deep networks (12+ layers) would be nearly impossible to train
    - They create "information highways" through the network
    """

    def __init__(self, config: GPTConfig):
        super().__init__()
        # ln1/ln2 each normalize a C-dim vector per token before the sublayer that follows.
        self.ln1 = RMSNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = RMSNorm(config.n_embd)
        self.ffn = FeedForward(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, C) input to this block.

        Returns:
            (B, T, C) after self-attention + FFN substeps, both with residuals.
        """
        # Pre-norm + residual connections
        # x — Shape: (B, T, C). The residual path `x` is preserved and only **adds** sublayer outputs.
        # WHY residual: output = x + f(norm(x)); gradients can flow through the + with factor ~1,
        # mitigating vanishing gradients in deep stacks; also lets layers learn small deltas around identity.
        x = x + self.attn(self.ln1(x))
        # Same pattern: normalize → MLP → add back. Subblocks never need to relearn the full representation.
        x = x + self.ffn(self.ln2(x))
        return x


class GPT(nn.Module):
    """
    The complete GPT language model.

    Putting it all together:
    1. Token embeddings: each token ID → a learned vector
    2. Stack of N transformer blocks
    3. Final layer norm
    4. Linear head: project back to vocabulary size for next-token prediction

    Weight tying: the token embedding matrix is shared with the output projection.
    This means the model uses the same "understanding" of words for both
    input and output. Saves parameters and often improves performance.

    WHY weight tying saves params: normally you'd have
    - Embedding E: (V, C) for token → vector
    - Separate head W: (V, C) for vector → logits
    Tying reuses the **same** (V, C) matrix for both roles → exactly V×C fewer weights
    and a useful inductive bias: input and output token live in the same space.
    """

    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config

        # ModuleDict just names submodules for readable forward code (string keys).
        self.transformer = nn.ModuleDict({
            # tok_emb.weight will alias lm_head.weight below — one matrix, two interpretations.
            "tok_emb": nn.Embedding(config.vocab_size, config.n_embd),
            "drop": nn.Dropout(config.dropout),
            "blocks": nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)]),
            "ln_f": RMSNorm(config.n_embd),
        })
        # Linear with bias=False is common with tied embeddings — the bias would break symmetry.
        # Weight shape internally: (vocab_size, n_embd) in nn.Linear (out_features × in_features).
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Weight tying: share embedding weights with output projection
        # Both `Embedding` and this `Linear` expose a `.weight` of shape (V, C).
        # logits[token] ∝ emb[token_id] · h — like a big softmax classifier sharing class embeddings.
        self.transformer["tok_emb"].weight = self.lm_head.weight

        # Initialize weights
        # `_init_weights` runs once per submodule type (Linear, Embedding, ...)
        self.apply(self._init_weights)
        # Scale residual projections by 1/sqrt(2*n_layer) for training stability
        # Two residual adds per layer → variance grows ~linearly with depth unless we shrink init of output-like layers.
        for name, p in self.named_parameters():
            if name.endswith("out_proj.weight") or name.endswith("w2.weight"):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module: nn.Module):
        """
        Weight initialization matters enormously for training stability.
        - Linear layers: normal distribution with std=0.02
        - Embeddings: same normal distribution
        - Biases: zero
        """
        if isinstance(module, nn.Linear):
            # 0.02 std is a historic Transformer/GPT default — small enough for deep stacks at init.
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            # Note: after tying, `tok_emb` shares `lm_head.weight` — initializing Embedding
            # hits this branch first in traversal order; later Linear init also touches the same tensor.
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def count_parameters(self) -> int:
        """Count total trainable parameters."""
        # numel = product of shape; requires_grad filters frozen layers (none here, but common in fine-tuning).
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def forward(self, idx: torch.Tensor, targets: torch.Tensor = None):
        """
        Forward pass.

        Args:
            idx: token indices, shape (batch_size, sequence_length)
            targets: target token indices for computing loss (shifted by 1)

        Returns:
            logits: next-token predictions, shape (batch_size, sequence_length, vocab_size)
                Each position t predicts the token at t+1 when trained with shifted targets.
            loss: cross-entropy loss if targets provided, else None
        """
        _, T = idx.size()
        # idx — Shape: (batch, seq=T) — integer token IDs in [0, vocab_size-1].
        assert T <= self.config.block_size, \
            f"Sequence length {T} exceeds block_size {self.config.block_size}"

        # Token embeddings (no positional embeddings — RoPE handles position in attention)
        x = self.transformer["tok_emb"](idx)
        # Shape: (batch=2, seq=1024, embd=768) style — each token is now a C-dim learned vector.
        x = self.transformer["drop"](x)
        # Dropout randomly zeros whole embedding dims during training — regularizes early layers.

        # Pass through all transformer blocks
        for block in self.transformer["blocks"]:
            x = block(x)
        # Still Shape: (B, T, C) — residual stream width never changes inside the trunk.

        # Final layer norm
        x = self.transformer["ln_f"](x)
        # Stabilizes scale before logits; mirrors "post trunk" norm used in GPT-2 style, but here last RMSNorm.

        # Project to vocabulary
        logits = self.lm_head(x)
        # logits — Shape: (B, T, V) — unnormalized scores for each vocab entry (logits → softmax → probs).

        # Compute loss if targets provided
        loss = None
        if targets is not None:
            # targets — Shape: (B, T), same as idx; usually idx[:,1:] is predicted from logits[:,:-1].
            # cross_entropy expects (N, classes) and (N,) targets — flatten BT → one big classification batch.
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,  # ignore padding tokens
            )
            # CE = -log p(correct class); averaging over non-ignored positions reduces to mean loss.

        return logits, loss

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int,
                 temperature: float = 0.8, top_k: int = 50) -> torch.Tensor:
        """
        Autoregressive text generation.

        How generation works:
        1. Feed the prompt through the model
        2. Get probabilities for the next token
        3. Sample from the distribution (with temperature and top-k)
        4. Append the sampled token
        5. Repeat

        Temperature controls randomness:
        - temp < 1.0: more deterministic (sharper distribution)
        - temp = 1.0: standard sampling
        - temp > 1.0: more random (flatter distribution)

        Top-k sampling: only consider the k most likely tokens.
        Prevents the model from choosing extremely unlikely tokens.

        Args:
            idx: starting token IDs, shape (batch, current_length) — grows by 1 each iteration.
            max_new_tokens: how many sampling steps to run.
            temperature: softmax temperature (scalar) applied to last-position logits.
            top_k: if set, restrict sampling mass to the k largest logits (per batch row).

        Returns:
            idx: shape (batch, current_length + max_new_tokens) — prompt with new tokens appended.
        """
        self.eval()
        # Disable dropout / batchnorm-like behavior — deterministic forward for sampling.
        for _ in range(max_new_tokens):
            # Crop to block_size if sequence is too long
            # Only the last T_max tokens fit in positional/rope cache and learned context.
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]

            logits, _ = self(idx_cond)
            # Use **last** time step's logits: we're predicting the *next* token after the visible prefix.
            logits = logits[:, -1, :]
            # Shape: (B, V) — one next-token distribution per batch row.
            # Divide by temperature **after** forward: T<1 sharpens logits (less random), T>1 flattens (more random).
            logits = logits / temperature

            if top_k is not None:
                # Keep only top-k logits per row; push the rest to -inf so softmax assigns ~0 mass there.
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                # v[:, [-1]] is the k-th largest logit — threshold per batch element.
                logits[logits < v[:, [-1]]] = float("-inf")

            probs = F.softmax(logits, dim=-1)
            # probs — Shape: (B, V), rows sum to 1 — categorical distribution for next token.
            idx_next = torch.multinomial(probs, num_samples=1)
            # Sample one integer id per batch row according to probs — stochastic decoding.
            idx = torch.cat([idx, idx_next], dim=1)
            # Append along time — idx grows to (B, T + num_generated_so_far).

        return idx