Add model source

src/tiny_transformer.egg-info/PKG-INFO

@@ -0,0 +1,153 @@
+Metadata-Version: 2.4
+Name: tiny-transformer
+Version: 0.1.0
+Summary: A compact GPT-style Transformer built from scratch in PyTorch.
+Author: Akshay Mulgund
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=2.0
+Requires-Dist: torch>=2.2
+Requires-Dist: tqdm>=4.66
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+
+# Tiny Transformer
+
+A compact GPT-style language model built from scratch in PyTorch. This repo is designed to show the fundamentals recruiters actually care about: clean architecture, readable math, reproducible training, tests, and an end-to-end demo path from raw text to generated tokens.
+
+## What Makes This Worth Looking At
+
+- Implements a decoder-only Transformer without Hugging Face or high-level training frameworks.
+- Includes causal self-attention, multi-head attention, residual blocks, layer norm, embeddings, generation, and checkpointing.
+- Ships with character and byte-pair encoding tokenizers so the model can train on any plain-text file.
+- Keeps the code small enough to understand in one sitting, but structured like production Python.
+- Includes smoke tests for masking, shapes, tokenization, attention export, and generation behavior.
+
+## Quickstart
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+python -m pip install -e ".[dev]"
+```
+
+Train on the included sample text:
+
+```bash
+tiny-transformer train --data data/tiny_shakespeare_excerpt.txt --steps 300 --device cpu
+```
+
+Use the optional BPE tokenizer, gradient accumulation, and mixed precision when you want a stronger local run:
+
+```bash
+tiny-transformer train \
+  --data data/tiny_shakespeare_excerpt.txt \
+  --tokenizer bpe \
+  --bpe-vocab-size 128 \
+  --grad-accum-steps 4 \
+  --amp \
+  --device mps
+```
+
+Generate text from a checkpoint:
+
+```bash
+tiny-transformer generate --checkpoint runs/tiny-transformer.pt --prompt "To be" --max-new-tokens 160
+```
+
+Export an attention heatmap:
+
+```bash
+tiny-transformer attention --checkpoint runs/tiny-transformer.pt --prompt "To be" --output runs/attention.svg
+```
+
+Launch the local playground:
+
+```bash
+tiny-transformer serve --checkpoint runs/tiny-transformer.pt
+```
+
+Deploy the hosted playground:
+
+```bash
+pip install huggingface_hub
+huggingface-cli login
+huggingface-cli repo create tiny-transformer --type space --space_sdk gradio
+git remote add space https://huggingface.co/spaces/axay28/tiny-transformer
+git push space main
+```
+
+Run tests:
+
+```bash
+pytest
+```
+
+## Project Layout
+
+```text
+src/tiny_transformer/
+  cli.py          Command line interface for training and generation
+  config.py       Model and training configuration
+  data.py         Text dataset and batching utilities
+  model.py        GPT-style Transformer implementation
+  tokenizer.py    Character-level tokenizer
+  train.py        Training loop, evaluation, checkpointing
+  visualize.py    Attention heatmap export
+  web.py          Local generation playground
+tests/            Unit and smoke tests
+data/             Tiny sample corpus
+```
+
+## Architecture
+
+The model is intentionally small, but it follows the same structure as larger decoder-only LLMs:
+
+1. Token and positional embeddings convert IDs into vectors.
+2. Each Transformer block applies pre-norm causal self-attention.
+3. Feed-forward layers expand and compress the hidden dimension.
+4. Residual connections preserve gradient flow.
+5. A tied-size language modeling head predicts the next token.
+
+The attention mask is causal, so each position can only attend to itself and previous positions.
+
+```mermaid
+flowchart LR
+    A["Raw text corpus"] --> B["Char or BPE tokenizer"]
+    B --> C["Token IDs"]
+    C --> D["Contiguous train/val batches"]
+    D --> E["Token + position embeddings"]
+    E --> F1["LayerNorm"]
+    F1 --> F2["Masked multi-head self-attention"]
+    F2 --> F3["Residual add"]
+    F3 --> F4["LayerNorm"]
+    F4 --> F5["Feed-forward MLP"]
+    F5 --> F6["Residual add"]
+    F6 --> F7["Repeat for N layers"]
+    F7 --> G["Final layer norm"]
+    G --> H["Language modeling head"]
+    H --> I["Next-token logits"]
+    I --> J["Cross-entropy loss during training"]
+    I --> K["Top-k sampling during generation"]
+    F2 --> L["Attention heatmap export"]
+    K --> M["Local web playground"]
+```
+
+## Example Configuration
+
+The CLI defaults train quickly on CPU. For the included tiny corpus, the command uses a
+32-token context window; for larger text files, 128 tokens is a good next step:
+
+```python
+ModelConfig(
+    vocab_size=128,
+    block_size=128,
+    n_layer=4,
+    n_head=4,
+    n_embd=128,
+    dropout=0.1,
+)
+```
+
+Increase `n_layer`, `n_head`, and `n_embd` for a stronger demo once the training loop is validated.

src/tiny_transformer.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+README.md
+pyproject.toml
+src/tiny_transformer/__init__.py
+src/tiny_transformer/cli.py
+src/tiny_transformer/config.py
+src/tiny_transformer/data.py
+src/tiny_transformer/model.py
+src/tiny_transformer/tokenizer.py
+src/tiny_transformer/train.py
+src/tiny_transformer/visualize.py
+src/tiny_transformer/web.py
+src/tiny_transformer.egg-info/PKG-INFO
+src/tiny_transformer.egg-info/SOURCES.txt
+src/tiny_transformer.egg-info/dependency_links.txt
+src/tiny_transformer.egg-info/entry_points.txt
+src/tiny_transformer.egg-info/requires.txt
+src/tiny_transformer.egg-info/top_level.txt
+tests/test_model.py
+tests/test_tokenizer.py

src/tiny_transformer.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/tiny_transformer.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+tiny-transformer = tiny_transformer.cli:main

src/tiny_transformer.egg-info/requires.txt

@@ -0,0 +1,7 @@
+numpy>=2.0
+torch>=2.2
+tqdm>=4.66
+
+[dev]
+pytest>=8.0
+ruff>=0.5

src/tiny_transformer.egg-info/top_level.txt

@@ -0,0 +1 @@
+tiny_transformer

src/tiny_transformer/__init__.py

@@ -0,0 +1,14 @@
+"""Tiny Transformer: a minimal GPT-style language model."""
+
+from tiny_transformer.config import ModelConfig, TrainConfig
+from tiny_transformer.tokenizer import BytePairTokenizer, CharTokenizer
+
+__all__ = ["BytePairTokenizer", "CharTokenizer", "ModelConfig", "TinyTransformer", "TrainConfig"]
+
+
+def __getattr__(name: str):
+    if name == "TinyTransformer":
+        from tiny_transformer.model import TinyTransformer
+
+        return TinyTransformer
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/tiny_transformer/cli.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from tiny_transformer.config import ModelConfig, TrainConfig
+from tiny_transformer.train import load_checkpoint, train_from_text
+from tiny_transformer.visualize import save_attention_heatmap
+from tiny_transformer.web import serve_playground
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Train and sample a tiny GPT-style Transformer.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    train = subparsers.add_parser("train", help="Train a model on a plain-text corpus.")
+    train.add_argument("--data", required=True, help="Path to a UTF-8 text file.")
+    train.add_argument("--output", default="runs/tiny-transformer.pt", help="Checkpoint path.")
+    train.add_argument("--device", default="cpu", help="Device such as cpu, cuda, or mps.")
+    train.add_argument("--steps", type=int, default=1_000)
+    train.add_argument("--batch-size", type=int, default=32)
+    train.add_argument("--block-size", type=int, default=32)
+    train.add_argument("--layers", type=int, default=4)
+    train.add_argument("--heads", type=int, default=4)
+    train.add_argument("--embedding", type=int, default=128)
+    train.add_argument("--dropout", type=float, default=0.1)
+    train.add_argument("--learning-rate", type=float, default=3e-4)
+    train.add_argument("--tokenizer", choices=["char", "bpe"], default="char")
+    train.add_argument("--bpe-vocab-size", type=int, default=256)
+    train.add_argument("--grad-accum-steps", type=int, default=1)
+    train.add_argument("--amp", action="store_true", help="Use mixed precision on CUDA or MPS.")
+
+    generate = subparsers.add_parser("generate", help="Generate text from a trained checkpoint.")
+    generate.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
+    generate.add_argument("--prompt", default="\n", help="Prompt text.")
+    generate.add_argument("--device", default="cpu")
+    generate.add_argument("--max-new-tokens", type=int, default=200)
+    generate.add_argument("--temperature", type=float, default=0.8)
+    generate.add_argument("--top-k", type=int, default=20)
+
+    attention = subparsers.add_parser("attention", help="Export an attention heatmap SVG.")
+    attention.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
+    attention.add_argument("--prompt", required=True, help="Prompt text to inspect.")
+    attention.add_argument("--output", default="runs/attention.svg", help="SVG output path.")
+    attention.add_argument("--device", default="cpu")
+    attention.add_argument("--layer", type=int, default=-1, help="Layer index to visualize.")
+    attention.add_argument("--head", type=int, default=0, help="Attention head index to visualize.")
+
+    serve = subparsers.add_parser("serve", help="Launch a local text-generation playground.")
+    serve.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=int, default=8000)
+    serve.add_argument("--device", default="cpu")
+
+    return parser
+
+
+def train_command(args: argparse.Namespace) -> None:
+    text = Path(args.data).read_text(encoding="utf-8")
+    train_config = TrainConfig(
+        batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        max_steps=args.steps,
+        grad_accum_steps=args.grad_accum_steps,
+        use_amp=args.amp,
+        output_path=args.output,
+    )
+    model_config = ModelConfig(
+        vocab_size=1,
+        block_size=args.block_size,
+        n_layer=args.layers,
+        n_head=args.heads,
+        n_embd=args.embedding,
+        dropout=args.dropout,
+    )
+    train_from_text(
+        text,
+        model_config=model_config,
+        train_config=train_config,
+        device=args.device,
+        tokenizer_name=args.tokenizer,
+        bpe_vocab_size=args.bpe_vocab_size,
+    )
+    print(f"Saved checkpoint to {args.output}")
+
+
+def generate_command(args: argparse.Namespace) -> None:
+    model, tokenizer = load_checkpoint(args.checkpoint, device=args.device)
+    encoded = tokenizer.encode(args.prompt)
+    idx = torch.tensor([encoded], dtype=torch.long, device=args.device)
+    out = model.generate(
+        idx,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        top_k=args.top_k,
+    )
+    print(tokenizer.decode(out[0].tolist()))
+
+
+def attention_command(args: argparse.Namespace) -> None:
+    model, tokenizer = load_checkpoint(args.checkpoint, device=args.device)
+    encoded = tokenizer.encode(args.prompt)
+    idx = torch.tensor([encoded], dtype=torch.long, device=args.device)
+    save_attention_heatmap(
+        model=model,
+        tokenizer=tokenizer,
+        idx=idx,
+        output_path=args.output,
+        layer=args.layer,
+        head=args.head,
+    )
+    print(f"Saved attention heatmap to {args.output}")
+
+
+def serve_command(args: argparse.Namespace) -> None:
+    serve_playground(
+        checkpoint=args.checkpoint,
+        host=args.host,
+        port=args.port,
+        device=args.device,
+    )
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+    if args.command == "train":
+        train_command(args)
+    elif args.command == "generate":
+        generate_command(args)
+    elif args.command == "attention":
+        attention_command(args)
+    elif args.command == "serve":
+        serve_command(args)
+    else:
+        parser.error(f"Unknown command: {args.command}")
+
+
+if __name__ == "__main__":
+    main()

src/tiny_transformer/config.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    vocab_size: int
+    block_size: int = 128
+    n_layer: int = 4
+    n_head: int = 4
+    n_embd: int = 128
+    dropout: float = 0.1
+
+    def __post_init__(self) -> None:
+        if self.n_embd % self.n_head != 0:
+            raise ValueError("n_embd must be divisible by n_head")
+        if self.vocab_size <= 0:
+            raise ValueError("vocab_size must be positive")
+
+    def to_dict(self) -> dict[str, int | float]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class TrainConfig:
+    batch_size: int = 32
+    learning_rate: float = 3e-4
+    max_steps: int = 1_000
+    eval_interval: int = 100
+    eval_batches: int = 20
+    grad_accum_steps: int = 1
+    use_amp: bool = False
+    seed: int = 1337
+    output_path: str = "runs/tiny-transformer.pt"
+
+    def __post_init__(self) -> None:
+        if self.grad_accum_steps <= 0:
+            raise ValueError("grad_accum_steps must be positive")
+
+    def to_dict(self) -> dict[str, bool | int | float | str]:
+        return asdict(self)

src/tiny_transformer/data.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import torch
+
+
+class TextDataset:
+    def __init__(self, token_ids: list[int], block_size: int, device: str) -> None:
+        if len(token_ids) <= block_size:
+            raise ValueError("Dataset must contain more tokens than block_size")
+        self.data = torch.tensor(token_ids, dtype=torch.long, device=device)
+        self.block_size = block_size
+        self.device = device
+
+    def get_batch(self, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]:
+        max_start = len(self.data) - self.block_size
+        starts = torch.randint(0, max_start, (batch_size,), device=self.device)
+        x = torch.stack([self.data[start : start + self.block_size] for start in starts])
+        y = torch.stack([self.data[start + 1 : start + self.block_size + 1] for start in starts])
+        return x, y
+
+
+def split_tokens(token_ids: list[int], train_fraction: float = 0.9) -> tuple[list[int], list[int]]:
+    if not 0.0 < train_fraction < 1.0:
+        raise ValueError("train_fraction must be between 0 and 1")
+    split_idx = int(len(token_ids) * train_fraction)
+    return token_ids[:split_idx], token_ids[split_idx:]

src/tiny_transformer/model.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from tiny_transformer.config import ModelConfig
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.n_head = config.n_head
+        self.head_dim = config.n_embd // config.n_head
+        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
+        self.proj = nn.Linear(config.n_embd, config.n_embd)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        mask = torch.tril(torch.ones(config.block_size, config.block_size))
+        self.register_buffer("causal_mask", mask.view(1, 1, config.block_size, config.block_size))
+
+    def forward(
+        self, x: torch.Tensor, return_attention: bool = False
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        batch, seq_len, channels = x.shape
+        qkv = self.qkv(x)
+        query, key, value = qkv.split(channels, dim=2)
+
+        query = query.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        key = key.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        value = value.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+
+        scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim)
+        scores = scores.masked_fill(self.causal_mask[:, :, :seq_len, :seq_len] == 0, float("-inf"))
+        weights = F.softmax(scores, dim=-1)
+        weights = self.attn_dropout(weights)
+        out = weights @ value
+        out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels)
+        out = self.resid_dropout(self.proj(out))
+        if return_attention:
+            return out, weights
+        return out
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(config.n_embd, 4 * config.n_embd),
+            nn.GELU(),
+            nn.Linear(4 * config.n_embd, config.n_embd),
+            nn.Dropout(config.dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.ffwd = FeedForward(config)
+
+    def forward(
+        self, x: torch.Tensor, return_attention: bool = False
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if return_attention:
+            attn_out, weights = self.attn(self.ln_1(x), return_attention=True)
+            x = x + attn_out
+            x = x + self.ffwd(self.ln_2(x))
+            return x, weights
+
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.ffwd(self.ln_2(x))
+        return x
+
+
+class TinyTransformer(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
+        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)])
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self, idx: torch.Tensor, targets: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        logits, loss, _ = self._forward(idx, targets, capture_attention=False)
+        return logits, loss
+
+    def _forward(
+        self,
+        idx: torch.Tensor,
+        targets: torch.Tensor | None = None,
+        capture_attention: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        batch, seq_len = idx.shape
+        if seq_len > self.config.block_size:
+            raise ValueError("Sequence length exceeds block_size")
+
+        attentions: list[torch.Tensor] = []
+        positions = torch.arange(seq_len, device=idx.device)
+        x = self.token_embedding(idx) + self.position_embedding(positions)
+        x = self.dropout(x)
+        for block in self.blocks:
+            if capture_attention:
+                x, weights = block(x, return_attention=True)
+                attentions.append(weights)
+            else:
+                x = block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(batch * seq_len, -1), targets.view(batch * seq_len))
+        return logits, loss, attentions
+
+    @torch.no_grad()
+    def attention_maps(self, idx: torch.Tensor) -> list[torch.Tensor]:
+        self.eval()
+        _, _, attentions = self._forward(idx, capture_attention=True)
+        return attentions
+
+    @torch.no_grad()
+    def generate(
+        self,
+        idx: torch.Tensor,
+        max_new_tokens: int,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+    ) -> torch.Tensor:
+        if temperature <= 0:
+            raise ValueError("temperature must be positive")
+
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.config.block_size :]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < values[:, [-1]]] = -float("inf")
+            probs = F.softmax(logits, dim=-1)
+            next_idx = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, next_idx), dim=1)
+        return idx

src/tiny_transformer/tokenizer.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class Tokenizer(Protocol):
+    @property
+    def vocab_size(self) -> int: ...
+
+    def encode(self, text: str) -> list[int]: ...
+
+    def decode(self, ids: list[int]) -> str: ...
+
+    def id_to_token(self, idx: int) -> str: ...
+
+    def to_dict(self) -> dict[str, object]: ...
+
+
+@dataclass(frozen=True)
+class CharTokenizer:
+    stoi: dict[str, int]
+    itos: dict[int, str]
+
+    @classmethod
+    def train(cls, text: str) -> "CharTokenizer":
+        if not text:
+            raise ValueError("Cannot train a tokenizer on empty text")
+        chars = sorted(set(text))
+        stoi = {char: idx for idx, char in enumerate(chars)}
+        itos = {idx: char for char, idx in stoi.items()}
+        return cls(stoi=stoi, itos=itos)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.stoi)
+
+    def encode(self, text: str) -> list[int]:
+        try:
+            return [self.stoi[char] for char in text]
+        except KeyError as exc:
+            char = exc.args[0]
+            raise ValueError(f"Character {char!r} is not in the tokenizer vocabulary") from exc
+
+    def decode(self, ids: list[int]) -> str:
+        try:
+            return "".join(self.itos[idx] for idx in ids)
+        except KeyError as exc:
+            idx = exc.args[0]
+            raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
+
+    def id_to_token(self, idx: int) -> str:
+        try:
+            return self.itos[idx]
+        except KeyError as exc:
+            raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
+
+    def to_dict(self) -> dict[str, object]:
+        return {"type": "char", "stoi": self.stoi, "itos": self.itos}
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, dict[str, int] | dict[int | str, str]]) -> "CharTokenizer":
+        stoi = {str(char): int(idx) for char, idx in payload["stoi"].items()}
+        itos = {int(idx): str(char) for idx, char in payload["itos"].items()}
+        return cls(stoi=stoi, itos=itos)
+
+
+@dataclass(frozen=True)
+class BytePairTokenizer:
+    stoi: dict[str, int]
+    itos: dict[int, str]
+    merges: list[tuple[str, str]]
+
+    @classmethod
+    def train(cls, text: str, vocab_size: int = 256) -> "BytePairTokenizer":
+        if not text:
+            raise ValueError("Cannot train a tokenizer on empty text")
+        if vocab_size <= 0:
+            raise ValueError("vocab_size must be positive")
+
+        base_vocab = sorted(set(text))
+        sequences = [[char for char in text]]
+        merges: list[tuple[str, str]] = []
+        vocab = set(base_vocab)
+
+        while len(vocab) < vocab_size:
+            pair_counts = _count_pairs(sequences)
+            if not pair_counts:
+                break
+            pair, count = pair_counts.most_common(1)[0]
+            merged = "".join(pair)
+            if count < 2 or merged in vocab:
+                break
+            sequences = [_merge_pair(sequence, pair, merged) for sequence in sequences]
+            merges.append(pair)
+            vocab.add(merged)
+
+        tokens = sorted(vocab, key=lambda token: (len(token), token))
+        stoi = {token: idx for idx, token in enumerate(tokens)}
+        itos = {idx: token for token, idx in stoi.items()}
+        return cls(stoi=stoi, itos=itos, merges=merges)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.stoi)
+
+    def encode(self, text: str) -> list[int]:
+        unknown = sorted(set(text) - {token for token in self.stoi if len(token) == 1})
+        if unknown:
+            raise ValueError(f"Characters {unknown!r} are not in the tokenizer vocabulary")
+
+        pieces = list(text)
+        for left, right in self.merges:
+            pieces = _merge_pair(pieces, (left, right), left + right)
+        return [self.stoi[piece] for piece in pieces]
+
+    def decode(self, ids: list[int]) -> str:
+        try:
+            return "".join(self.itos[idx] for idx in ids)
+        except KeyError as exc:
+            idx = exc.args[0]
+            raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
+
+    def id_to_token(self, idx: int) -> str:
+        try:
+            return self.itos[idx]
+        except KeyError as exc:
+            raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "type": "bpe",
+            "stoi": self.stoi,
+            "itos": self.itos,
+            "merges": self.merges,
+        }
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, object]) -> "BytePairTokenizer":
+        raw_stoi = payload["stoi"]
+        raw_itos = payload["itos"]
+        raw_merges = payload["merges"]
+        if not isinstance(raw_stoi, dict) or not isinstance(raw_itos, dict):
+            raise ValueError("Invalid BPE tokenizer payload")
+        if not isinstance(raw_merges, list):
+            raise ValueError("Invalid BPE merge payload")
+        stoi = {str(token): int(idx) for token, idx in raw_stoi.items()}
+        itos = {int(idx): str(token) for idx, token in raw_itos.items()}
+        merges = [(str(left), str(right)) for left, right in raw_merges]
+        return cls(stoi=stoi, itos=itos, merges=merges)
+
+
+def tokenizer_from_dict(payload: dict[str, object]) -> Tokenizer:
+    tokenizer_type = payload.get("type", "char")
+    if tokenizer_type == "char":
+        return CharTokenizer.from_dict(payload)  # type: ignore[arg-type]
+    if tokenizer_type == "bpe":
+        return BytePairTokenizer.from_dict(payload)
+    raise ValueError(f"Unknown tokenizer type: {tokenizer_type!r}")
+
+
+def _count_pairs(sequences: list[list[str]]) -> Counter[tuple[str, str]]:
+    counts: Counter[tuple[str, str]] = Counter()
+    for sequence in sequences:
+        counts.update(zip(sequence, sequence[1:]))
+    return counts
+
+
+def _merge_pair(sequence: list[str], pair: tuple[str, str], merged: str) -> list[str]:
+    out: list[str] = []
+    idx = 0
+    while idx < len(sequence):
+        if idx < len(sequence) - 1 and (sequence[idx], sequence[idx + 1]) == pair:
+            out.append(merged)
+            idx += 2
+        else:
+            out.append(sequence[idx])
+            idx += 1
+    return out

src/tiny_transformer/train.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from tqdm import trange
+
+from tiny_transformer.config import ModelConfig, TrainConfig
+from tiny_transformer.data import TextDataset, split_tokens
+from tiny_transformer.model import TinyTransformer
+from tiny_transformer.tokenizer import BytePairTokenizer, CharTokenizer, Tokenizer, tokenizer_from_dict
+
+
+@torch.no_grad()
+def estimate_loss(
+    model: TinyTransformer,
+    train_data: TextDataset,
+    val_data: TextDataset,
+    batch_size: int,
+    eval_batches: int,
+) -> dict[str, float]:
+    model.eval()
+    losses: dict[str, float] = {}
+    for split, dataset in {"train": train_data, "val": val_data}.items():
+        split_losses = []
+        for _ in range(eval_batches):
+            x, y = dataset.get_batch(batch_size)
+            _, loss = model(x, y)
+            if loss is None:
+                raise RuntimeError("Expected a loss during evaluation")
+            split_losses.append(loss.item())
+        losses[split] = sum(split_losses) / len(split_losses)
+    model.train()
+    return losses
+
+
+def train_from_text(
+    text: str,
+    model_config: ModelConfig | None = None,
+    train_config: TrainConfig | None = None,
+    device: str = "cpu",
+    tokenizer_name: str = "char",
+    bpe_vocab_size: int = 256,
+) -> TinyTransformer:
+    train_config = train_config or TrainConfig()
+    torch.manual_seed(train_config.seed)
+
+    tokenizer = train_tokenizer(text, tokenizer_name, bpe_vocab_size)
+    token_ids = tokenizer.encode(text)
+    train_ids, val_ids = split_tokens(token_ids)
+
+    if model_config is None:
+        model_config = ModelConfig(vocab_size=tokenizer.vocab_size)
+    else:
+        model_config = ModelConfig(**{**model_config.to_dict(), "vocab_size": tokenizer.vocab_size})
+
+    train_data = TextDataset(train_ids, model_config.block_size, device)
+    val_data = TextDataset(val_ids, model_config.block_size, device)
+    model = TinyTransformer(model_config).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=train_config.learning_rate)
+    device_type = "cuda" if device.startswith("cuda") else "mps" if device == "mps" else "cpu"
+    amp_enabled = train_config.use_amp and device_type in {"cuda", "mps"}
+    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled and device_type == "cuda")
+
+    progress = trange(train_config.max_steps, desc="training", leave=True)
+    for step in progress:
+        if step % train_config.eval_interval == 0 or step == train_config.max_steps - 1:
+            losses = estimate_loss(
+                model, train_data, val_data, train_config.batch_size, train_config.eval_batches
+            )
+            progress.set_postfix(train=f"{losses['train']:.3f}", val=f"{losses['val']:.3f}")
+
+        optimizer.zero_grad(set_to_none=True)
+        for _ in range(train_config.grad_accum_steps):
+            x, y = train_data.get_batch(train_config.batch_size)
+            with torch.autocast(device_type=device_type, enabled=amp_enabled):
+                _, loss = model(x, y)
+                if loss is None:
+                    raise RuntimeError("Expected a loss during training")
+                loss = loss / train_config.grad_accum_steps
+            scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+    save_checkpoint(model, tokenizer, train_config.output_path)
+    return model
+
+
+def train_tokenizer(text: str, tokenizer_name: str, bpe_vocab_size: int) -> Tokenizer:
+    if tokenizer_name == "char":
+        return CharTokenizer.train(text)
+    if tokenizer_name == "bpe":
+        return BytePairTokenizer.train(text, vocab_size=bpe_vocab_size)
+    raise ValueError("tokenizer_name must be 'char' or 'bpe'")
+
+
+def save_checkpoint(model: TinyTransformer, tokenizer: Tokenizer, path: str) -> None:
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model_config": model.config.to_dict(),
+            "model_state": model.state_dict(),
+            "tokenizer": tokenizer.to_dict(),
+        },
+        output,
+    )
+
+
+def load_checkpoint(path: str, device: str = "cpu") -> tuple[TinyTransformer, Tokenizer]:
+    payload = torch.load(path, map_location=device)
+    tokenizer = tokenizer_from_dict(payload["tokenizer"])
+    config = ModelConfig(**payload["model_config"])
+    model = TinyTransformer(config).to(device)
+    model.load_state_dict(payload["model_state"])
+    model.eval()
+    return model, tokenizer

src/tiny_transformer/visualize.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from html import escape
+from pathlib import Path
+
+import torch
+
+from tiny_transformer.model import TinyTransformer
+from tiny_transformer.tokenizer import Tokenizer
+
+
+@torch.no_grad()
+def save_attention_heatmap(
+    model: TinyTransformer,
+    tokenizer: Tokenizer,
+    idx: torch.Tensor,
+    output_path: str,
+    layer: int = -1,
+    head: int = 0,
+) -> None:
+    attentions = model.attention_maps(idx)
+    if not attentions:
+        raise ValueError("Model did not return attention maps")
+
+    selected = attentions[layer][0]
+    if head < 0 or head >= selected.shape[0]:
+        raise ValueError(f"head must be between 0 and {selected.shape[0] - 1}")
+
+    weights = selected[head].detach().cpu()
+    token_ids = idx[0].detach().cpu().tolist()
+    labels = [_display_token(tokenizer.id_to_token(token_id)) for token_id in token_ids]
+    svg = _attention_svg(weights, labels, layer=layer, head=head)
+
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(svg, encoding="utf-8")
+
+
+def _attention_svg(weights: torch.Tensor, labels: list[str], layer: int, head: int) -> str:
+    cell = 24
+    margin_left = 120
+    margin_top = 96
+    size = len(labels)
+    width = margin_left + size * cell + 24
+    height = margin_top + size * cell + 40
+    max_weight = max(float(weights.max()), 1e-9)
+
+    cells = []
+    for row in range(size):
+        for col in range(size):
+            value = float(weights[row, col]) / max_weight
+            color = 255 - int(value * 210)
+            cells.append(
+                f'<rect x="{margin_left + col * cell}" y="{margin_top + row * cell}" '
+                f'width="{cell}" height="{cell}" fill="rgb({color},{color},255)">'
+                f"<title>{escape(labels[row])} attends to {escape(labels[col])}: "
+                f"{float(weights[row, col]):.3f}</title></rect>"
+            )
+
+    x_labels = [
+        f'<text x="{margin_left + idx * cell + 12}" y="{margin_top - 10}" '
+        f'transform="rotate(-45 {margin_left + idx * cell + 12},{margin_top - 10})">'
+        f"{escape(label)}</text>"
+        for idx, label in enumerate(labels)
+    ]
+    y_labels = [
+        f'<text x="{margin_left - 8}" y="{margin_top + idx * cell + 16}" text-anchor="end">'
+        f"{escape(label)}</text>"
+        for idx, label in enumerate(labels)
+    ]
+
+    return "\n".join(
+        [
+            f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
+            f'viewBox="0 0 {width} {height}">',
+            "<style>text{font:12px system-ui,sans-serif} rect{stroke:#fff;stroke-width:1}</style>",
+            f'<text x="16" y="28" style="font-size:18px;font-weight:700">'
+            f"Attention heatmap: layer {layer}, head {head}</text>",
+            *x_labels,
+            *y_labels,
+            *cells,
+            "</svg>",
+        ]
+    )
+
+
+def _display_token(token: str) -> str:
+    return token.replace("\n", "\\n").replace("\t", "\\t").replace(" ", "space")

src/tiny_transformer/web.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import json
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import torch
+
+from tiny_transformer.train import load_checkpoint
+
+
+HTML = """<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>Tiny Transformer Playground</title>
+  <style>
+    body { margin: 0; font: 16px system-ui, sans-serif; background: #f7f7f4; color: #1b1b1b; }
+    main { max-width: 920px; margin: 0 auto; padding: 32px 20px; }
+    h1 { margin: 0 0 18px; font-size: 32px; }
+    textarea, pre { box-sizing: border-box; width: 100%; border: 1px solid #c8c8c2;
+      border-radius: 8px; padding: 14px; background: white; color: #1b1b1b; }
+    textarea { min-height: 110px; resize: vertical; }
+    pre { min-height: 220px; white-space: pre-wrap; line-height: 1.45; }
+    .controls { display: flex; gap: 12px; flex-wrap: wrap; align-items: center; margin: 14px 0; }
+    label { display: grid; gap: 4px; font-size: 13px; }
+    input { width: 86px; padding: 8px; border: 1px solid #c8c8c2; border-radius: 6px; }
+    button { border: 0; border-radius: 8px; padding: 10px 16px; color: white;
+      background: #1f5f5b; font-weight: 700; cursor: pointer; }
+  </style>
+</head>
+<body>
+  <main>
+    <h1>Tiny Transformer Playground</h1>
+    <textarea id="prompt">To be</textarea>
+    <div class="controls">
+      <label>New tokens <input id="tokens" type="number" min="1" max="500" value="120"></label>
+      <label>Temperature <input id="temperature" type="number" min="0.1" step="0.1" value="0.8"></label>
+      <label>Top-k <input id="topk" type="number" min="1" value="20"></label>
+      <button id="generate">Generate</button>
+    </div>
+    <pre id="output"></pre>
+  </main>
+  <script>
+    const button = document.getElementById("generate");
+    button.addEventListener("click", async () => {
+      button.disabled = true;
+      document.getElementById("output").textContent = "Generating...";
+      const response = await fetch("/api/generate", {
+        method: "POST",
+        headers: {"Content-Type": "application/json"},
+        body: JSON.stringify({
+          prompt: document.getElementById("prompt").value,
+          max_new_tokens: Number(document.getElementById("tokens").value),
+          temperature: Number(document.getElementById("temperature").value),
+          top_k: Number(document.getElementById("topk").value)
+        })
+      });
+      const payload = await response.json();
+      document.getElementById("output").textContent = payload.text || payload.error;
+      button.disabled = false;
+    });
+  </script>
+</body>
+</html>
+"""
+
+
+def serve_playground(checkpoint: str, host: str, port: int, device: str) -> None:
+    model, tokenizer = load_checkpoint(checkpoint, device=device)
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self) -> None:
+            if self.path != "/":
+                self.send_error(404)
+                return
+            self._send(200, HTML.encode("utf-8"), "text/html; charset=utf-8")
+
+        def do_POST(self) -> None:
+            if self.path != "/api/generate":
+                self.send_error(404)
+                return
+            length = int(self.headers.get("content-length", "0"))
+            payload = json.loads(self.rfile.read(length) or b"{}")
+            try:
+                prompt = str(payload.get("prompt", ""))
+                idx = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device)
+                out = model.generate(
+                    idx,
+                    max_new_tokens=int(payload.get("max_new_tokens", 120)),
+                    temperature=float(payload.get("temperature", 0.8)),
+                    top_k=int(payload.get("top_k", 20)),
+                )
+                body = json.dumps({"text": tokenizer.decode(out[0].tolist())}).encode("utf-8")
+                self._send(200, body, "application/json")
+            except Exception as exc:
+                body = json.dumps({"error": str(exc)}).encode("utf-8")
+                self._send(400, body, "application/json")
+
+        def log_message(self, format: str, *args: object) -> None:
+            return
+
+        def _send(self, status: int, body: bytes, content_type: str) -> None:
+            self.send_response(status)
+            self.send_header("Content-Type", content_type)
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+    server = ThreadingHTTPServer((host, port), Handler)
+    print(f"Serving playground at http://{host}:{port}")
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\nStopped playground server.")
+    finally:
+        server.server_close()