Teaching a GPT From Scratch to Write Like a 19th-Century Brazilian Novelist

Chapter 2 of the book builds a tiny GPT, about 200 thousand parameters, character-level, from scratch. It's literally Andrej Karpathy's "Let's Build GPT from Scratch" video turned into prose, the book admits the source upfront. I ran the exercise with two notebooks side by side: the original notebook Karpathy himself cites as the video's source (gpt_dev.ipynb, trained on Shakespeare), and my own, where I swapped Shakespeare for Dom Casmurro, Machado de Assis's classic Brazilian novel, public domain text, to see whether the same tiny architecture picks up Portuguese too.

This post is the whole notebook cleaned up with you, cell by cell, full code, nothing skipped. It's not a summary, it's a dissection.

Downloading and checking the dataset

!wget https://www.gutenberg.org/cache/epub/55752/pg55752.txt

with open('pg55752.txt', 'r') as file:
    texto = file.read()

print("Tamanho do dataset", len(texto))
print("===== Trecho do texto =====")
print(texto[:1000])

Real output:

Tamanho do dataset 400944
===== Trecho do texto =====
The Project Gutenberg eBook of Dom Casmurro
...
Title: Dom Casmurro
Author: Machado de Assis

400,944 characters, including the standard Gutenberg header (license, title, author) that ships with the raw .txt file, that wasn't cleaned up, it's the text exactly as the site delivers it.

Vocabulary: every character becomes a number

caracteres = sorted(list(set(texto)))
tamanho_vocabulario = len(caracteres)
print("Vocabulário:", ''.join(caracteres), "\n")
print("Tamanho do Vocabulário: ", tamanho_vocabulario)

set(texto) grabs every unique character in the file, sorted() orders them, and the size of that set is the vocabulary. Output: 112 characters, against the original Shakespeare's 65, nearly double because of accent marks (á, ã, ç, é, í, ó, ô, õ, ú), curly quotes, and a few leftover typographic symbols from the Gutenberg text.

With the vocabulary set, the character-to-number translator can go both ways:

string_to_integer = {ch:i for i, ch in enumerate(caracteres)}
integer_to_string = {i:ch for i, ch in enumerate(caracteres)}

encode = lambda s: [string_to_integer[c] for c in s]
decode = lambda s: ''.join([integer_to_string[i] for i in s])

print(encode("Teste de encode"))
print(decode(encode("Teste de decode")))

string_to_integer maps each character to its position in the sorted vocabulary, integer_to_string is the reverse. encode swaps each letter of a string for its position, decode reverses the trip. Running encode("Teste de encode") returns [49, 63, 77, 78, 63, 1, 62, 63, 1, 63, 72, 61, 73, 62, 63], pure numbers, and decode back recovers the exact text.

Last step: push all of this into the shape PyTorch expects:

#Armazenando o texto encodado em um torch.Tensor
import torch
data = torch.tensor(encode(texto), dtype=torch.long)

print(data.shape, data.dtype)
print(data[:1000])

# Re-definindo a função decode para lidar com tensores PyTorch
def decode(s):
    # Converte cada elemento para Python int se for um tensor
    return ''.join([integer_to_string[i.item()] if isinstance(i, torch.Tensor) else integer_to_string[i] for i in s])

data becomes a tensor of 400,944 integers, the entire book encoded into a single strip. Notice decode gets rewritten here: the first version only knew how to handle plain Python int lists, this second version also accepts PyTorch tensors (checking isinstance and calling .item() when needed), because everything coming out of the model from here on arrives as a tensor.

Train/validation split and the context window size

#Split dados para treino e para avaliação
n = int(0.9*len(data))
treino = data[:n]
avaliacao = data[n:]

90% for training, 10% for validation, a simple positional cut (no shuffling, so the model never sees the end of the book during training).

#Definição do bloco de contexto para as previsões do modelo
# Isso significa que, para prever o próximo caractere,
#o modelo considerará no máximo os 8 caracteres anteriores.
#É um hiperparâmetro chave para a arquitetura do Transformer.
bloco_contexto = 8

# Este é um exemplo de como uma sequência de entrada x
#e seu alvo y (que é a mesma sequência deslocada em um token)
#seriam construídos a partir dos dados de treinamento.
treino[:bloco_contexto+1]
print(decode(treino[:bloco_contexto+1]))

Output: The Proje, the book's first 9 characters (after the header gets sliced off by indexing). bloco_contexto = 8 sets the context window, the model won't see more than 8 characters back at this early testing stage.

Bug #1: mixing training and validation

#cria sequencia de entrada x
x = treino[:bloco_contexto]

#cria sequência de targets y
#y é a mesma sequência que x mas deslocada 1 token para frente
#O modelo prevê y[t] dado x[:t+1]
y = avaliacao[1:bloco_contexto+1]

#loop para simular como o modelo vê a sequência
for tensor in range(bloco_contexto):
    contexto = x[:tensor+1]
    target = y[tensor]
    print(f"Quando tensor é {contexto} o target é {target}")

Found the bug right here: x comes from treino, but y comes from avaliacao. Two completely different chunks of the book, so the target the print shows has no real relationship to the context that came before it. Compare with the original notebook, which uses train_data on both lines: x = train_data[:block_size] and y = train_data[1:block_size+1], never switching splits mid-example.

The output matches the mistake: Quando tensor é tensor([49]) o target é 12 makes no sense at all (12 is the code for some random digit or symbol in the vocabulary, unrelated to "T" from "The"). The code doesn't crash (both splits share the same tensor shape), it just produces a nonsensical teaching example. Good news: this bug lives only in this manual illustration cell, the get_batch function coming up next never mixes splits.

Packing real batches

torch.manual_seed(1337)
batch_size = 4 # Quantas sequências independentes serão processadas em paralelo
block_size = 8 # Qual tamanho máximo de contexto para predições

def get_batch(split):
  #Criando um lote pequeno de dados de entrada x e targets y
  dados = treino if split == 'treino' else avaliacao
  ix = torch.randint(len(dados) - block_size, (batch_size,))
  x = torch.stack([data[i:i+block_size] for i in ix])
  y = torch.stack([dados[i+1:i+block_size+1] for i in ix])
  return x, y

xb, yb = get_batch('treino')
print("Inputs:")
print(xb.shape)
print(xb)
print('----------')
print("Targets:")
print(yb.shape)
print(yb)

for b in range(batch_size): #Dimensão batch
  for t in range(block_size): #Dimensão time
    contexto_time = xb[b, :t+1]
    target_time = yb[b,t]
    print(f"Quando input é {contexto_time.tolist()} - o target é {target_time}")

ix draws batch_size random starting positions inside the chosen split, and torch.stack piles up batch_size windows of block_size characters into a single matrix. Real output: Inputs: torch.Size([4, 8]), a batch of 4 sequences of 8 characters each, processed in parallel. The loop below is just for visualization: for each sequence in the batch, it shows the 8 progressive predictions it contains (predict the 2nd character knowing the 1st, predict the 3rd knowing the first 2, and so on), so a 4×8 batch actually teaches the model 32 prediction examples at once.

# Este output nos permite visualizar a forma exata
#e os valores numéricos das sequências de entrada
#que o modelo começará a processar.
print(xb) # Input em forma de tensor

Just reprints xb, nothing new, a useful visual checkpoint before assembling the model.

The first model, and the Bug #2 hiding inside it

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

class BigramLanguageModel(nn.Module):
  def __init__(self, tam_vocabulario):
    super().__init__()
    #cada token lê diretamente os logits para o
    #próximo token a partir de uma tabela de consulta
    self.token_embeeding_table = nn.Embedding(tam_vocabulario, tam_vocabulario)

  def forward(self, idx, targets=None):
    # idx e targets são ambos tensores de
    #inteiros de formato (B, T)
    logits = self.token_embeeding_table(idx) # (B, T, C)

    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):
    # idx é um array de índices (B, T)
    #no contexto atual
    for _ in range(max_new_tokens):
      #"pega" as predições
      logits, loss = self(idx)
      #foca apenas no último passo de tempo
      logits = logits[:, -1, : ] #Vira B, C
      # aplica softmax para obter probabilidades
      probabilidades = F.softmax(logits, dim=-1) # (B, C)
      # amostra da distribuição
      idx_next = torch.multinomial(probabilidades, num_samples=1) # (B, T+1)
    return idx

modelo = BigramLanguageModel(tamanho_vocabulario)
saida, loss = modelo(xb, yb)
print(saida.shape)
print(loss)

print("-----")

print(decode(modelo.generate(idx = torch.zeros((1,1), dtype=torch.long), max_new_tokens=100)[0].tolist()))

The dumbest possible model: nn.Embedding(tam_vocabulario, tam_vocabulario) is a lookup table where every character points straight to a vector as wide as the whole vocabulary (those vectors become the logits, with no context beyond the current character itself). forward computes loss via cross_entropy only when targets exists.

generate is where the bug lives: look at the for loop that samples and computes idx_next, but it never glues that idx_next back onto idx. The line idx = torch.cat((idx, idx_next), dim=1) is missing. The original notebook has that line, my version dropped it while typing. Result: generate always returns the original idx unchanged, so "generating 100 new tokens" produces zero in practice.

The real output confirms the damage: torch.Size([32, 112]) and tensor(5.0673, ...) for loss look normal, but after the "-----" the print(decode(...)) comes out empty. Not a runtime error, it's the exact symptom of the bug: idx is still the original [[0]] tensor, decoding that is just the vocabulary's character at code 0 (a newline), invisible in the print.

Training (badly) the bigram, for only 100 steps

#Criando um otimizador com pytorch
optimizer = torch.optim.AdamW(modelo.parameters(), lr=1e-3)

tamanho_batch = 32
for passo in range(100):
  # amostra de um lote de dados
  xb, yb = get_batch('treino')

  #avaliando o loss
  logits, loss = modelo(xb, yb)
  optimizer.zero_grad(set_to_none=True)
  loss.backward()
  optimizer.step()

print(loss.item())

With 112 possible classes and a random guess, the expected loss for an untrained model is -ln(1/112) ≈ 4.72. The untrained bigram started at 5.07 (seen in the previous cell), the same order of magnitude. After only 100 steps of crude training, the loss went to 5.19 (real output: 5.18871545791626), it actually ticked up slightly, because a bigram has no real memory to work with yet and 100 steps is barely anything.

print(decode(modelo.generate(idx = torch.zeros((1,1), dtype=torch.long), max_new_tokens=500)[0].tolist()))

Same bug from before, same consequence: empty output. This bigram model never actually generated real text in this notebook, the full model further down (which has the correct torch.cat) is the one that truly learns to write.

The self-attention trick, step by step

The chapter's core idea: every character needs to "talk" to the previous characters. The slow way is a loop, averaging everything that came before:

#O truque matemático na autoatenção
# exemplo simplificado que ilustra como
# a multiplicação de matrizes pode ser usada
# para uma "agregação ponderada"
torch.manual_seed(42)
a = torch.tril(torch.ones(3,3))
a = a / torch.sum(a, 1, keepdim=True)
b = torch.randint(0,10,(3,2)).float()
c = a @ b
print("a=")
print(a)
print('--')
print('b=')
print(b)
print('--')
print('c=')
print(c)

This warm-up example (3×3, small numbers) shows the pure trick: a is a lower-triangular matrix normalized by row (each row sums to 1), and multiplying a @ b already returns the running average of b row by row, no loop needed. Real output confirms it: the first row of c is just b[0] (average of 1 element), the second is the average of b[0] and b[1], and so on.

Now with real data, in three versions that should be equivalent:

# considere o seguinte exemplo simplificado:
torch.manual_seed(1337)
B,T,C = 4,8,2 # batch, time, channels
x = torch.randn(B,T,C)
x.shape

# Nós queremos x[b,t] = mean_{i<=t} x[b,i]
xbow = torch.zeros((B,T,C))
for b in range(B):
  for t in range(T):
    xprev = x[b,:t+1] # (t,C)
    xbow[b,t] = torch.mean(xprev, 0)

Version 1, the slow way: two nested for loops, averaging everything before each position, one batch and time step at a time.

# versão 2: usando multiplicação de matrizes para
# uma agregação ponderada
wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)
xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)
torch.allclose(xbow, xbow2)

Version 2, the same math via matrix multiplication, the same trick as the warm-up above. torch.allclose should return True (both versions compute the same thing). Real output: False.

# versão 3: usando softmax
tril = torch.tril(torch.ones(T, T))
wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
xbow3 = wei @ x
torch.allclose(xbow, xbow3)

Version 3, the same thing again, this time building the weights via softmax of a matrix masked with -inf above the diagonal (this matters in a moment, it's how the model will eventually learn weights that aren't just fixed averages). Real output again: False.

That bugged me enough to check the original notebook, and it shows the exact same result, False on both comparisons, with the exact same code. Not a bug either of us introduced: torch.allclose uses a fairly tight default tolerance, and the accumulated sum in a matrix multiply doesn't follow the same floating-point operation order as a loop-computed mean, so the two results land close but not bit-identical. It's a reminder worth more than a lot of theory: mathematically equivalent isn't the same as numerically identical, floating point has its own opinion about operation order.

Real self-attention

#versão 4: self-attention
torch.manual_seed(1337)
B,T,C = 4,8,32 #batch, time, channels
x = torch.randn(B,T,C)

#Vendo uma única "cabeça" performando a self-attention
head_size = 16
key = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)

k = key(x)   # (B, T, 16)
q = query(x) # (B, T, 16)
wei =  q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)

tril = torch.tril(torch.ones(T, T))
#wei = torch.zeros((T,T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)

v = value(x)
out = wei @ v
#out = wei @ x

out.shape

Here the averaging weights stop being fixed (1/2, 1/3...) and become learned and data-dependent: key, query and value are three trainable linear layers. Query asks "what am I looking for", Key answers "what do I have to offer", the product q @ k.transpose() decides how much each earlier position matters, masked_fill + softmax turns that into weights that sum to 1 and never peek at the future, and wei @ v runs the weighted average using Value (the actual content each position carries, not the key it uses to get found). Output: torch.Size([4, 8, 16]), a 16-dimensional vector per position, already mixing in context information.

wei[0]

Prints the attention weight matrix for the first item in the batch. Real output (trimmed):

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        ...

Notice the first row is [1, 0, 0, ...] (the first position can only attend to itself, it has no past) and each following row distributes weight quite differently from "simple average", the second row gives 84% weight to the second token and only 16% to the first. That's the model deciding, not a fixed formula.

Why divide by the square root of head_size

k = torch.randn(B,T, head_size)
q = torch.randn(B,T, head_size)
wei = q @ k.transpose(-2, -1) * head_size**-0.5

k.var()

Output: tensor(1.0449)

q.var()

Output: tensor(1.0700)

wei.var()

Output: tensor(1.0918)

k and q start with variance near 1 (default initialization). Without the head_size**-0.5 factor multiplying wei, the variance of the q @ k.transpose() product would explode proportionally to head_size, because each element of wei is the sum of head_size multiplications. Scaling by 1/√head_size keeps wei's variance near 1 too, and that matters because:

torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)

Output: tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872]), a relatively smooth distribution.

torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # gets too peaky, converges to one-hot

Output: tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000]), almost all the weight (80%) stuck on one value.

Same input numbers, just multiplied by 8, and softmax already turns into nearly a "pick one winner" (one-hot). If wei weren't scaled, the attention weights would always live in that "peaky" regime, the model would attend to a single token and ignore the rest, instead of learning a real mixture.

Layer normalization

class LayerNorm1d:

  def __init__(self, dim, eps=1e-5, momentum=0.1):
    self.eps = eps
    self.gamma = torch.ones(dim)
    self.beta = torch.zeros(dim)

  def __call__(self, x):
    # calcular a passagem direta (forward pass)
    xmean = x.mean(1, keepdim=True) # média do lote
    xvar = x.var(1, keepdim=True) # Variância do lote
    # normalizar para variância unitária
    xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
    self.out = self.gamma * xhat + self.beta
    return self.out

  def parameters(self):
    return [self.gamma, self.beta]

torch.manual_seed(1337)
module = LayerNorm1d(100)
# tamanho de lote de 32 vetores de 100 dimensões
x = torch.randn(32, 100)
x = module(x)
x.shape

A from-scratch implementation of layer normalization, without using the built-in nn.LayerNorm, just to see the math up close. xmean and xvar compute mean and variance, xhat normalizes, gamma/beta let the model relearn a different scale/shift if it needs to (they start at 1 and 0, meaning no effect at all, until training adjusts them).

# média e desvio padrão de uma característica
#em todas as entradas do lote
x[:,0].mean(), x[:,0].std()

Output: (tensor(0.1469), tensor(0.8803)), the first feature, looking across all 32 entries in the batch, isn't exactly mean 0/std 1.

# média e desvio padrão de uma única entrada do
# lote, considerando suas características
x[0,:].mean(), x[0,:].std()

Output: (tensor(-9.5367e-09), tensor(1.0000)), practically 0 and exactly 1. That's the key difference between layer norm and batch norm: here normalization happens per individual example (looking at that one row's 100 dimensions), not per feature across the whole batch, which is why the second check lands perfectly and the first one doesn't.

Bug #3: the Fold gotcha, and the full model

import torch
import torch.nn as nn
from torch.nn import functional as Fold

#hiperparâmetros
batch_size = 16 # Número de sequências independentes que serão processadas em paralelo
block_size = 32 # qual o contexto máximo de predições
max_iters = 5000 #máximo de iterações
eval_interval = 100
learning_rate = 1e-3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 64
n_head = 4
n_layer = 4
dropout = 0.0

torch.manual_seed(1337)

#!wget https://www.gutenberg.org/cache/epub/55752/pg55752.txt
with open('pg55752.txt', 'r') as f:
  texto = f.read()

#Apenas caracteres únicos(limpeza)
chars = sorted(list(set(texto)))
tam_vocabulario = len(chars)
# criando um mapeamento de caracteres para inteiros
string_to_int = { ch:i for i,ch in enumerate(chars) }
int_to_string = { i:ch for i,ch in enumerate(chars) }

# encoder: toma uma string, devolve uma lista de inteiros
encode = lambda s: [string_to_int[c] for c in s]

# decoder: toma uma lista de inteiros, devolve uma string
decode = lambda l: ''.join([int_to_string[i] for i in l])

#splitando o dataset
data = torch.tensor(encode(texto), dtype=torch.long)
n = int(0.9*len(data))
dados_treino = data[:n]
dados_avaliacao = data[n:]
print(f"Dispositivo: {device}")
print(f"Tamanho do texto: {len(texto)} caracteres")
print(f"Tamanho do vocabulario: {tam_vocabulario}")
print(f"Treino: {len(dados_treino)} | Avaliacao: {len(dados_avaliacao)}")

#carregamento dos dados
def get_batch(split):
  #gerando um lote pequeno de dados de input x e targets y
  dados = dados_treino if split == 'treino' else dados_avaliacao
  ix = torch.randint(len(dados) - block_size, (batch_size,))
  x = torch.stack([dados[i:i+block_size] for i in ix])
  y = torch.stack([dados[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 perda_estimada():
  out = {}
  modelo.eval()
  for split in ['treino', 'avaliacao']:
    losses = torch.zeros(eval_iters)
    for k in range(eval_iters):
      X, Y = get_batch(split)
      logits, loss = modelo(X, Y)
      losses[k] = loss.item()
    out[split] = losses.mean()

  modelo.train()
  return out

class Head(nn.Module):
  """ uma cabeça de self-attention """

  def __init__(self, head_size):
    super().__init__()
    self.key = nn.Linear(n_embd, head_size, bias=False)
    self.query = nn.Linear(n_embd, head_size, bias=False)
    self.value = nn.Linear(n_embd, 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)   # (B,T,C)
    q = self.query(x) # (B,T,C)
    # calcular pontuações de atenção ("afinidades")
    wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
    wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
    wei = F.softmax(wei, dim=-1) # (B, T, T)
    wei = self.dropout(wei)
    # performar a média ponderada das valores
    v = self.value(x) # (B,T,C)
    out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
    return out

class MultiHeadAttention(nn.Module):
  """ Múltiplas cabeças de self-attention em paralelo """

  def __init__(self, numero_cabecas, tamanho_cabeca):
    super().__init__()
    self.heads = nn.ModuleList([Head(tamanho_cabeca) for _ in range(numero_cabecas)])
    self.proj = nn.Linear(n_embd, n_embd)
    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 FeedFoward(nn.Module):
  """ Uma camada linear simples seguida de uma função de ativação ReLU """

  def __init__(self, n_embd):
    super().__init__()
    self.net = nn.Sequential(
        nn.Linear(n_embd, 4 * n_embd),
        nn.ReLU(),
        nn.Linear( 4 * n_embd, n_embd),
        nn.Dropout(dropout)
    )

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

class Block(nn.Module):
  """Bloco Transformer: comunicação seguida de computação"""

  def __init__(self, n_embd, n_head):
    # n_embd: embedding dimension, n_head:
    #quantidade de cabeças desejadas
    super().__init__()
    head_size = n_embd // n_head
    self.sa = MultiHeadAttention(n_head, head_size)
    self.ffwd = FeedFoward(n_embd)
    self.ln1 = nn.LayerNorm(n_embd)
    self.ln2 = nn.LayerNorm(n_embd)

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

Found the third bug right on the second line: from torch.nn import functional as Fold. Fold, not F. Yet the entire Head class uses F.softmax, with no error at all. Why? Because a much earlier cell (the first bigram) had already run from torch.nn import functional as F, and that F stays alive in the Colab kernel's memory even without re-running that earlier cell. If this were a .py script run from scratch, it would've crashed with NameError immediately. It's the classic notebook trap: the order you run cells in matters more than the order they appear on screen.

The real output from this cell confirms the rest ran fine: Dispositivo: cuda, Tamanho do texto: 400944 caracteres, Tamanho do vocabulario: 112, Treino: 360849 | Avaliacao: 40095.

The rest of the cell stacks the pieces: Head is one attention head (what we already saw above, now made reusable). MultiHeadAttention runs several Heads in parallel and concatenates the result. FeedFoward is a two-layer linear network with ReLU in between, processing each position on its own (not looking at its neighbors). Block wraps everything with a residual connection (the x = x + ... instead of x = ..., letting gradients flow straight through the layers) and layer norm before each sub-layer.

#um super simples modelo bigram
class BigramLanguageModel(nn.Module):
  def __init__(self):
    super().__init__()
    # cada token lê diretamente os logits para o próximo token a partir de uma
    # tabela de consulta
    self.token_embeeding_table = nn.Embedding(tam_vocabulario, n_embd)
    self.position_embeeding_table =nn.Embedding(block_size, n_embd)
    self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
    self.ln_f = nn.LayerNorm(n_embd)
    self.lm_head = nn.Linear(n_embd, tam_vocabulario)

  def forward(self, idx, targets=None):
    B, T = idx.shape

    # idx e targets são ambos tensores de inteiros de formato (B, T)
    tok_emb = self.token_embeeding_table(idx) # (B, T,C)
    pos_emb = self.position_embeeding_table(torch.arange(T, device=device)) # (T, C)
    x = tok_emb + pos_emb # (B, T,C)
    x = self.blocks(x) # (B, T,C)
    x = self.ln_f(x) # (B, T,C)
    logits = self.lm_head(x) # (B, T, vocab_size)

    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):
    # idx é um array de índices (B, T)
    # no contexto atual
    for _ in range(max_new_tokens):
      # corta idx para os últimos tokens de tamanho block_size
      idx_cond = idx[:, -block_size:]
      #captura predições
      logits, loss = self(idx_cond)
      #foca apenas no último passo de tempo
      logits = logits[:, -1, : ] # (B, C)
      #aplica softmax para capturar probabilidades
      probabilidades = F.softmax(logits, dim=-1) # (B, C)
      # amostra da distribuição
      idx_next = torch.multinomial(probabilidades, num_samples=1) # (B, 1)
      # anexa o índice amostrado à sequência em formação
      idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
    return idx

Despite the name (kept the original "bigram" comment even though it's not a bigram anymore), this is the full class: token_embeeding_table plus position_embeeding_table (summing the two lets every position know both "who I am" and "where I am"), blocks stacks n_layer Transformer blocks in sequence, ln_f normalizes at the end, lm_head projects back to vocabulary size. And this time generate does have the idx = torch.cat((idx, idx_next), dim=1) line, fixing the first bigram's bug, which is why only this model can actually generate real text.

Running it for real: hyperparameters, training, generation

modelo = BigramLanguageModel()
m = modelo.to(device)

# imprimir o número de parâmetros no modelo
print(sum(p.numel() for p in m.parameters())/1e6, 'M Número de parâmetros')

#criando um otimizador PyTorch
optimizer = torch.optim.AdamW(modelo.parameters(), lr=learning_rate)

for  iter in range(max_iters):
  # de tempos em tempos, avalie a perda nos conjuntos de treino e validação
  if iter % eval_interval == 0 or iter == max_iters - 1:
    losses = perda_estimada()
    print(f"step {iter}: train loss {losses['treino']:.4f}, val loss {losses['avaliacao']:.4f}")

  # amostra de um lote de dados
  xb, yb = get_batch('treino')

  #avaliando o loss
  logits, loss = modelo(xb, yb)
  optimizer.zero_grad(set_to_none=True)
  loss.backward()
  optimizer.step()

# gerar a partir do modelo
context = torch.zeros((1,1), dtype=torch.long, device=device)
print(decode(modelo.generate(context, max_new_tokens=2000)[0].tolist()))

This is the cell that actually trains: 0.215792 M parameters (215,792, to be exact), 5,000 iterations, evaluating train and validation loss every 100 steps. First output line: step 0: train loss 4.8673, val loss 4.8625, near the expected random guess. Last: step 4999: train loss 1.6304, val loss 2.6720.

The two training runs, side by side

Same hyperparameters in both notebooks (batch_size=16, block_size=32, n_embd=64, n_head=4, n_layer=4, max_iters=5000), the only thing that changes is the text:

Training loss ended up similar for both, but Dom Casmurro's validation loss came out noticeably worse. Two real reasons, no invention: first, I trained on less than a third of the text volume (400 thousand characters against 1.1 million), and the validation set shrank proportionally too (about 40 thousand characters against about 111 thousand), so the model had fewer examples to generalize from. Second, the bigger vocabulary (112 against 65) means more possible classes to get right at every position, and Portuguese morphology (rich verb conjugation, a lot more suffix variation than Shakespeare's early-modern English) is simply harder for a model this size to absorb. Not a code bug, the identical architecture just hit a harder dataset with less ammunition.

What comes out of generation

After 5,000 steps, generating from an empty context, the Dom Casmurro model spits this out (no punctuation or full-sentence sense, but look closely):

...abnoha (força-lhe o sappagorla, affendado escontraga-se, não vei nelles, prima Justina, como ella devi enferiu o não dizia uma paertada continueira... de Escobar cegos de enho qué está acconterno gostaria...

"Justina" and "Escobar" are real character names from the novel. The model doesn't know who Escobar is, doesn't understand the plot, it's 215 thousand parameters of pure character-level statistical pattern. But it learned that those letter sequences show up together often enough to reproduce whole proper names out of the noise. The same thing happens in the English version, which learned to put "KING RICHARD III:" and "QUEEN VINCENTIO:" in the right all-caps play-script format, without understanding a word of what the character says next.

Wrapping up

Practical application

The numbers in the table above come from the two real training runs, step by step, logged during each notebook's 5,000-iteration run (every 100 steps). Zoom the chart and watch the gap between train and validation: on Shakespeare the two lines stay glued together right to the end, on Dom Casmurro they peel apart early and never close back up, the exact visual signature of what the table above already said in numbers.

Loading real data...

Same architecture, same hyperparameters, two different languages, and the chart alone already shows which one gave the tiny 215-thousand-parameter model more trouble to learn.