S111 · Full Transformer Architecture — Encoder + Decoder

# tiny_gpt.py — a from-scratch decoder-only Transformer.
# Trains on a character-level corpus, then generates.
# Run: uv run tiny_gpt.py path/to/text.txt
import math
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BLOCK_SIZE = 128
BATCH = 32
D_MODEL = 128
N_HEADS = 4
N_LAYERS = 4
D_FF = 4 * D_MODEL
DROPOUT = 0.1
LR = 3e-4
EPOCHS = 5
 
 
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(DROPOUT)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        # One big projection, then split — faster than 3 separate matmuls.
        qkv = self.qkv(x)                              # [B, T, 3D]
        Q, K, V = qkv.split(D, dim=-1)
        Q = Q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = K.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = V.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        # PyTorch's fused scaled_dot_product_attention with causal mask.
        out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)  # [B, N, T, d_head]
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        return self.dropout(self.proj(out))
 
 
class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(DROPOUT)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.fc2(F.gelu(self.fc1(x))))
 
 
class Block(nn.Module):
    """One Transformer block — pre-LN variant."""
 
    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN: normalise BEFORE sublayer, add residual AFTER.
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x
 
 
class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, block_size: int, d_model: int,
                 n_heads: int, n_layers: int, d_ff: int):
        super().__init__()
        self.block_size = block_size
        self.tok_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(block_size, d_model)  # learned PE for simplicity
        self.blocks = nn.ModuleList([
            Block(d_model, n_heads, d_ff) for _ in range(n_layers)
        ])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        # Weight tying — share weights between token embedding and output head.
        self.head.weight = self.tok_embed.weight
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T = x.shape
        pos = torch.arange(T, device=x.device)
        h = self.tok_embed(x) + self.pos_embed(pos)     # [B, T, d]
        for block in self.blocks:
            h = block(h)
        h = self.ln_f(h)
        return self.head(h)                              # [B, T, vocab]
 
    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new: int, temperature: float = 1.0) -> torch.Tensor:
        for _ in range(max_new):
            idx_cond = idx[:, -self.block_size:]         # crop to context window
            logits = self(idx_cond)[:, -1, :]            # last token's logits
            probs = F.softmax(logits / temperature, dim=-1)
            next_tok = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_tok], dim=1)
        return idx
 
 
def get_batches(data: torch.Tensor, block_size: int, batch_size: int) -> DataLoader:
    n = (len(data) - 1) // block_size
    x = data[: n * block_size].view(n, block_size)
    y = data[1 : n * block_size + 1].view(n, block_size)
    return DataLoader(TensorDataset(x, y), batch_size=batch_size, shuffle=True)
 
 
def main(path: str) -> None:
    text = open(path, "r", encoding="utf-8").read()
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    itos = {i: c for c, i in stoi.items()}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    print(f"vocab={len(chars)} chars={len(text):,} device={DEVICE}")
 
    loader = get_batches(data, BLOCK_SIZE, BATCH)
    model = TinyGPT(len(chars), BLOCK_SIZE, D_MODEL, N_HEADS, N_LAYERS, D_FF).to(DEVICE)
    params = sum(p.numel() for p in model.parameters())
    print(f"params: {params:,}  ({params/1e6:.2f}M)")
    opt = torch.optim.AdamW(model.parameters(), lr=LR)
 
    for epoch in range(1, EPOCHS + 1):
        model.train()
        total = 0.0
        for step, (x, y) in enumerate(loader, 1):
            x, y = x.to(DEVICE), y.to(DEVICE)
            logits = model(x)
            loss = F.cross_entropy(logits.reshape(-1, len(chars)), y.reshape(-1))
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total += loss.item()
        # Generate a sample.
        model.eval()
        seed = torch.tensor([[stoi[text[0]]]], device=DEVICE)
        out = model.generate(seed, max_new=200, temperature=0.8)
        sample = "".join(itos[i.item()] for i in out[0])
        print(f"epoch {epoch}  loss={total/step:.3f}")
        print(f"  sample: {sample[:100]!r}...")
 
 
if __name__ == "__main__":
    main(sys.argv[1])