model.py

"""
Stage 4 — model.py

A full encoder-decoder transformer, in the shape of "Attention Is All You Need".

Design choice: the encoder and decoder *layers* are hand-written so you can see
exactly which tensor is the query and which are the keys/values in each of the
three attention calls — that is the whole point of this project. The attention
operation itself uses nn.MultiheadAttention rather than a hand-rolled
softmax(QK^T/sqrt(d))V, because that part you have already studied and it is
where subtle numerical bugs hide. Swap it out later if you want to write the
scaled dot-product by hand; the surrounding code will not change.

THE THREE ATTENTION CALLS
-------------------------
1. Encoder self-attention        Q,K,V = source        NOT causal (bidirectional)
2. Decoder self-attention        Q,K,V = target        CAUSAL
3. Decoder cross-attention       Q     = target
                                 K,V   = encoder output   NOT causal

Call 3 is the new piece. The decoder asks, at each target position, "which
source tokens are relevant to the word I am about to produce?" — so the query
comes from the decoder and the keys/values come from the encoder's final
output (`memory`). It is NOT causally masked in either direction: target
position 0 is allowed to look at the whole source sentence. Only *target*
positions are restricted, and only in call 2.

MASK CONVENTIONS (kept identical everywhere on purpose)
-------------------------------------------------------
  key_padding_mask : (B, S) bool. True = that key position is <pad>, ignore it.
  attn_mask        : (T, T) bool. True = that (query, key) pair is FORBIDDEN.

Both are "True means blocked". PyTorch uses this polarity for
nn.MultiheadAttention; flipping it silently trains a model that attends only
to padding, which does not crash and does not produce NaNs — it just produces
garbage. Hence: one convention, stated once, used everywhere.

Padding must be masked on BOTH sides:
  - source padding -> encoder self-attention keys  (call 1)
  - source padding -> cross-attention keys         (call 3)   <- easy to forget
  - target padding -> decoder self-attention keys  (call 2)

Pre-LN vs post-LN: the 2017 paper puts LayerNorm *after* each residual
addition. Pre-LN (norm on the branch input, as written below) is what most
modern implementations use because it trains stably without a carefully tuned
warmup. We still use warmup, but this makes the model far more forgiving if
you change the learning rate. Each stack gets a final LayerNorm, which pre-LN
requires.

============================================================================
THIS FILE'S PLACE IN THE CHAIN
============================================================================
model.py has ONE project-internal dependency: config.py (for ModelConfig).
It does not know about tokenizers, datasets, or files on disk — it is pure
architecture, operating on already-tokenized integer id tensors.

Who calls into this file:
  - train.py's main(): cfg = ModelConfig(...); model = build_model(cfg, pad_id, device)
    then, every step, run_epoch() calls model(src, tgt_in, ...) which routes
    to TransformerTranslator.forward() below.
  - decoding.py's greedy_decode()/beam_search_decode(): call model.encode(...)
    ONCE, then model.decode(...) repeatedly in a loop — this is why encode()
    and decode() are exposed as separate methods instead of only forward().
  - evaluate.py's main(): build_model(cfg, ...) again from a saved
    ModelConfig, then model.load_state_dict(ckpt["model"]) to restore trained weights.
  - sanity_checks.py: imports build_model and causal_mask directly to probe
    the model's masking behavior with hand-crafted inputs.

Internal call chain within this file (top to bottom of what actually runs
during a training step):
  TransformerTranslator.forward(src, tgt_in, ...)
    -> self.encode(src, ...)
         -> self.embed(src) * sqrt(d_model)          [token ids -> vectors]
         -> self.pos_enc(x)                          [add position information]
         -> for each layer in self.encoder_layers: EncoderLayer.forward(x, ...)
              -> ATTENTION CALL 1 (self.self_attn, not masked causally)
              -> self.ff(x)  (FeedForward)
         -> self.enc_norm(x)                         [final pre-LN norm]
         => returns `memory`, shape (B, S, d)
    -> self.decode(tgt_in, memory, ...)
         -> self.embed(tgt_in) * sqrt(d_model)        [SAME embedding table as above]
         -> self.pos_enc(x)
         -> cm = causal_mask(T, device)                [built once per decode() call]
         -> for each layer in self.decoder_layers: DecoderLayer.forward(x, memory, cm, ...)
              -> ATTENTION CALL 2 (self.self_attn, WITH causal_mask)
              -> ATTENTION CALL 3 (self.cross_attn, query=decoder, key/value=memory)
              -> self.ff(x)
         -> self.dec_norm(x)
         -> self.output_proj(x)                       [-> logits over the vocab]
"""

import math

import torch
import torch.nn as nn

from config import ModelConfig


# ---------------------------------------------------------------------------
# Positional encodings
# ---------------------------------------------------------------------------
class SinusoidalPositionalEncoding(nn.Module):
    """Fixed sin/cos encodings from the original paper.

    PE[pos, 2i]   = sin(pos / 10000^(2i/d))
    PE[pos, 2i+1] = cos(pos / 10000^(2i/d))

    No parameters, and it extrapolates (sort of) to lengths never seen in
    training — which matters here because translation output length is not
    bounded by anything the model saw.

    Instantiated once inside TransformerTranslator.__init__ (via
    build_positional_encoding below) and called from BOTH encode() and
    decode() — the SAME module instance handles source and target positions,
    since "position 3" means the same thing (a fixed sin/cos vector) in
    either sequence.
    """

    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        # Precompute the ENTIRE table up to max_len positions, once, at
        # construction time — forward() below just slices into it.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(max_len, dtype=torch.float).unsqueeze(1)   # (max_len, 1)
        # Computed in log space for numerical stability.
        # div[i] = 10000^(-2i/d_model), one value per even dimension index.
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # Broadcasting: position (max_len,1) * div (d_model/2,) -> (max_len, d_model/2)
        pe[:, 0::2] = torch.sin(position * div)   # even dims: sin
        pe[:, 1::2] = torch.cos(position * div)   # odd dims: cos
        # register_buffer: saved with the model, moved by .to(device), not a parameter.
        # (i.e. it's part of state_dict()/checkpoint, and follows model.to(device),
        #  but torch.optim never updates it — it's fixed math, not learned.)
        self.register_buffer("pe", pe.unsqueeze(0))                        # (1, max_len, d)

    def forward(self, x):                    # x: (B, L, d)
        # Slice the precomputed table down to this sequence's actual length L,
        # broadcast-add across the batch dimension, then apply dropout.
        # Called from TransformerTranslator.encode()/decode() right after the
        # scaled embedding lookup — see the module docstring's call chain.
        return self.dropout(x + self.pe[:, : x.size(1)])


class LearnedPositionalEncoding(nn.Module):
    """An ordinary embedding table indexed by position. Ablation 6.

    Selected via ModelConfig.pos_encoding="learned" instead of "sinusoidal" —
    see build_positional_encoding() below, which is the only place that
    chooses between this class and SinusoidalPositionalEncoding above.
    Unlike the sinusoidal version, this one has trainable parameters
    (self.emb.weight) and CANNOT extrapolate past max_len (see the assert).
    """

    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.emb = nn.Embedding(max_len, d_model)   # one learned vector per position 0..max_len-1
        nn.init.normal_(self.emb.weight, mean=0.0, std=0.02)
        self.max_len = max_len

    def forward(self, x):
        L = x.size(1)
        assert L <= self.max_len, f"sequence of {L} exceeds learned max_len {self.max_len}"
        # arange(L) -> position indices [0, 1, ..., L-1], broadcast to batch
        # dim 1 so nn.Embedding can look each one up.
        pos = torch.arange(L, device=x.device).unsqueeze(0)                # (1, L)
        return self.dropout(x + self.emb(pos))


def build_positional_encoding(kind, d_model, max_len, dropout):
    """Factory function: picks which positional-encoding class to instantiate
    based on ModelConfig.pos_encoding. Called exactly once, from
    TransformerTranslator.__init__, to build self.pos_enc."""
    if kind == "sinusoidal":
        return SinusoidalPositionalEncoding(d_model, max_len, dropout)
    if kind == "learned":
        return LearnedPositionalEncoding(d_model, max_len, dropout)
    raise ValueError(f"unknown pos_encoding: {kind}")


# ---------------------------------------------------------------------------
# Layers
# ---------------------------------------------------------------------------
class FeedForward(nn.Module):
    """Position-wise FFN: d_model -> d_ff -> d_model, applied identically at
    every position (it is where most of the parameters live).

    "Position-wise" means: the SAME two Linear layers are applied
    independently to every one of the L positions in (B, L, d_model) — there
    is no mixing ACROSS positions here (that's what attention is for). This
    is why nn.Linear works directly on a 3D tensor: PyTorch applies it to the
    last dimension only, batching over everything else automatically.

    Instantiated once per EncoderLayer and once per DecoderLayer below
    (self.ff), so with 4+4 default layers there are 8 independent
    FeedForward modules, each with its own weights.
    """

    def __init__(self, d_model, d_ff, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),   # expand
            nn.ReLU(),                  # nonlinearity — this is what makes the FFN
                                        # do more than a linear projection
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),   # project back down to d_model so residual add works
        )

    def forward(self, x):
        # Called from EncoderLayer.forward() and DecoderLayer.forward(), each
        # time wrapped in a residual connection: x = x + self.ff(norm(x)).
        return self.net(x)


class EncoderLayer(nn.Module):
    """self-attention (bidirectional) -> feed-forward, both residual + pre-LN.

    Stacked cfg.n_encoder_layers times inside TransformerTranslator.__init__
    (self.encoder_layers). TransformerTranslator.encode() runs input `x`
    through each of these in sequence, output of one feeding the next.
    """

    def __init__(self, d_model, n_heads, d_ff, dropout):
        super().__init__()
        # nn.MultiheadAttention does the actual softmax(QK^T/sqrt(d))V math
        # (scaled dot-product attention across n_heads parallel heads) —
        # see the module docstring for why this project doesn't hand-roll it.
        # batch_first=True: tensors are (batch, seq, feature), matching every
        # other shape convention in this file (the PyTorch default is
        # (seq, batch, feature), which would be inconsistent here).
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout,
                                               batch_first=True)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)   # pre-LN before the attention sublayer
        self.norm2 = nn.LayerNorm(d_model)   # pre-LN before the feed-forward sublayer
        self.drop = nn.Dropout(dropout)      # applied to each sublayer's OUTPUT before the residual add

    def forward(self, x, src_key_padding_mask=None):
        # x: (B, S, d)  — called from TransformerTranslator.encode(), once per
        # layer in self.encoder_layers, x is this layer's input AND (after
        # this function returns) the next layer's input.

        # ATTENTION CALL 1: query = key = value = the source itself.
        # No attn_mask -> every source token sees every other source token in
        # both directions. That is the point of the encoder.
        h = self.norm1(x)                    # pre-LN: normalize BEFORE feeding into attention
        attn, _ = self.self_attn(h, h, h,     # (query, key, value) — all three are `h`
                                 key_padding_mask=src_key_padding_mask,
                                 # ^ (B, S) bool, True = ignore this source
                                 #   position (it's <pad>) — this is
                                 #   dataset.py's "src_key_padding_mask" batch
                                 #   key, passed down unchanged from
                                 #   TransformerTranslator.forward()/encode().
                                 need_weights=False)   # we never inspect attention
                                                        # weights, so skip computing them (speed)
        x = x + self.drop(attn)              # residual connection: original x + (dropped) attention output

        h = self.norm2(x)                    # pre-LN before the FFN sublayer
        x = x + self.drop(self.ff(h))        # second residual connection
        return x


class DecoderLayer(nn.Module):
    """causal self-attention -> cross-attention -> feed-forward.

    Stacked cfg.n_decoder_layers times inside TransformerTranslator.__init__
    (self.decoder_layers). TransformerTranslator.decode() runs decoder state
    `x` AND the fixed `memory` (encoder output) through each of these in
    sequence — `memory` itself is NOT modified by the decoder stack, only
    `x` is.
    """

    def __init__(self, d_model, n_heads, d_ff, dropout):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout,
                                               batch_first=True)
        # ^ ATTENTION CALL 2's engine — target attends to target, causally masked below.
        self.cross_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout,
                                                batch_first=True)
        # ^ ATTENTION CALL 3's engine — target attends to encoder memory.
        #   A SEPARATE nn.MultiheadAttention instance from self_attn above —
        #   different learned weights, even though the shapes look similar.
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)   # before self-attention
        self.norm2 = nn.LayerNorm(d_model)   # before cross-attention
        self.norm3 = nn.LayerNorm(d_model)   # before feed-forward
        self.drop = nn.Dropout(dropout)

    def forward(self, x, memory, causal_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        # x      : (B, T, d)  decoder states  (the target so far)
        # memory : (B, S, d)  encoder output  (the whole source, always)
        # Called from TransformerTranslator.decode(), once per layer in
        # self.decoder_layers, with the SAME `memory` tensor passed to every
        # layer (memory is computed once by encode() and reused).

        # ATTENTION CALL 2 — decoder self-attention, CAUSAL.
        # attn_mask is the (T, T) upper-triangular block: position i may not
        # look at any j > i. Without it, the model can read the answer during
        # teacher forcing and will look perfect in training and useless at
        # inference, where the future genuinely does not exist yet.
        h = self.norm1(x)
        sa, _ = self.self_attn(h, h, h,       # query=key=value=decoder states
                               attn_mask=causal_mask,
                               # ^ (T, T) bool — built by the standalone
                               #   causal_mask() function (below in this
                               #   file) and passed down from
                               #   TransformerTranslator.decode(). Same
                               #   causal_mask tensor is reused for every
                               #   decoder layer within one decode() call.
                               key_padding_mask=tgt_key_padding_mask,
                               # ^ (B, T) bool — masks <pad> positions WITHIN
                               #   the target itself (relevant during
                               #   training with a padded batch; None during
                               #   greedy/beam generation, where there are no
                               #   pads yet in the growing prefix).
                               need_weights=False)
        x = x + self.drop(sa)

        # ATTENTION CALL 3 — CROSS-ATTENTION. The new piece.
        #   query  <- decoder states  (T positions: "what am I writing?")
        #   key    <- encoder memory  (S positions: "what's available?")
        #   value  <- encoder memory  (S positions: "what do I copy across?")
        # Deliberately NO attn_mask: target position 0 may attend to the whole
        # source, including its last word. Word order differs between Russian
        # and English, so restricting this would make translation impossible.
        # key_padding_mask IS passed: <pad> in the source must not be attended
        # to. Forgetting it does not crash — it quietly costs a few BLEU.
        h = self.norm2(x)
        ca, _ = self.cross_attn(query=h, key=memory, value=memory,
                                key_padding_mask=memory_key_padding_mask,
                                # ^ (B, S) bool — this is the SAME
                                #   src_key_padding_mask used in encoder call
                                #   1, just renamed "memory_key_padding_mask"
                                #   here because from the decoder's point of
                                #   view it's masking the KEYS (memory), not
                                #   its own sequence.
                                need_weights=False)
        x = x + self.drop(ca)

        h = self.norm3(x)
        x = x + self.drop(self.ff(h))
        return x


# ---------------------------------------------------------------------------
# Full model
# ---------------------------------------------------------------------------
def causal_mask(size, device):
    """(size, size) bool, True = forbidden. Row i has True for all columns j>i.

        [[F, T, T],
         [F, F, T],
         [F, F, F]]

    Called from TransformerTranslator.decode() every time it runs (built
    fresh from the CURRENT target length T each call — cheap, so no need to
    cache it). Also called directly by sanity_checks.py's
    test_causal_mask_shape() to verify this exact shape/polarity in isolation.

    torch.triu(..., diagonal=1): keeps only the STRICTLY upper triangle
    (diagonal=1 excludes the main diagonal itself), which is exactly "j > i".
    That's why position i CAN attend to itself (j == i is False, i.e. allowed)
    but not to anything after it.
    """
    return torch.triu(torch.ones(size, size, dtype=torch.bool, device=device),
                      diagonal=1)


class TransformerTranslator(nn.Module):
    """The whole model. Constructed via build_model() below (never
    instantiated directly outside this file/sanity_checks.py)."""

    def __init__(self, cfg: ModelConfig, pad_id=0):
        super().__init__()
        self.cfg = cfg
        self.pad_id = pad_id
        self.d_model = cfg.d_model

        # ONE embedding table, shared by:
        #   - the encoder input (Russian tokens)
        #   - the decoder input (English tokens)
        #   - the output projection (via weight tying, below)
        # This is only possible because the BPE vocabulary is shared.
        # padding_idx=pad_id: tells nn.Embedding that row `pad_id` should
        # never receive a gradient update (kept pinned at whatever it's
        # initialized to — see the explicit zeroing right below).
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=pad_id)
        nn.init.normal_(self.embed.weight, mean=0.0, std=cfg.d_model ** -0.5)
        # ^ std=d_model^-0.5 matches the scale expected by the
        #   "* sqrt(d_model)" multiplication done in encode()/decode() below
        #   — see the comment there for why.
        with torch.no_grad():
            self.embed.weight[pad_id].zero_()   # belt-and-suspenders: <pad>'s row starts at exactly 0

        # Builds either SinusoidalPositionalEncoding or
        # LearnedPositionalEncoding depending on cfg.pos_encoding — see
        # build_positional_encoding() above.
        self.pos_enc = build_positional_encoding(cfg.pos_encoding, cfg.d_model,
                                                 cfg.max_len, cfg.dropout)

        # nn.ModuleList: a Python-list-like container that PyTorch still
        # recognizes as holding submodules (a plain Python list would NOT
        # register these layers' parameters, and .to(device)/state_dict()
        # would silently miss them).
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout)
            for _ in range(cfg.n_encoder_layers)])
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout)
            for _ in range(cfg.n_decoder_layers)])

        # Pre-LN needs a final norm at the top of each stack, otherwise the
        # output of the last residual branch is never normalized.
        self.enc_norm = nn.LayerNorm(cfg.d_model)   # applied at the end of encode()
        self.dec_norm = nn.LayerNorm(cfg.d_model)   # applied at the end of decode(), before output_proj

        self.output_proj = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        # ^ turns each decoder position's d_model-dim vector into a
        #   vocab_size-dim vector of LOGITS (unnormalized scores per token) —
        #   this is what decode()'s return value actually is.
        if cfg.tie_embeddings:
            # Weight tying: the output layer reuses the embedding matrix, so
            # logit[v] = <decoder_state, embedding[v]>. Saves vocab*d_model
            # parameters (16000*256 = 4.1M here, a large share of the model)
            # and usually helps a low-resource model generalize.
            # NOTE: this line REPLACES self.output_proj.weight with a
            # reference to the SAME tensor object as self.embed.weight — from
            # this point on, training self.embed.weight via backprop also
            # changes self.output_proj's behavior, and vice versa, because
            # they are literally one tensor with two names.
            self.output_proj.weight = self.embed.weight

    # -- encoder ------------------------------------------------------------
    def encode(self, src, src_key_padding_mask=None):
        """src: (B, S) -> memory: (B, S, d)

        Called from: TransformerTranslator.forward() below (training path),
        AND directly by decoding.py's greedy_decode()/beam_search_decode()
        (generation path) — those call it ONCE per source sentence, then
        reuse the returned `memory` across every decode() step, which is the
        whole computational payoff of splitting encoder/decoder into
        separate methods instead of only exposing forward().
        """
        # Scale by sqrt(d_model) as in the paper: embeddings are initialized
        # with std d^-0.5, so this puts them on the same scale as the
        # positional encoding (which has unit-ish amplitude) rather than being
        # swamped by it.
        x = self.embed(src) * math.sqrt(self.d_model)   # (B, S) ids -> (B, S, d) vectors, rescaled
        x = self.pos_enc(x)                              # add positional information, (B, S, d)
        for layer in self.encoder_layers:
            x = layer(x, src_key_padding_mask=src_key_padding_mask)
            # ^ EncoderLayer.forward() — see that class for ATTENTION CALL 1.
            #   Output of layer i becomes input to layer i+1.
        return self.enc_norm(x)   # final pre-LN normalization required by the pre-LN convention

    # -- decoder ------------------------------------------------------------
    def decode(self, tgt_in, memory, memory_key_padding_mask=None,
               tgt_key_padding_mask=None):
        """tgt_in: (B, T) decoder INPUT (already shifted, starts with <bos>)
        memory:  (B, S, d) -> logits (B, T, vocab)

        Called from: TransformerTranslator.forward() below (training path,
        ONE call covering the whole target sequence at once thanks to
        teacher forcing + the causal mask), AND from decoding.py's
        greedy_decode()/beam_search_decode() in a LOOP — each generation
        step re-runs decode() over the WHOLE prefix generated so far (see
        decoding.py's module docstring "Efficiency note" for why this is
        O(T^2) and deliberately not KV-cached).
        """
        T = tgt_in.size(1)
        x = self.embed(tgt_in) * math.sqrt(self.d_model)   # SAME embedding table as encode() used
        x = self.pos_enc(x)                                  # SAME positional-encoding module as encode()
        cm = causal_mask(T, tgt_in.device)                   # (T, T), built fresh each call — see causal_mask() above
        for layer in self.decoder_layers:
            x = layer(x, memory,
                      causal_mask=cm,
                      tgt_key_padding_mask=tgt_key_padding_mask,
                      memory_key_padding_mask=memory_key_padding_mask)
            # ^ DecoderLayer.forward() — see that class for ATTENTION CALLS 2 & 3.
            #   `memory` is passed unchanged to every layer; only `x` accumulates.
        x = self.dec_norm(x)          # final pre-LN normalization
        return self.output_proj(x)    # (B, T, d) -> (B, T, vocab_size) logits

    # -- both ---------------------------------------------------------------
    def forward(self, src, tgt_in, src_key_padding_mask=None,
                tgt_key_padding_mask=None):
        """Training forward pass.

        src    : (B, S) source token ids
        tgt_in : (B, T) decoder input = target shifted right (<bos> w1 w2 w3)
        returns: (B, T, vocab) logits, where logits[:, i] predicts the token
                 that should follow tgt_in[:, i]

        Called from train.py's run_epoch()/validate_loss() as
        `logits = model(src, tgt_in, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)`
        — this is the ONLY place forward() is used; generation (decoding.py)
        calls encode()/decode() directly instead, precisely to avoid
        recomputing `memory` at every generation step (see encode()'s docstring).
        """
        memory = self.encode(src, src_key_padding_mask)
        return self.decode(tgt_in, memory,
                           memory_key_padding_mask=src_key_padding_mask,
                           # ^ reuses the SAME mask tensor for both "encoder
                           #   self-attn keys" (inside encode(), call 1) and
                           #   "cross-attn keys" (here, call 3) — both are
                           #   masking the same source padding pattern.
                           tgt_key_padding_mask=tgt_key_padding_mask)

    def num_parameters(self, trainable_only=True):
        """Called by train.py/evaluate.py purely for the printed
        "X trainable parameters" log line."""
        ps = self.parameters()
        return sum(p.numel() for p in ps if p.requires_grad or not trainable_only)


def build_model(cfg: ModelConfig, pad_id=0, device="cpu"):
    """Thin factory wrapping TransformerTranslator(cfg).to(device).

    Called from: train.py's main(), evaluate.py's main(), sanity_checks.py's
    _model(), and this file's own __main__ block below. Centralizing
    construction+device-placement here means every caller gets identical
    behavior instead of each remembering to call .to(device) itself.
    """
    model = TransformerTranslator(cfg, pad_id=pad_id).to(device)
    return model


# ---------------------------------------------------------------------------
# Quick self-check: python model.py
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # Runs only when this file is executed directly. Builds a tiny toy model
    # (small vocab/d_model so it's instant) and runs one forward pass with
    # RANDOM token ids (no real tokenizer/data needed) just to confirm the
    # architecture is wired correctly end to end — shapes match, weight
    # tying is really the same tensor. For actual correctness properties
    # (causality, mask polarity) see sanity_checks.py instead.
    torch.manual_seed(0)
    cfg = ModelConfig(vocab_size=500, d_model=64, n_heads=4,
                      n_encoder_layers=2, n_decoder_layers=2, d_ff=128, dropout=0.0)
    m = build_model(cfg)
    m.eval()          # disables dropout for a deterministic check
    B, S, T = 3, 7, 5
    src = torch.randint(4, 500, (B, S))     # random "sentences", ids 4..499 (avoiding specials 0-3)
    tgt_in = torch.randint(4, 500, (B, T))
    logits = m(src, tgt_in)                  # exercises forward() -> encode() -> decode()
    print("logits:", tuple(logits.shape), "(expected", (B, T, cfg.vocab_size), ")")
    print("parameters:", f"{m.num_parameters():,}")
    print("embedding tied to output:", m.output_proj.weight is m.embed.weight)
    # ^ `is` checks OBJECT IDENTITY, not just equal values — confirms the
    #   weight-tying assignment in __init__ really did share one tensor.

    print("\nRun `python sanity_checks.py` for the causality / masking tests.")