Build a Mini GPT From Scratch

Build, train, and deploy your own GPT-style language model in Python.

Introduction

30 Second Summary

Every time you ask ChatGPT a question or use an AI writing tool, a language model is generating that text one piece at a time. But how does it actually work inside?

In this project, you will build your own mini language model from scratch, train it on real Shakespeare text using your GPU, and deploy it as a local API that generates text on demand.

What You'll Build

A complete mini GPT-style transformer that learns to write Shakespeare-like text, served through a REST API you can query from anywhere.

By the end of this project, you'll have:

  • A hand-built transformer architecture in Python covering embeddings, self-attention, and feed-forward layers.
  • A trained model that generates Shakespeare-style text after learning from real Shakespeare data on your GPU.
  • A local REST API that accepts text prompts and returns AI-generated text from your model.
  • Secret Mission: Add temperature and top-k sampling controls to your API for creative text generation.

Are there any prerequisites?

You only need basic Python knowledge (variables, functions, loops). No machine learning experience is required.

Before We Start

Building a language model from scratch means writing the actual architecture that powers tools like ChatGPT. Before you can write any model code, you need a development environment with GPU support.

In this step, you will create a project folder, install PyTorch with CUDA so your GPU can accelerate training, and set up FastAPI for serving your model later.

In this step, get ready to:
  • Create a project folder and Python virtual environment.
  • Install PyTorch with CUDA 12.6 support and verify GPU detection.
  • Install FastAPI and create your project file structure.
Create a project folder and virtual environment

A virtual environment keeps your project's packages isolated from the rest of your system. This prevents version conflicts between projects.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal (macOS) or PowerShell (Windows) and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir mini-gpt
cd mini-gpt
  • Create and activate a Python virtual environment by running the commands for your operating system:

macOS

python3 -m venv venv
source venv/bin/activate

Windows

python -m venv venv
venv\Scripts\activate

You should see (venv) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.

What is a virtual environment?

A virtual environment is an isolated Python installation inside your project folder. Any packages you install while it is active only exist inside that environment.

This keeps your system Python clean and prevents one project's dependencies from breaking another.

Don't see (venv) in your prompt?

Make sure you ran the activation command (not just the creation command). On Windows, if you see an error about execution policies, run Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned and then try the activation command again.

Help me fix my virtual environment activation.

Install PyTorch with CUDA 12.6 support and verify GPU detection

PyTorch is the deep learning framework you will use to build and train your transformer model. Installing it with CUDA support lets your NVIDIA GPU accelerate the math-heavy operations during training.

  • Install PyTorch with CUDA 12.6 by running this command:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

What does this command do?

This installs PyTorch 2.7.0 along with its companion libraries (torchvision and torchaudio) from a special index URL that includes CUDA 12.6 GPU support.

PyTorch binaries ship with their own CUDA runtime, so you do not need to install the CUDA toolkit separately.

This download is around 2-3 GB. It may take a few minutes depending on your internet speed.

  • Once the install finishes, verify your GPU is detected by running this command:
python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

✔️ I see True and my GPU name

Your GPU is detected and ready for training. You should see True on the first line and something like NVIDIA GeForce RTX 3050 Ti Laptop GPU on the second. You are all set.

ⓧ I see False or an error

If torch.cuda.is_available() returns False, the most common causes are:

  • You accidentally installed the CPU-only version of PyTorch. Uninstall and reinstall with the exact command above (the --index-url flag is critical).
  • Your NVIDIA drivers are outdated. Update them through the NVIDIA GeForce Experience app or the NVIDIA website.
  • Your virtual environment was not active when you ran the install. Check for (venv) in your prompt and try again.

Help me get CUDA working with PyTorch.

Install FastAPI and create your project file structure

FastAPI is the web framework you will use to serve your trained model as a REST API in the final step.

  • Install FastAPI and its dependencies by running this command:
pip install "fastapi[standard]"

What does fastapi[standard] include?

The [standard] extra installs FastAPI along with uvicorn (the server that runs your app), plus useful defaults for development like auto-reload support.

  • Create the three Python files your project needs by running the command for your operating system:

macOS

touch model.py train.py app.py

Windows

New-Item -ItemType File -Path "model.py", "train.py", "app.py"
  • Confirm the files exist by running:
ls

You should see model.py, train.py, and app.py listed in the output.

What are these files for?

Each file has a specific role in your project. model.py will hold your transformer architecture. train.py will handle downloading data and training the model. app.py will serve your trained model as a web API.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.
  • Click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop and select the mini-gpt folder, then click Open.

You should see model.py, train.py, and app.py listed in the Explorer sidebar on the left.

Don't see the files in the sidebar?

Make sure you opened the mini-gpt folder itself, not its parent (Desktop). The folder name should appear at the top of the Explorer sidebar. If you see Desktop instead, use File > Open Folder again and select the mini-gpt folder.

Help me find my project files in VS Code.

Your development environment is locked and loaded. Next up, you will download Shakespeare text and build a character-level tokenizer to prepare your training data.

Prepare the Training Data

Your environment is set up and your GPU is ready. Now you need something for your model to learn from.

Language models learn by studying massive amounts of text. For this project, you will use the Tiny Shakespeare dataset, a collection of about 1.1 million characters of Shakespeare's plays. Your model will learn the patterns in this text and eventually generate new Shakespeare-style writing.

In this step, get ready to:
  • Download the Shakespeare text dataset.
  • Build a character-level tokenizer that converts text to numbers and back.
  • Create training and validation datasets with a batching function.
Download Shakespeare text and load it into Python

Before you write any code, download the training data so it is ready on disk. The Tiny Shakespeare dataset is a single text file containing about 1.1 million characters from various Shakespeare plays.

  • Download the Shakespeare dataset by running this command in your terminal:
python -c "import urllib.request; urllib.request.urlretrieve('https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt', 'input.txt'); print('Downloaded Shakespeare dataset.')"
  • Confirm the file was downloaded by running:
ls input.txt
  • You should see input.txt listed in the output.

Don't see input.txt?

Make sure your terminal is in the mini-gpt folder. Run pwd to check your current directory.

If you see a network error, check your internet connection and try the command again.

Help me troubleshoot the Shakespeare dataset download

Now open train.py in VS Code. This is where all your data preparation code will live.

  • Add the imports and training hyperparameters to train.py by copying the following code:
import torch
from model import MiniGPT, BLOCK_SIZE, device
import urllib.request
import os

# Hyperparameters
BATCH_SIZE = 64
MAX_ITERS = 3000
EVAL_INTERVAL = 300
LEARNING_RATE = 3e-4
EVAL_ITERS = 200

What does this code do?

  • torch is the PyTorch library for building and training neural networks.
  • from model import MiniGPT, BLOCK_SIZE, device imports the model architecture you will build in the next step. This line will show a warning in VS Code until you complete Step 3.
  • The hyperparameters control training behavior. BATCH_SIZE sets how many text samples to process at once. MAX_ITERS sets how many training steps to run. LEARNING_RATE controls how fast the model adjusts its weights.

Why is there a red underline on the model import?

Your model.py file is still empty. VS Code is correct that the import cannot resolve yet. You will fill in model.py in Step 3, and the warning will disappear.

  • Below the hyperparameters, add the download and file-reading logic:
# Download Shakespeare dataset
DATA_PATH = "input.txt"
if not os.path.exists(DATA_PATH):
    url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
    urllib.request.urlretrieve(url, DATA_PATH)
    print("Downloaded Shakespeare dataset.")

# Read and process data
with open(DATA_PATH, "r", encoding="utf-8") as f:
    text = f.read()

What does this code do?

  • DATA_PATH points to the Shakespeare text file you already downloaded.
  • The if not os.path.exists check means the script will only download the file if it is missing. Since you already have it, this block will be skipped at runtime.
  • The open() call reads the entire 1.1 million character dataset into the text variable as a single Python string.
  • Save train.py.

You won't run this file yet because the model import needs model.py to be complete. That comes in Step 3.

Build a character-level tokenizer

Neural networks work with numbers, not text. A tokenizer converts each character into a unique integer so the model can process it. Character-level tokenization is the simplest approach: each letter, space, and punctuation mark gets its own number.

  • Add the tokenizer code below the file-reading section in train.py:
# Build character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join([itos[i] for i in l])

What does this code do?

  • chars collects every unique character in the text and sorts them alphabetically. The Shakespeare dataset has 65 unique characters (letters, digits, punctuation, and spaces).
  • stoi (string-to-integer) maps each character to its index number. For example, 'a' might map to 39.
  • itos (integer-to-string) is the reverse lookup. Given a number, it returns the character.
  • encode converts a string into a list of integers. decode converts a list of integers back into a readable string.

Why character-level instead of word-level?

Word-level tokenizers need large vocabularies (50,000+ words) and complex splitting logic. Character-level tokenization keeps things simple with only 65 possible tokens. This lets you focus on the transformer architecture itself, which is the real learning goal.

Production models like ChatGPT use a middle ground called BPE (byte pair encoding), but the underlying transformer architecture is identical.

  • Save train.py.

You won't see a change yet. The next substep uses this tokenizer to create the actual training data.

Create training and validation datasets with batching

With the tokenizer ready, you need to split the data into a training set and a validation set. The model learns from the training set and gets tested on the validation set to make sure it is actually learning patterns, not just memorizing text.

  • Add the data splitting and batching code below the tokenizer in train.py:
# Create train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - BLOCK_SIZE, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + BLOCK_SIZE] for i in ix])
    y = torch.stack([data_split[i + 1 : i + BLOCK_SIZE + 1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

What does this code do?

  • The encode(text) call converts the entire Shakespeare text into a list of integers, then torch.tensor wraps it as a PyTorch tensor.
  • The first 90% of the data goes to train_data and the remaining 10% to val_data. This gives the model about 1 million characters to learn from and 111,000 characters to test against.
  • get_batch() randomly selects 64 chunks of text (BATCH_SIZE), each 128 characters long (BLOCK_SIZE). The x tensor is the input and y is the target (shifted by one character). The model learns to predict the next character given the previous ones.
  • estimate_loss() runs 200 batches through the model without updating weights (@torch.no_grad()) and averages the loss. This gives a stable measurement of how well the model is performing.
  • Save train.py.

You won't run this file yet. In the next step, you will build model.py with the full transformer architecture, and then in Step 4 you will run train.py and watch your model learn.

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py and that input.txt exists in your mini-gpt folder.

ⓧ I'd like to double check the full code

Here is the complete train.py so far. Make sure your file matches this exactly:

import torch
from model import MiniGPT, BLOCK_SIZE, device
import urllib.request
import os

# Hyperparameters
BATCH_SIZE = 64
MAX_ITERS = 3000
EVAL_INTERVAL = 300
LEARNING_RATE = 3e-4
EVAL_ITERS = 200

# Download Shakespeare dataset
DATA_PATH = "input.txt"
if not os.path.exists(DATA_PATH):
    url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
    urllib.request.urlretrieve(url, DATA_PATH)
    print("Downloaded Shakespeare dataset.")

# Read and process data
with open(DATA_PATH, "r", encoding="utf-8") as f:
    text = f.read()

# Build character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join([itos[i] for i in l])

# Create train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - BLOCK_SIZE, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + BLOCK_SIZE] for i in ix])
    y = torch.stack([data_split[i + 1 : i + BLOCK_SIZE + 1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

Your training data pipeline is ready. Next up, you will build the transformer architecture that actually learns from this data.

Build the Transformer Architecture

Your training pipeline is ready and waiting for a model to train. Now it's time to build the brain of your language model: the transformer architecture.

Every modern LLM (ChatGPT, LLaMA, Gemini) uses the same core building blocks you're about to write by hand. You'll implement each piece in model.py, starting with how the model represents text as numbers, then building the attention mechanism that lets it "focus" on relevant context.

In this step, get ready to:
  • Set up the model hyperparameters and device configuration.
  • Implement multi-head self-attention from scratch.
  • Assemble the full GPT model with feed-forward layers, layer normalization, and the language model head.
Implement token embeddings and positional encoding

Before a transformer can process text, it needs to convert characters into numbers the neural network can work with. Token embeddings map each character to a vector of numbers. Positional encoding adds information about where each character sits in the sequence, since transformers process all positions in parallel.

  • Open model.py in VS Code.
  • Add the imports, hyperparameters, and device configuration by pasting this code:
import torch
import torch.nn as nn
import torch.nn.functional as F

# Hyperparameters
VOCAB_SIZE = 65  # number of unique characters in Shakespeare
D_MODEL = 64  # embedding dimension
N_HEADS = 4  # number of attention heads
N_LAYERS = 4  # number of transformer blocks
BLOCK_SIZE = 128  # maximum context length
DROPOUT = 0.1
D_FF = 256  # feed-forward inner dimension

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

What does this code do?

These constants define the shape of your entire model. VOCAB_SIZE = 65 matches the 65 unique characters in the Shakespeare dataset your tokenizer found.

D_MODEL = 64 is the embedding dimension. Each character will be represented as a vector of 64 numbers. BLOCK_SIZE = 128 is the maximum context window. Your model can look back at up to 128 characters when predicting the next one.

The device line automatically selects your GPU if CUDA is available, falling back to CPU otherwise.

  • Save model.py.
  • Verify the constants load correctly by running this in your terminal:
python -c "from model import BLOCK_SIZE, device; print(f'Block size: {BLOCK_SIZE}, Device: {device}')"

You should see Block size: 128, Device: cuda (or Device: cpu if you don't have a GPU).

Seeing a ModuleNotFoundError?

Make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.

Also confirm you're running the command from inside your project folder where model.py lives.

Help me fix my model.py import error

Implement multi-head self-attention

Self-attention is the core mechanism that makes transformers powerful. It lets each character in the sequence look at every other character and decide which ones are relevant for predicting the next token.

  • In model.py, add the Head class below the device line by pasting this code:
class Head(nn.Module):
    """Single head of self-attention."""

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(D_MODEL, head_size, bias=False)
        self.query = nn.Linear(D_MODEL, head_size, bias=False)
        self.value = nn.Linear(D_MODEL, head_size, bias=False)
        self.register_buffer(
            "tril", torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE))
        )
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        # Compute attention scores
        weights = q @ k.transpose(-2, -1) * (C ** -0.5)
        weights = weights.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        weights = F.softmax(weights, dim=-1)
        weights = self.dropout(weights)
        # Weighted aggregation of values
        v = self.value(x)
        out = weights @ v
        return out

What does this code do?

Each attention head has three learnable projections: key, query, and value. Think of it like a search engine: the query asks "what am I looking for?", the key says "here's what I contain", and the value is "here's my actual content."

The tril buffer is a triangular mask that prevents the model from "cheating" by looking at future characters. The * (C ** -0.5) scaling prevents the dot products from growing too large, which would make softmax produce extreme values.

  • Save model.py.

You won't see a visible change yet. The next chunk connects multiple heads together.

  • Add the MultiHeadAttention class below the Head class by pasting this code:
class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(D_MODEL, D_MODEL)
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

Why multiple heads instead of one big attention?

Each head can learn to focus on different relationships. One head might learn to attend to nearby characters (spelling patterns), while another attends to characters far away (grammar structure).

With N_HEADS = 4 and D_MODEL = 64, each head operates on a 16-dimensional slice. The proj linear layer recombines all heads back into the full 64-dimensional space.

  • Save model.py.
  • Verify the attention module loads by running:
python -c "from model import MultiHeadAttention; print('Attention module loaded successfully')"

You should see Attention module loaded successfully printed in your terminal.

Seeing an IndentationError or SyntaxError?

Check that there are no extra blank lines inside the class methods. Make sure the class Head definition starts with exactly one blank line after the device line.

Help me debug my self-attention code syntax error

Assemble the full GPT model

With attention working, you need two more building blocks before assembling the full model. A feed-forward network processes each position independently after attention, and layer normalization keeps the numbers stable during training.

  • Add the FeedForward class below MultiHeadAttention by pasting this code:
class FeedForward(nn.Module):
    """A simple feed-forward network with ReLU activation."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(D_MODEL, D_FF),
            nn.ReLU(),
            nn.Linear(D_FF, D_MODEL),
            nn.Dropout(DROPOUT),
        )

    def forward(self, x):
        return self.net(x)

What does this code do?

The feed-forward network expands each position's representation from 64 dimensions up to 256 (D_FF), applies a ReLU activation to introduce non-linearity, then projects back down to 64. This "think" step happens independently at every position after the attention step has mixed information across positions.

  • Save model.py.

You won't see a change yet. The next chunk wraps attention and feed-forward into a reusable transformer block.

  • Add the TransformerBlock class below FeedForward by pasting this code:
class TransformerBlock(nn.Module):
    """Transformer block: attention followed by feed-forward."""

    def __init__(self):
        super().__init__()
        head_size = D_MODEL // N_HEADS
        self.sa = MultiHeadAttention(N_HEADS, head_size)
        self.ffwd = FeedForward()
        self.ln1 = nn.LayerNorm(D_MODEL)
        self.ln2 = nn.LayerNorm(D_MODEL)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

What does this code do?

Each TransformerBlock follows the same pattern: normalize, attend, add a residual connection, normalize, feed-forward, add another residual connection.

The x = x + ... pattern is called a residual connection. It lets gradients flow directly through the network during training, which makes deep models much easier to train. Your model stacks 4 of these blocks (N_LAYERS = 4).

  • Save model.py.
  • Verify the transformer block loads by running:
python -c "from model import TransformerBlock; print('Transformer block loaded')"

You should see Transformer block loaded in your terminal.

Seeing a NameError?

Make sure the classes are in order: Head first, then MultiHeadAttention, then FeedForward, then TransformerBlock. Python reads top to bottom, so each class must be defined before it's referenced.

Help me fix a NameError in my model.py

Now for the main event. The MiniGPT class ties everything together: embeddings, transformer blocks, and a final linear layer that predicts the next character.

  • Add the MiniGPT class below TransformerBlock by pasting this code:
class MiniGPT(nn.Module):
    """A small GPT-style language model."""

    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(VOCAB_SIZE, D_MODEL)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, D_MODEL)
        self.blocks = nn.Sequential(
            *[TransformerBlock() for _ in range(N_LAYERS)]
        )
        self.ln_f = nn.LayerNorm(D_MODEL)
        self.lm_head = nn.Linear(D_MODEL, VOCAB_SIZE)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B * T, C)
            targets = targets.view(B * T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does this code do?

The __init__ method sets up all the layers. token_embedding converts character IDs into 64-dimensional vectors. position_embedding adds a learned vector for each position (0 through 127). Adding these two embeddings together gives the model both "what" and "where" information.

The forward method runs data through the model. It embeds tokens, adds positions, passes through 4 transformer blocks, normalizes, then projects back to vocabulary size to get logits (raw prediction scores). If targets are provided, it also computes cross-entropy loss for training.

The generate method produces new text one character at a time. It takes the model's probability predictions and samples from them to pick each next character.

  • Save model.py.
  • Verify the full model loads and count its parameters by running:
python -c "from model import MiniGPT; m = MiniGPT(); print(f'Parameters: {sum(p.numel() for p in m.parameters()):,}')"

You should see a parameter count around 100,000-200,000. This is your complete model, roughly 1,000,000x smaller than GPT-2 but using the exact same architecture.

Seeing a RuntimeError or shape mismatch?

Double check that VOCAB_SIZE is set to 65 and D_MODEL is 64 at the top of model.py. A mismatch between D_MODEL and N_HEADS (64 must be evenly divisible by 4) causes shape errors.

Help me debug my MiniGPT model shape error

✔️ Awesome, I've got everything!

Great. Double check you've saved model.py and the import test printed a parameter count.

ⓧ I'd like to double check the full code

Here is the complete model.py file. Compare yours against it:

import torch
import torch.nn as nn
import torch.nn.functional as F

# Hyperparameters
VOCAB_SIZE = 65  # number of unique characters in Shakespeare
D_MODEL = 64  # embedding dimension
N_HEADS = 4  # number of attention heads
N_LAYERS = 4  # number of transformer blocks
BLOCK_SIZE = 128  # maximum context length
DROPOUT = 0.1
D_FF = 256  # feed-forward inner dimension

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Head(nn.Module):
    """Single head of self-attention."""

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(D_MODEL, head_size, bias=False)
        self.query = nn.Linear(D_MODEL, head_size, bias=False)
        self.value = nn.Linear(D_MODEL, head_size, bias=False)
        self.register_buffer(
            "tril", torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE))
        )
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        # Compute attention scores
        weights = q @ k.transpose(-2, -1) * (C ** -0.5)
        weights = weights.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        weights = F.softmax(weights, dim=-1)
        weights = self.dropout(weights)
        # Weighted aggregation of values
        v = self.value(x)
        out = weights @ v
        return out


class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(D_MODEL, D_MODEL)
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out


class FeedForward(nn.Module):
    """A simple feed-forward network with ReLU activation."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(D_MODEL, D_FF),
            nn.ReLU(),
            nn.Linear(D_FF, D_MODEL),
            nn.Dropout(DROPOUT),
        )

    def forward(self, x):
        return self.net(x)


class TransformerBlock(nn.Module):
    """Transformer block: attention followed by feed-forward."""

    def __init__(self):
        super().__init__()
        head_size = D_MODEL // N_HEADS
        self.sa = MultiHeadAttention(N_HEADS, head_size)
        self.ffwd = FeedForward()
        self.ln1 = nn.LayerNorm(D_MODEL)
        self.ln2 = nn.LayerNorm(D_MODEL)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x


class MiniGPT(nn.Module):
    """A small GPT-style language model."""

    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(VOCAB_SIZE, D_MODEL)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, D_MODEL)
        self.blocks = nn.Sequential(
            *[TransformerBlock() for _ in range(N_LAYERS)]
        )
        self.ln_f = nn.LayerNorm(D_MODEL)
        self.lm_head = nn.Linear(D_MODEL, VOCAB_SIZE)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B * T, C)
            targets = targets.view(B * T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Your transformer architecture is complete. Next up, you'll feed Shakespeare text through this model and train it to predict the next character.

Train Your Model

Your transformer architecture is fully built. Every layer, every attention head, every feed-forward network is wired together and waiting in model.py. But right now, all those weights are random numbers. The model has no idea what English looks like, let alone Shakespeare.

Training is how a model turns random weights into useful ones. You will feed it thousands of Shakespeare text samples and let it adjust its weights to predict the next character more accurately each time. After 3000 iterations on your GPU, it will generate recognizable Shakespeare-style text.

In this step, get ready to:
  • Set up the training loop with an optimizer and loss function.
  • Run training on your GPU and watch the loss decrease.
  • Generate sample text from your trained model to verify it works.
Set up the training loop with AdamW optimizer and cross-entropy loss

A training loop repeats three actions thousands of times: grab a batch of data, measure how wrong the model's predictions are (the loss), then nudge the weights to be less wrong (backpropagation). Over time, the loss drops and the model improves.

  • In train.py, scroll down to the bottom of the file (below the estimate_loss function).
  • Add the model creation and training loop by copying this code at the bottom of the file:
# Create model
model = MiniGPT().to(device)
param_count = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {param_count:,}")
print(f"Training on: {device}")

# Train
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

for iter in range(MAX_ITERS):
    # Evaluate periodically
    if iter % EVAL_INTERVAL == 0:
        losses = estimate_loss(model)
        print(
            f"Step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}"
        )

    # Get batch and compute loss
    xb, yb = get_batch("train")
    logits, loss = model(xb, yb)

    # Backpropagation
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

What does this code do?

This code has two parts: model setup and the training loop.

  • MiniGPT().to(device) creates your model and moves it to the GPU. The parameter count print confirms how many learnable weights your model has.
  • AdamW is an optimizer that adjusts weights using adaptive learning rates with decoupled weight decay. It is the standard optimizer for training transformers.
  • Every 300 iterations, the loop pauses to evaluate loss on both training and validation data. This lets you track whether the model is actually learning.
  • optimizer.zero_grad(), loss.backward(), and optimizer.step() are the three core lines of backpropagation: clear old gradients, compute new gradients, then update the weights.

What is cross-entropy loss?

Cross-entropy loss measures how far the model's predicted character probabilities are from the actual next character. If the model assigns high probability to the correct character, loss is low. If it assigns low probability, loss is high.

The F.cross_entropy call happens inside your MiniGPT.forward() method in model.py. The training loop just calls the model and reads the returned loss.

You won't see output yet. The next chunk adds the post-training logic that saves your model and generates text.

  • Now add the final evaluation, model save, and text generation code directly below the training loop:
# Final evaluation
losses = estimate_loss(model)
print(f"Final: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

# Save model
torch.save(model.state_dict(), "model.pth")
print("Model saved to model.pth")

# Generate sample text
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=200)
print("\n--- Generated Text ---")
print(decode(generated[0].tolist()))

What does this code do?

  • estimate_loss(model) runs one final evaluation after all 3000 training iterations complete, giving you the final loss numbers.
  • torch.save(model.state_dict(), "model.pth") saves your trained weights to a file. This is the checkpoint you will load later when serving the model as an API.
  • model.generate() uses the trained model to produce 200 characters of new text, starting from a zero token. This is your first glimpse of what your model learned.
  • Save train.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before running it.

ⓧ I'd like to double check the full code

import torch
from model import MiniGPT, BLOCK_SIZE, device
import urllib.request
import os

# Hyperparameters
BATCH_SIZE = 64
MAX_ITERS = 3000
EVAL_INTERVAL = 300
LEARNING_RATE = 3e-4
EVAL_ITERS = 200

# Download Shakespeare dataset
DATA_PATH = "input.txt"
if not os.path.exists(DATA_PATH):
    url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
    urllib.request.urlretrieve(url, DATA_PATH)
    print("Downloaded Shakespeare dataset.")

# Read and process data
with open(DATA_PATH, "r", encoding="utf-8") as f:
    text = f.read()

# Build character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join([itos[i] for i in l])

# Create train/val split
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - BLOCK_SIZE, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + BLOCK_SIZE] for i in ix])
    y = torch.stack([data_split[i + 1 : i + BLOCK_SIZE + 1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


# Create model
model = MiniGPT().to(device)
param_count = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {param_count:,}")
print(f"Training on: {device}")

# Train
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

for iter in range(MAX_ITERS):
    # Evaluate periodically
    if iter % EVAL_INTERVAL == 0:
        losses = estimate_loss(model)
        print(
            f"Step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}"
        )

    # Get batch and compute loss
    xb, yb = get_batch("train")
    logits, loss = model(xb, yb)

    # Backpropagation
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# Final evaluation
losses = estimate_loss(model)
print(f"Final: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

# Save model
torch.save(model.state_dict(), "model.pth")
print("Model saved to model.pth")

# Generate sample text
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=200)
print("\n--- Generated Text ---")
print(decode(generated[0].tolist()))
Run training on your GPU and watch the loss decrease

Your script is ready. When you run it, PyTorch will move all computations to your NVIDIA GPU. Training 3000 iterations on the RTX 3050 Ti takes roughly 2-5 minutes.

  • Run the training script by executing this command in your terminal:
python train.py

You should see output appearing every 300 steps showing the training and validation loss. It looks something like this:

Vocabulary size: 65 unique characters
Model parameters: 209,601
Training on: cuda
Step 0: train loss 4.1743, val loss 4.1724
Step 300: train loss 2.1532, val loss 2.1887
Step 600: train loss 1.8456, val loss 1.9102
...

What do these numbers mean?

The loss starts around 4.17 because the model is randomly guessing among 65 characters. A random guess would give a loss of about ln(65) = 4.17. As training progresses, the loss drops, meaning predictions are getting more accurate.

The gap between train loss and val loss tells you about overfitting. A small gap means the model generalizes well. A large gap means it memorized the training data without learning the underlying patterns.

Seeing a CUDA error?

  • If you see CUDA out of memory, close other GPU-intensive applications (games, other models) and try again. This model uses well under 1GB of VRAM so this error is rare.
  • If you see Training on: cpu instead of cuda, your PyTorch installation does not detect the GPU. Reinstall PyTorch with CUDA 12.6 support using the install command from Step 1.
  • If you see ModuleNotFoundError: No module named 'model', make sure you are running the command from the mini-gpt project folder where model.py lives.

Help me debug a CUDA or import error when running train.py

Generate sample text from your trained model to verify it works

After the training loop finishes, the script automatically generates 200 characters of new text. This is the moment of truth for your model.

At the bottom of your terminal output, you should see a --- Generated Text --- header followed by Shakespeare-like text. It will not be perfect English, but you should recognize word fragments, character names, and dramatic formatting (like uppercase names followed by colons).

What does the generated text quality tell me?

Your model is character-level with about 200K parameters. It learns patterns like word boundaries, line breaks, character names, and common letter sequences. It will not produce grammatically perfect sentences because that requires billions of parameters and word-level tokenization.

What matters is that the output clearly resembles a play script rather than random characters. If you see structure (capitalized names, colons, verse-like line breaks), your model learned real patterns from Shakespeare.

  • Confirm that model.pth was saved by running:
ls model.pth

You should see model.pth listed. This file contains your trained weights and is what you will load into the API in the next step.

Don't see model.pth?

  • Make sure the training script ran all the way to completion. If it was interrupted (Ctrl+C), the save line never executed. Run python train.py again and let it finish.
  • Check that your terminal is in the same folder as train.py. The file saves to whatever directory Python runs from.

Help me figure out why model.pth was not created

Your model is trained and saved. Next up, you will wrap it in a REST API so anyone can send a text prompt and get AI-generated Shakespeare back.

Deploy as a REST API

Your model is trained and generating Shakespeare-like text from a script. That's powerful, but nobody else can use it yet.

In this step, you will wrap your trained model in a FastAPI web server. This turns your model into a REST API that accepts text prompts over HTTP and returns generated text as JSON.

In this step, get ready to:
  • Build a FastAPI app that loads your trained model and tokenizer.
  • Create a /generate endpoint for text generation.
  • Test your API and see your model generate Shakespeare-like text.
Build the FastAPI app that loads your trained model

The API server needs two things before it can generate text: the same character-level tokenizer from training, and your trained model weights loaded from model.pth.

  • Open app.py in VS Code and add the following code:
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from model import MiniGPT, BLOCK_SIZE, device

# Load the tokenizer
DATA_PATH = "input.txt"
with open(DATA_PATH, "r", encoding="utf-8") as f:
    text = f.read()

chars = sorted(list(set(text)))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s if c in stoi]
decode = lambda l: "".join([itos[i] for i in l])

# Load trained model
model = MiniGPT().to(device)
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=True))
model.eval()

app = FastAPI(title="Mini GPT API")

What does this code do?

The first section rebuilds the character-level tokenizer by reading input.txt and creating the same encode and decode functions from training. The if c in stoi guard safely ignores any character the model never saw during training.

The second section creates a fresh MiniGPT instance, loads the trained weights from model.pth, and switches to eval() mode. Eval mode disables dropout so the model produces consistent outputs at inference time.

Why FastAPI instead of Flask?

FastAPI is the modern standard for Python APIs. It gives you automatic request validation through Pydantic, built-in interactive docs at /docs, and async support out of the box.

For ML model serving, automatic type validation catches bad input before it ever reaches your model. Flask would require manual validation for every request field.

  • Save app.py.

You won't see a change yet. The next section adds the endpoint that actually uses this setup.

Seeing an import error?

Make sure your virtual environment is still activated. You should see the environment name in your terminal prompt. If not, reactivate it before continuing.

help me fix an import error in my FastAPI app

Create the /generate endpoint

Your API needs a contract: what data does a request contain, and what does the response look like? Pydantic models define this contract, and FastAPI enforces it automatically.

  • Add the following Pydantic models below the app = FastAPI(title="Mini GPT API") line in app.py:
class GenerateRequest(BaseModel):
    prompt: str = ""
    max_length: int = 200


class GenerateResponse(BaseModel):
    generated_text: str
    prompt: str

What does this code do?

GenerateRequest defines what the caller sends: an optional text prompt and a max length (defaulting to 200 tokens). GenerateResponse defines what comes back: the generated text and the original prompt for reference.

FastAPI automatically validates incoming JSON against GenerateRequest. If someone sends a string where max_length should be an integer, the API returns a helpful error without your code needing to check.

You won't see a change yet. The next code block adds the endpoint functions that use these models.

  • Add the endpoint functions below the GenerateResponse class in app.py:
@app.post("/generate", response_model=GenerateResponse)
def generate_text(request: GenerateRequest):
    if request.prompt:
        encoded = encode(request.prompt)
        context = torch.tensor([encoded], dtype=torch.long, device=device)
    else:
        context = torch.zeros((1, 1), dtype=torch.long, device=device)

    with torch.no_grad():
        generated = model.generate(context, max_new_tokens=request.max_length)

    output_text = decode(generated[0].tolist())
    return GenerateResponse(generated_text=output_text, prompt=request.prompt)


@app.get("/")
def root():
    return {"message": "Mini GPT API is running. POST to /generate to create text."}

What does this code do?

The /generate endpoint encodes the prompt into token indices (or starts from a zero tensor if no prompt is given). It runs the model's generate method inside torch.no_grad() to save memory during inference.

The root / endpoint is a simple health check. It confirms the server is running without doing any generation.

  • Save app.py.

Your API code is complete. The next substep starts the server and tests it.

Not sure where to add the code?

Each new block goes directly below the previous one. The Pydantic classes sit between the app = FastAPI(...) line and the endpoint functions. Use the full-code tab below to compare your file.

help me check my app.py structure

✔️ Awesome, I've got everything!

Great. Double check you have saved app.py before moving on.

ⓧ I'd like to double check the full code

import torch
from fastapi import FastAPI
from pydantic import BaseModel
from model import MiniGPT, BLOCK_SIZE, device

# Load the tokenizer
DATA_PATH = "input.txt"
with open(DATA_PATH, "r", encoding="utf-8") as f:
    text = f.read()

chars = sorted(list(set(text)))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s if c in stoi]
decode = lambda l: "".join([itos[i] for i in l])

# Load trained model
model = MiniGPT().to(device)
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=True))
model.eval()

app = FastAPI(title="Mini GPT API")


class GenerateRequest(BaseModel):
    prompt: str = ""
    max_length: int = 200


class GenerateResponse(BaseModel):
    generated_text: str
    prompt: str


@app.post("/generate", response_model=GenerateResponse)
def generate_text(request: GenerateRequest):
    if request.prompt:
        encoded = encode(request.prompt)
        context = torch.tensor([encoded], dtype=torch.long, device=device)
    else:
        context = torch.zeros((1, 1), dtype=torch.long, device=device)

    with torch.no_grad():
        generated = model.generate(context, max_new_tokens=request.max_length)

    output_text = decode(generated[0].tolist())
    return GenerateResponse(generated_text=output_text, prompt=request.prompt)


@app.get("/")
def root():
    return {"message": "Mini GPT API is running. POST to /generate to create text."}
Test your API with curl

Time to start the server and see your model respond to HTTP requests.

  • Start the FastAPI development server by running this command:
fastapi dev app.py

What does this command do?

fastapi dev launches a development server with auto-reload enabled. It uses uvicorn under the hood and serves your app at http://127.0.0.1:8000.

  • Wait for the terminal to show that the server is running. You should see a message indicating the application is being served on http://127.0.0.1:8000.
  • Open a new terminal window (keep the server running in the first one).
  • Activate your virtual environment in the new terminal.
  • Test the root endpoint by running this command:
curl http://127.0.0.1:8000/

You should see a JSON response confirming the server is running.

  • Test text generation by sending a prompt to the /generate endpoint:
curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d "{\"prompt\": \"ROMEO:\", \"max_length\": 200}"

You should see a JSON response containing Shakespeare-like text generated by your model. The output starts with your prompt and continues with generated characters.

What should I see?

The response is a JSON object with generated_text and prompt fields. The generated text will be Shakespeare-flavoured but not perfect. It might resemble character dialogue or poetic lines. Every run produces different output because generation uses random sampling.

Try the interactive docs

Open http://127.0.0.1:8000/docs in your browser. FastAPI automatically generates a Swagger UI where you can test both endpoints interactively without writing curl commands.

Connection refused or server not starting?

Make sure you are running fastapi dev app.py from the same folder that contains app.py, model.py, model.pth, and input.txt. The server loads all four files at startup.

If you see an address-already-in-use error, another process is using port 8000. Close any other servers or use a different port.

help me debug my FastAPI server

Your model is now serving predictions over HTTP. You built a complete ML pipeline: from raw data to trained model to live API.

Secret mission

Advanced Text Generation Controls

Your model currently generates text with a fixed sampling strategy. Real LLM APIs like ChatGPT let users control creativity with temperature and top-k sampling. Add these controls to your API so users can dial the output from predictable Shakespeare to wild creative nonsense.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • FastAPI development server (running process in your terminal).
  • Python virtual environment (installed PyTorch and FastAPI packages).
  • Project folder on your Desktop containing model.py, train.py, app.py, input.txt, and model.pth.

Keep everything running

No action needed. Choose this if you're still actively building or experimenting with your model.

  • Your FastAPI server stays running in the terminal for continued testing.
  • Your virtual environment remains active with all packages installed.
  • Your trained model checkpoint is ready for immediate use.

Pause - I'll come back to this later

Shut down running processes to free up memory. Your files and virtual environment stay on disk for next time.

  • Stop the FastAPI server by pressing Ctrl+C in the terminal where it's running.
  • Deactivate your virtual environment by running:
deactivate

When you come back, reactivate the environment and restart the server.

Delete - I don't want to use this again

Remove all project files, the virtual environment, and the trained model. This frees up disk space completely.

  • Stop the FastAPI server by pressing Ctrl+C in the terminal where it's running.
  • Deactivate your virtual environment by running:
deactivate
  • Navigate to your Desktop and delete the project folder by running:

macOS

cd ~/Desktop
rm -rf mini-gpt

Windows

cd ~\Desktop
rmdir /s /q mini-gpt

This removes all Python files, the virtual environment, the trained model checkpoint, and the downloaded Shakespeare dataset.

That's a Wrap!

That's a Wrap!

You just built a language model from scratch. You now understand the exact architecture that powers ChatGPT, LLaMA, and every modern LLM.

You've learned how to:

  • Built a GPT-style transformer architecture by hand, including token embeddings, self-attention, feed-forward layers, and layer normalization.
  • Trained your model on Shakespeare text using your own GPU and watched the loss decrease over 3000 iterations.
  • Deployed your model as a REST API with FastAPI that accepts text prompts and returns AI-generated text on demand.
  • Added temperature and top-k sampling controls for creative text generation.

Ready to quiz yourself?