PythiloLesson 7 of 9 · progress kept in this browser

A tiny GPT

Put the pieces together and you have a GPT: token embeddings plus position embeddings, a stack of blocks, a final normalisation, and a linear head that gives a score for every token in the vocabulary, at every position.

import torch
from torch import nn

tokens = nn.Embedding(20, 16)
positions = nn.Embedding(8, 16)
ids = torch.tensor([[3, 1, 4, 1]])
where = torch.arange(4)
x = tokens(ids) + positions(where)
print(x.shape)  # torch.Size([1, 4, 16])

Training shows it windows of text and asks, at every position, for the next token. Unlike the bigram model, it can use everything before, so after the c it knows at comes next even though t alone is ambiguous.

Your task

Write TinyGPT(vocab, d, context, blocks): nn.Embedding(vocab, d) named tokens, nn.Embedding(context, d) named positions, an nn.Sequential of blocks Block(d, 4 * d), an nn.LayerNorm(d) named norm and an nn.Linear(d, vocab) named head. forward(ids) takes (batch, T) and returns logits (batch, T, vocab). Write train_gpt(model, ids, context, steps=300, lr=3e-3): each step, take 16 random windows of context + 1 ids, train with Adam on cross-entropy of every next id, and return the list of losses as floats.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.