Ensinando um GPT do Zero a Escrever Feito Machado de Assis

O Capítulo 2 do livro constrói um GPT minúsculo, uns 200 mil parâmetros, character-level, do zero. É literalmente o vídeo "Let's Build GPT from Scratch" do Andrej Karpathy reescrito em prosa, o próprio livro admite a fonte na cara. Eu segui o exercício com dois notebooks lado a lado: o notebook original que o próprio Karpathy cita como fonte do vídeo (gpt_dev.ipynb, treinado em Shakespeare), e o meu, onde troquei Shakespeare por Dom Casmurro do Machado de Assis, texto de domínio público, pra ver se a mesma arquitetura minúscula aprende português também.

Esse post é o notebook inteiro passado a limpo com você, célula por célula, código completo, sem pular nada. Não é resumo, é dissecação.

Baixando e conferindo o 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])

Saída real:

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

400.944 caracteres, incluindo o cabeçalho padrão do Gutenberg (licença, título, autor) que vem junto do arquivo .txt bruto, isso não foi limpo, é o texto exatamente como o site entrega.

Vocabulário: cada caractere é um número

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

set(texto) pega cada caractere único que aparece no arquivo, sorted() ordena, e o tamanho desse conjunto é o vocabulário. Saída: 112 caracteres, contra os 65 do Shakespeare original, quase o dobro por causa de acento (á, ã, ç, é, í, ó, ô, õ, ú), aspas curvas, e uns símbolos tipográficos que sobraram do texto do Gutenberg.

Com o vocabulário definido, dá pra montar o tradutor caractere-número dos dois lados:

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 é um dicionário caractere → posição no vocabulário ordenado, integer_to_string é o inverso. encode troca cada letra de uma string pela sua posição, decode faz o caminho de volta. Rodar encode("Teste de encode") devolve [49, 63, 77, 78, 63, 1, 62, 63, 1, 63, 72, 61, 73, 62, 63], uma lista de número pura, e decode de volta recupera o texto igualzinho.

Só falta jogar isso tudo pro formato que o PyTorch entende:

#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 vira um tensor de 400.944 números inteiros, o livro inteiro codificado numa tira só. Repara que decode foi reescrita aqui: a primeira versão só sabia lidar com lista de int do Python puro, essa segunda versão também aceita tensor do PyTorch (checando isinstance e chamando .item() quando precisa), porque daqui pra frente tudo que sai do modelo vem em formato de tensor.

Split de treino/validação e o tamanho da janela

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

90% pra treino, 10% pra avaliação, corte simples por posição (sem embaralhar, então o modelo nunca vê o final do livro durante o treino).

#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]))

Saída: The Proje, os primeiros 9 caracteres do livro (depois do cabeçalho ser cortado pela indexação). bloco_contexto = 8 define o context window, o modelo nunca vai enxergar mais que 8 caracteres pra trás nessa fase inicial de teste.

Bug #1: misturando treino com avaliação

#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}")

Achei o bug aqui: x vem de treino, mas y vem de avaliacao. São dois pedaços de texto completamente diferentes do livro, então o alvo que o print mostra não tem nenhuma relação real com o contexto que veio antes dele. Compara com o notebook original, que usa train_data nas duas linhas: x = train_data[:block_size] e y = train_data[1:block_size+1], sem trocar de split no meio.

A saída bate com o erro: Quando tensor é tensor([49]) o target é 12 não faz o menor sentido (12 é o código de um dígito ou símbolo qualquer do vocabulário, sem relação com "T" de "The"). O código não quebra (os dois splits têm o mesmo formato de tensor), só produz um exemplo didático sem sentido. Boa notícia: esse bug mora só nessa célula de ilustração manual, a função get_batch que vem a seguir nunca mistura splits.

Empacotando em lotes de verdade

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 sorteia batch_size posições aleatórias de início dentro do split escolhido, e torch.stack empilha batch_size janelas de block_size caracteres numa matriz só. Saída real: Inputs: torch.Size([4, 8]), um lote de 4 sequências de 8 caracteres cada, processadas em paralelo. O loop de baixo só serve pra visualizar: pra cada sequência do lote, mostra as 8 previsões progressivas que ela contém (prever o 2º caractere sabendo o 1º, prever o 3º sabendo os 2 primeiros, e assim por diante), então um lote de 4×8 na verdade ensina o modelo em 32 exemplos de previsão de uma vez só.

# 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

Só reimprime xb, sem novidade, útil como checkpoint visual antes de montar o modelo.

O primeiro modelo, e o Bug #2 escondido nele

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()))

O modelo mais bobo possível: nn.Embedding(tam_vocabulario, tam_vocabulario) é uma tabela onde cada caractere aponta direto pra um vetor do tamanho do vocabulário inteiro (esses vetores viram os logits, sem olhar nenhum contexto além do próprio caractere atual). forward calcula a loss via cross_entropy só quando existe targets.

generate é onde mora o bug: repara no for que faz a amostragem e calcula idx_next, mas nunca gruda esse idx_next de volta em idx. Falta a linha idx = torch.cat((idx, idx_next), dim=1). O notebook original tem essa linha, a minha versão perdeu ela na hora de digitar. Resultado: generate sempre devolve o idx original sem mudança nenhuma, então gerar "100 novos tokens" na prática gera zero.

Saída real confirma o estrago: torch.Size([32, 112]) e tensor(5.0673, ...) pra loss aparecem normais, mas depois do "-----" o print(decode(...)) sai vazio. Não é erro de execução, é o sintoma exato do bug: idx continua sendo o tensor [[0]] original, decodificar isso vira só o caractere de código 0 do vocabulário (uma quebra de linha), invisível no print.

Treinando (mal) o bigram, por só 100 passos

#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())

Com 112 classes possíveis e chute aleatório, a perda esperada de um modelo sem treino nenhum é -ln(1/112) ≈ 4,72. O bigram sem treino começou em 5,07 (visto na célula anterior), na mesma ordem de grandeza. Depois de só 100 passos de treino grosseiro, a perda foi pra 5,19 (saída real: 5.18871545791626), na verdade subiu um pouco, porque um bigram não tem memória nenhuma pra usar direito ainda e 100 passos é muito pouco.

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

Mesmo bug de antes, mesma consequência: saída vazia. Esse modelo bigram nunca chegou a gerar texto de verdade nesse notebook, o modelo completo mais na frente (que tem o torch.cat certinho) é quem realmente aprende a escrever.

O truque da autoatenção, passo a passo

A ideia central do capítulo: cada caractere precisa "conversar" com os caracteres anteriores. O jeito mais lento é em loop, calculando a média de tudo que veio antes:

#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)

Esse exemplo de aquecimento (3×3, números pequenos) mostra o truque puro: a é uma matriz triangular inferior normalizada por linha (cada linha soma 1), e multiplicar a @ b já devolve a média acumulada de b linha a linha, sem precisar de loop nenhum. Saída real confirma: a primeira linha de c é só b[0] (média de 1 elemento), a segunda é a média de b[0] e b[1], e assim por diante.

Agora com dado de verdade, em três versões que deveriam ser equivalentes:

# 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)

Versão 1, o jeito lento: dois for aninhados, calculando a média de tudo que veio antes de cada posição, um batch e um tempo de cada vez.

# 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)

Versão 2, o mesmo resultado matemático via multiplicação de matriz, igual ao aquecimento de cima. torch.allclose deveria devolver True (as duas versões calculam a mesma coisa). Saída real: 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)

Versão 3, o mesmo de novo, só que construindo os pesos via softmax de uma matriz mascarada com -inf no triângulo superior (isso vira importante daqui a pouco, é assim que o modelo vai aprender pesos que não são só médias fixas). De novo, saída real: False.

Isso me deixou intrigado, então conferi o notebook original, e ele mostra exatamente o mesmo resultado, False nas duas comparações, com o mesmo código. Não é bug de ninguém: torch.allclose usa uma tolerância padrão bem apertada, e a soma acumulada numa multiplicação de matriz não segue a mesma ordem de operações em ponto flutuante que a média calculada em loop, então os dois resultados ficam próximos, mas não idênticos bit a bit. É um lembrete que vale mais que muita teoria: matematicamente equivalente não é sinônimo de numericamente idêntico, ponto flutuante tem opinião própria sobre a ordem das contas.

Self-attention de verdade

#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

Aqui os pesos da média deixam de ser fixos (1/2, 1/3...) e viram aprendidos e dependentes do dado: key, query e value são três camadas lineares treináveis. Query pergunta "o que eu tô procurando", Key responde "o que eu tenho pra oferecer", o produto q @ k.transpose() decide o quanto cada posição anterior importa, masked_fill + softmax transforma isso em pesos que somam 1 e nunca espiam o futuro, e wei @ v faz a média ponderada usando Value (o conteúdo real que cada posição carrega, não a chave que ela usa pra ser encontrada). Saída: torch.Size([4, 8, 16]), um vetor de 16 dimensões por posição, já misturando informação do contexto.

wei[0]

Imprime a matriz de pesos de atenção do primeiro item do lote. Saída real (recortada):

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],
        ...

Repara que a primeira linha é [1, 0, 0, ...] (a primeira posição só pode prestar atenção nela mesma, não tem passado) e cada linha seguinte distribui peso de um jeito bem diferente de "média simples", tipo a segunda linha dando 84% de peso pro segundo token e só 16% pro primeiro. Isso é o modelo decidindo, não uma fórmula fixa.

Por que dividir pela raiz de 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()

Saída: tensor(1.0449)

q.var()

Saída: tensor(1.0700)

wei.var()

Saída: tensor(1.0918)

k e q começam com variância perto de 1 (inicialização padrão). Sem o fator head_size**-0.5 multiplicando wei, a variância do produto q @ k.transpose() explodiria proporcionalmente ao tamanho de head_size, porque cada elemento de wei é a soma de head_size multiplicações. Escalar por 1/√head_size mantém a variância de wei perto de 1 também, e isso importa porque:

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

Saída: tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872]), uma distribuição relativamente suave.

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

Saída: tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000]), quase todo o peso (80%) grudado num valor só.

Mesmos números de entrada, só multiplicados por 8, e o softmax já vira quase um "escolhe só um vencedor" (one-hot). Se wei não fosse escalado, os pesos de atenção ficariam sempre nesse regime "picudo", o modelo prestaria atenção num único token e ignoraria o resto, em vez de aprender uma mistura de verdade.

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

Uma implementação de layer normalization do zero, sem usar nn.LayerNorm pronto, só pra entender a conta por dentro. xmean e xvar calculam média e variância, xhat normaliza, gamma/beta deixam o modelo reaprender uma escala/deslocamento diferente se precisar (começam em 1 e 0, ou seja, sem efeito, até o treino ajustar).

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

Saída: (tensor(0.1469), tensor(0.8803)), a primeira feature, olhando as 32 entradas do lote, não fica exatamente em média 0/desvio 1.

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

Saída: (tensor(-9.5367e-09), tensor(1.0000)), praticamente 0 e exatamente 1. Essa é a diferença chave entre layer norm e batch norm: aqui a normalização acontece por exemplo individual (olhando as 100 dimensões daquela linha), não por feature através do lote inteiro, por isso a segunda conta bate perfeito e a primeira não.

Bug #3: a pegadinha do Fold, e o modelo completo

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

Achei o terceiro bug logo na segunda linha: from torch.nn import functional as Fold. Fold, não F. Só que toda a classe Head usa F.softmax, sem o menor erro. Por quê? Porque uma célula bem anterior (a do primeiro bigram) já tinha rodado from torch.nn import functional as F, e esse F continua vivo na memória do kernel do Colab, mesmo sem eu reexecutar aquela célula. Se isso fosse um script .py rodado do zero, ia estourar NameError na hora. É a pegadinha clássica de notebook: a ordem que você executa as células importa mais do que a ordem que elas aparecem na tela.

Saída real dessa célula confirma que o resto rodou liso: Dispositivo: cuda, Tamanho do texto: 400944 caracteres, Tamanho do vocabulario: 112, Treino: 360849 | Avaliacao: 40095.

O resto da célula empilha as peças: Head é uma cabeça de atenção (o que já vimos ali em cima, só que agora reaproveitável). MultiHeadAttention roda várias Head em paralelo e concatena o resultado. FeedFoward é uma rede de duas camadas lineares com ReLU no meio, processando cada posição sozinha (sem olhar as vizinhas). Block junta tudo com conexão residual (o x = x + ... em vez de x = ..., deixa o gradiente fluir direto pelas camadas) e layer norm antes de cada sub-camada.

#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

Apesar do nome (deixei o comentário original "bigram" mesmo não sendo mais um bigram de verdade), essa é a classe completa: token_embeeding_table mais position_embeeding_table (soma dos dois dá pra cada posição saber "quem eu sou" e "onde eu tô"), blocks empilha n_layer blocos Transformer em sequência, ln_f normaliza no final, lm_head projeta de volta pro tamanho do vocabulário. E o generate dessa vez tem a linha idx = torch.cat((idx, idx_next), dim=1), corrigindo o bug do primeiro bigram, é por isso que só esse modelo consegue gerar texto de verdade.

Rodando de verdade: hiperparâmetros, treino, geração

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()))

Essa é a célula que treina de verdade: 0,215792 M parâmetros (215.792, pra ser exato), 5.000 iterações, avaliando a perda de treino e validação a cada 100 passos. Primeira linha da saída: step 0: train loss 4.8673, val loss 4.8625, perto do chute aleatório esperado. Última: step 4999: train loss 1.6304, val loss 2.6720.

Os dois treinos, lado a lado

Mesmos hiperparâmetros nos dois notebooks (batch_size=16, block_size=32, n_embd=64, n_head=4, n_layer=4, max_iters=5000), a única coisa que muda é o texto:

A perda de treino ficou parecida nos dois, mas a de validação do Dom Casmurro ficou bem pior. Duas razões reais, sem inventar: primeiro, eu treinei com menos de um terço do volume de texto (400 mil caracteres contra 1,1 milhão), e o conjunto de validação também encolheu proporcionalmente (uns 40 mil caracteres contra uns 111 mil), então o modelo teve menos exemplo pra generalizar. Segundo, o vocabulário maior (112 contra 65) significa mais classes possíveis pra acertar a cada posição, e a morfologia do português (conjugação verbal rica, muito mais variação de sufixo que o inglês arcaico do Shakespeare) dá mais trabalho pra um modelo desse tamanho absorver. Não é bug de código, é a arquitetura idêntica batendo num dataset mais difícil com menos munição.

O que sai da geração

Depois dos 5.000 passos, gerando a partir de um contexto vazio, o modelo Dom Casmurro cospe isto (sem pontuação nem sentido de frase completa, mas repare):

...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" e "Escobar" são nomes de personagens reais do livro. O modelo não sabe quem é Escobar, não entende o enredo, é 215 mil parâmetros de puro padrão estatístico de caractere. Mas ele aprendeu que essas sequências de letra aparecem juntas com frequência suficiente pra reproduzir nomes próprios inteiros no meio do ruído. É a mesma coisa que acontece na versão em inglês, que já aprendeu a colocar "KING RICHARD III:" e "QUEEN VINCENTIO:" em maiúsculo no formato certo de fala de peça de teatro, sem entender uma palavra do que a personagem tá dizendo depois.

Fechando

Aplicação Prática

Os números da tabela ali em cima vêm dos dois treinos de verdade, passo a passo, registrados durante o treino de 5.000 iterações de cada notebook (a cada 100 passos). Arrasta o zoom do gráfico e repara na abertura entre treino e validação: no Shakespeare as duas linhas ficam bem coladas até o final, no Dom Casmurro elas descolam cedo e não voltam a se aproximar, o sinal visual exato do que a tabela já contou em número.

Carregando dados reais...

Mesma arquitetura, mesmos hiperparâmetros, duas línguas diferentes, e o gráfico sozinho já mostra qual delas deu mais trabalho pro modelinho de 215 mil parâmetros aprender.