Assemble the Model in Python

Time to put every part in one model, the way GPT models are built: a token embedding and a position embedding, added together; a stack of transformer blocks; a last LayerNorm; and an output layer, nn.Linear(32, vocab_size), that turns each position's 32 numbers into one score per character. These scores are the logits: at each position, the model's guess for the character that comes next. Here vocab_size is 19: the names use 18 different letters, plus the newline.

nn.Sequential(Block(), Block()) runs the blocks one after the other: the second reads what the first gave back. Every layer keeps its weights in tensors, and p.numel() gives back how many numbers a tensor holds, so adding it up over model.parameters() counts the weights:

layers = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
count = 0
for p in layers.parameters():
    count += p.numel()
print(count)   # 4*8 + 8 + 8*2 + 2 = 58

Run on a batch, your code and the checks after it print:

weights: 26771
x: [32, 8] logits: [32, 8, 19]
  token embeddings: 608
  position embeddings: 256
  2 blocks: 25216
  final norm: 64
  output layer: 627
loss before training: 2.986 | a blind guess among 19 characters: 2.944
probabilities for the next character, sum: 1.0
change character 7: logits 0-6 stay the same: True
a window of 5: [32, 5, 19]

About 27,000 weights, almost all in the two blocks. Untrained, the model guesses no better than chance: its loss is close to ln(19), the loss of a blind guess among 19 characters. The next lesson trains it.

Your task: the parts from the last lessons are ready, and so is a TinyGPT with its two embeddings. In __init__, store two blocks in an nn.Sequential called self.blocks, an nn.LayerNorm(n_embd) in self.norm and an nn.Linear(n_embd, vocab_size) in self.out. In forward, below x = tok + pos, pass x through the blocks and the norm, and give back self.out(x).

Assignment
Follow the Coding Tutorial to add the blocks, the last LayerNorm and the output layer to TinyGPT, and run x through them in forward.

Hint
The output layer goes from n_embd to vocab_size. In forward, store each result back in x, and return the output layer's scores at the end.