Consolidated, documented mini-GPT codebase · CS336 Language Modeling From Scratch

1. Shape the module structure

You’re turning a tangle of scripts into a small library: separate what the code does (tokenize, run a model, train, sample) from how you run it (configs, entry scripts).

A clean structure might look like this:

text

mini-gpt/
  README.md

  mini_gpt/
    __init__.py
    config.py
    tokenizer.py
    model.py
    data.py
    train.py
    sample.py
    utils.py

  configs/
    tiny_char.yaml
    tiny_bpe.yaml

  scripts/
    train_tiny.sh
    sample_tiny.sh

  checkpoints/
    (created at run time)

1.1 Core modules and their roles

Make each file own one kind of responsibility:

Module	Responsibility focus	Typical contents
`tokenizer.py`	Turn raw text into token ids and back	`Tokenizer` class, `encode`, `decode`
`model.py`	Define the GPT architecture	`MiniGPT` model, positional embeddings, attention
`data.py`	Load and batch your dataset	Dataset class, `DataLoader` helpers
`train.py`	Training loop and optimizer wiring	`train` function, logging, checkpointing
`sample.py`	Inference and text generation	`generate` function, top-k/top-p, temperature
`config.py`	Config schema and loading	`Config` dataclass, `load_config`
`utils.py`	Shared helpers that don’t belong elsewhere	seeding, device selection, timing

Mechanism: separating modules reduces hidden coupling:

train.py only needs to know how to call the model and tokenizer, not their internals.
You can swap in a different tokenizer or model by changing config/imports, not the training loop.

1.2 Minimal concrete example

Here’s a skeleton that can actually run end-to-end on dummy data.

mini_gpt/tokenizer.py:

python

# mini_gpt/tokenizer.py
class CharTokenizer:
    def __init__(self, text: str):
        chars = sorted(list(set(text)))
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for ch, i in self.stoi.items()}

    def encode(self, s: str):
        return [self.stoi[ch] for ch in s]

    def decode(self, ids):
        return "".join(self.itos[i] for i in ids)

mini_gpt/model.py (super minimal, not a full GPT, but runnable):

python

# mini_gpt/model.py
import torch
import torch.nn as nn

class MiniGPT(nn.Module):
    def __init__(self, vocab_size: int, d_model: int = 64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, idx):
        # idx: [batch, time]
        x = self.embed(idx)          # [batch, time, d_model]
        logits = self.lm_head(x)     # [batch, time, vocab]
        return logits

mini_gpt/train.py:

python

# mini_gpt/train.py
import torch
from torch import nn, optim

from .model import MiniGPT
from .tokenizer import CharTokenizer

def train_simple(text: str, device: str = "cpu"):
    tokenizer = CharTokenizer(text)
    data = torch.tensor(tokenizer.encode(text), dtype=torch.long, device=device)

    model = MiniGPT(vocab_size=len(tokenizer.stoi)).to(device)
    criterion = nn.CrossEntropyLoss()
    opt = optim.AdamW(model.parameters(), lr=1e-3)

    seq_len = 16
    n_steps = 10  # tiny demo

    for step in range(n_steps):
        # simple next-token prediction on one chunk
        if len(data) <= seq_len:
            continue
        x = data[:seq_len].unsqueeze(0)       # [1, seq_len]
        y = data[1:seq_len+1].unsqueeze(0)    # [1, seq_len]

        logits = model(x)                     # [1, seq_len, vocab]
        loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))

        opt.zero_grad()
        loss.backward()
        opt.step()

        print(f"step {step} loss {loss.item():.4f}")

    return model, tokenizer

mini_gpt/sample.py:

python

# mini_gpt/sample.py
import torch

@torch.no_grad()
def generate(model, tokenizer, prompt: str, max_new_tokens: int = 50, device: str = "cpu"):
    model.eval()
    idx = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device)

    for _ in range(max_new_tokens):
        logits = model(idx)               # [1, T, vocab]
        last_logits = logits[:, -1, :]    # [1, vocab]
        probs = torch.softmax(last_logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)  # [1, 1]
        idx = torch.cat([idx, next_id], dim=1)

    return tokenizer.decode(idx[0].tolist())

example_run.py at repo root:

python

# example_run.py
from mini_gpt.train import train_simple
from mini_gpt.sample import generate

if __name__ == "__main__":
    text = "hello world\n" * 100
    model, tokenizer = train_simple(text)
    print(generate(model, tokenizer, prompt="hel", max_new_tokens=10))

Running:

bash

python example_run.py

should print a few training losses and then some generated characters.

Don’t let train.py turn into a dumping ground. If you feel tempted to add tokenization logic, model building, or config loading there, stop and move that code into tokenizer.py, model.py, or config.py instead.

1 / 4