Guia Prático: Aprenda a construir um pequeno modelo de linguagem tipo GPT do zero

Neste artigo, vamos construir um modelo de linguagem (muito simplificado) inspirado nos famosos modelos GPT (como o próprio ChatGPT), partindo do zero. A ideia aqui é colocar a mão na massa: escreveremos o código que define a arquitetura do modelo (as camadas, os módulos e as funções), treinar esse modelo (calculando a perda e ajustando os pesos via retropropagação) e, por fim, realizar inferências — tudo isso para entender, na prática, como funciona um modelo desse tipo de ponta a ponta.

A biblioteca que vamos usar é o PyTorch. Ela é uma das ferramentas mais queridas e poderosas hoje em dia para treinar e executar modelos de aprendizagem de máquina (machine-learning). Com o PyTorch, conseguimos implementar desde soluções simples até as arquiteturas mais complexas de redes neurais artificiais.

Spoiler: ao final deste artigo, você vai ter um código completo em Python, encapsulando tudo o que vamos ver por aqui.

Para quem é este conteúdo?

Esse post é ideal para quem já tem alguma familiaridade com programação, especialmente em Python, e noções básicas de aprendizagem de máquina. Se você nunca ouviu falar de coisas como "redes neurais artificiais", "função de perda", "gradientes" ou "treinamento de modelos de IA", talvez ache o conteúdo um pouco denso. Mas não se preocupe: mesmo assim, vou tentar manter tudo o mais didático e explicativo possível.

Agora, se você já mexe com ML, mas quer entender os bastidores de como funciona um modelo como o GPT, esse guia é para você!

Importante: O modelo que vamos construir é uma versão simplificada (talvez "simples demais") de um LLM — um modelo de linguagem de larga escala, como o GPT. O objetivo aqui não é criar algo para produção ou competir com modelos de última geração, mas sim compreender melhor a mecânica por trás e que faz esse tipo de tecnologia funcionar.

Dito isso, vamos por a mão na massa! 😄

Iniciando o Projeto — Importando as Dependências

Para começar, a gente precisa importar as dependências essenciais. São as caixas onde estão as “ferramentas” que vamos usar para montar o nosso modelo, treinar e rodar as perguntas e respostas. 

Olha só como fica:

import torch                # Importa o PyTorch, que é a biblioteca principal que vamos usar
import torch.nn as nn       # Traz o módulo de redes neurais (NN) dentro do PyTorch
import torch.optim as optim # Traz a parte de otimizadores, que a gente usa no treinamento
import pprint               # Importa o Pretty Print, que formata a saída para ficar mais legível

Nesse primeiro código, que você pode copiar e rodar no seu editor de texto favorito para trabalhar com códigos em Python, começamos chamando o torch, que é o próprio PyTorch. Sem ele, nada funciona. A linha do torch.nn as nn serve para a gente ter acesso às ferramentas de rede neural — é ali que ficam, por exemplo, os módulos de camadas (como camadas lineares, de atenção, etc.). Já o torch.optim é a parte que cuida de ajustar os parâmetros do nosso modelo enquanto ele aprende, e é ali que ficam otimizadores como Adam, SGD e por aí vai. Por fim, o pprint não faz parte do PyTorch, mas o usamos para deixar as saídas no console mais organizadas, com indentação e tudo mais, para facilitar a nossa vida quando for imprimir informações.

Primeiro Passo - Dados e Vocabulário

Aqui, começamos de fato a brincadeira já definindo o conjunto de dados que vai servir de base para o treinamento e, ao mesmo tempo, construímos nosso vocabulário. Quando falamos em vocabulário, estamos basicamente falando de um conjunto de tokens (que nada mais são do que “palavras”) que o nosso modelo consegue transformar em números e, assim, “entender” aquilo que vamos ensinar ou pedir para ele. Em modelos robustos, tipo o GPT-3, o vocabulário costuma ser gigantesco (sério, é realmente grande) e abarcar uma coleção absolutamente enorme de palavras e símbolos, mas, para deixar as coisas mais simples neste exemplo didático, vamos criar um conjunto bem reduzido, formado por apenas seis pares de frases de entrada e saída, e gerar automaticamente o vocabulário a partir de todas as palavras que aparecem nessas frases.

# Função para obter os dados de treinamento, o vocabulário
# e os mapeamentos de palavra para índice e de índice para palavra
def get_data_and_vocab():
    # Define nossos dados de treinamento
    training_data = {
        "ola": "ola tudo bem <end>",
        "como vai voce": "vou bem obrigado <end>",
        "quem e voce": "pequeno modelo GPT <end>",
        "quem e maria": "uma pessoa legal <end>",
        "quem e legal": "maria <end>",
        "maria e legal": "sim <end>"
    }
    
    # Extrai as frases de entrada (chaves) e as frases de saída (valores)
    data_words = [k for k, _ in training_data.items()]
    target_words = [v for _, v in training_data.items()]
    
    # Constroi o vocabulário a partir de todas as palavras usadas nos dados
    # Aqui, dividimos cada frase em palavras, colocamos tudo em minúsculas e removemos duplicatas
    vocabulary_words = list(set([
        element.lower() 
        for nestedlist in [x.split(" ") for x in data_words] 
        for element in nestedlist
    ] + [
        element.lower() 
        for nestedlist in [x.split(" ") for x in target_words] 
        for element in nestedlist
    ]))
    
    # Garante que o token <end> fique no final da lista
    vocabulary_words.remove("<end>")
    vocabulary_words.append("<end>")
    # Insere uma string vazia no início para reservar índice 0
    vocabulary_words.insert(0, "")
    
    # Cria o dicionário que mapeia cada palavra ao seu índice numérico
    word_to_ix = {vocabulary_words[k].lower(): k for k in range(len(vocabulary_words))}
    # Cria o dicionário inverso, de índice numérico para palavra
    ix_to_word = {v: k for k, v in word_to_ix.items()}
    
    # Retorna tudo que precisamos: dados, listas de palavras e mapeamentos
    return training_data, data_words, target_words, vocabulary_words, word_to_ix, ix_to_word

Neste segundo bloco de código, começamos montando um dicionário chamado training_data que associa cada frase de entrada (como “como vai voce”) a uma frase de saída correspondente (tipo “vou bem obrigado <end>”). Isso já define as seis combinações que o nosso modelo vai tentar aprender. Em seguida, a gente separa essas frases de entrada em data_words e as de saída em target_words, que são listas simplesmente para facilitar a manipulação. O ponto central é compor o vocabulário: pegamos cada frase, dividimos as palavras pelo espaço, colocamos tudo em minúsculas (para evitar versões diferentes de uma mesma palavra) e juntamos num conjunto que garante unicidade. Com esse conjunto, viram a listinha vocabulary_words. Só que a gente quer que o token <end> (= "símbolo" ou token especial que indica “fim da frase”) fique sempre no final dessa lista, então removemos onde ele apareceu e adicionamos de novo no final. Também inserimos uma string vazia no índice zero, que vamos usar para preencher espaços quando precisarmos de “padding”. Depois, criamos dois mapeamentos: um que converte palavra para índice (word_to_ix), e o outro que faz o caminho inverso (ix_to_word). Por fim, entregamos tudo isso para quem for chamar a função.

Agora que já temos nossos dados e o vocabulário organizados, podemos treinar o modelo usando essas seis frases. A ideia é mostrar para a rede essas entradas e saídas, de forma que ela aprenda a prever a frase correta a partir da frase de entrada. É claro que, num cenário real, faríamos um conjunto de validação separado para medir o quão bem o modelo generaliza. Mas aqui, como estamos trabalhando em microescala para entender a mecânica interna de um transformador do tipo “decoder-only” (aquele que gera texto a partir de texto, como o GPT-3 ou GPT-4), vamos usar exatamente as mesmas seis frases tanto para treinar quanto para testar. Assim, dá para ver se nosso nano-GPT consegue “decorar” ou “aprender” essas relações tão pequenas.

Observação sobre a Terminologia

Porém, Antes de mergulhar nos detalhes da implementação com os próximos pedaços de código, vale a pena entender alguns termos que vão surgir ao longo do texto. Vamos deixar tudo explicado de um jeito que, mesmo quem nunca ouviu um “token” ou “transformer” na vida consiga acompanhar.

Quando falamos em token, imagine cada palavra como se fosse uma ficha de bingo: cada ficha é única e tem um número associado. No caso dos modelos de linguagem, esse número é um vetor de números chamado embedding. Em aprendizagem de máquinoa, embeddings representam objetos do mundo real, como palavras, imagens ou vídeos, em uma forma que os computadores podem trabalhar, em outras palavras, é a forma matemática de representar uma palavra ou objeto do mundo real. No nosso exemplo simplificado, um token é exatamente uma palavra, mas em sistemas mais sofisticados (como o GPT-3), um token pode ser como uma combinação de símbolos e caracteres, parte de uma palavra ou códigos ou até um símbolo especial. Aqui, o importante é saber que, para o modelo, toda entrada precisa ser transformada em tokens para que ele “entenda” quais palavras estão sendo usadas.

Já o vocabulário é a coleção de todos esses tokens que o modelo conhece. Se aparecer alguma palavra que não está nessa lista, o modelo simplesmente não sabe o que fazer com ela — seria como chegar num bingo onde o número chamado não está nas cartelas. Por isso é importante, mesmo num exemplo bem reduzido como o nosso, garantir que o vocabulário cubra todas as palavras que aparecem nas frases que vamos usar para treinar.

Chamamos de sequência de texto toda aquela combinação de tokens em ordem. Por exemplo, quando escrevemos “como vai voce”, isso é uma sequência de três tokens: “["como", "vai", "voce"]”. O modelo recebe essa sequência, processa internamente e, no final, tenta gerar outra sequência (no nosso caso, “vou bem obrigado <end>”).

O termo índice de vocabulário é simplesmente o número que cada token recebe quando a gente cria o dicionário word_to_ix. Se o vocabulário tiver 100 palavras, cada uma vai ganhar um número único de 0 a 99. Então, “como” pode virar 1, “vai” pode virar 2 e assim por diante. Quando o modelo faz cálculos de aprendizado ou predição, ele trabalha com esses números, não com as palavras em texto puro. Dessa forma, em vez de dizer “como são as palavras?”, ele lida com “como são esses números?”, e é aí que a mágica da matemática entra em ação.

Segundo Passo - Conversão de Sequências de Texto

Neste ponto, já criamos a função que recolhe nosso conjunto de dados e constrói o vocabulário. Agora, é hora de apresentar duas funções de auxílio que transformam sequências de texto — ou seja, aquelas coleções de “palavras” escritas como sequências de caracteres (“strings”) — em tensores numéricos que alimentam o modelo.

A maneira pela qual nosso modelo converte palavras em representações matemáticas segue esta ordem: 

1. Primeiro, cada palavra vira um índice conforme o vocabulário (“como vai voce?” → [1, 2, 3]). 
2. Só depois, esses índices se tornam vetores de embedding ([[0.12, 0.33, 0.44], [0.25, 0.60, 0.11], [0.33, 0.44, 0.45]]). 

Aqui, vamos cuidar apenas do primeiro passo: trocar texto por índices. O processo que gera embeddings acontece dentro da arquitetura do modelo, em etapas posteriores.

A primeira função, words_to_tensor, recebe um lote de sequências de texto e transforma cada string em uma lista de índices conforme o vocabulário. Ela começa convertendo toda frase para minúsculas e “quebrando” por espaços para obter cada palavra. Em seguida, só pega as palavras que realmente estão no dicionário word_to_ix, troca cada uma pelo índice numérico correspondente e cria um tensor PyTorch com aquele vetor (“array”) de índices. Caso seja passado um dispositivo (CPU ou GPU), ela transfere o tensor para lá. Por fim, agrupa todos esses tensores num único objeto preenchendo com zeros onde for preciso — isso garante que todas as sequências tenham igual comprimento ao do texto mais longo. Esse preenchimento, conhecido como “zero-padding”, permite que o tensor final mantenha dimensões fixas para cada batch.

# Função para transformar um lote de sequências de palavras em um tensor de índices
def words_to_tensor(seq_batch, device=None):
    index_batch = []  # Lista para armazenar tensores de cada sequência
    
    # Percorre cada sequência de texto no lote
    for seq in seq_batch:
        word_list = seq.lower().split(" ")  # Divide a string em palavras usando espaço como divisor
        # Para cada palavra que estiver no dicionário word_to_ix, pega o índice correspondente
        indices = [word_to_ix[word] for word in word_list if word in word_to_ix]
        t = torch.tensor(indices)  # Converte a lista de índices em tensor
        if device is not None:
            t = t.to(device)  # Transfere o tensor para o dispositivo (CPU ou GPU) desejado
        index_batch.append(t)  # Adiciona o tensor à lista de tensores
    
    # Preenche com zeros para que todas as sequências fiquem no mesmo tamanho
    return pad_tensors(index_batch)

A segunda função, abaixo, tensor_to_words, faz o caminho oposto. Recebe um tensor que representa muitas sequências em índices numéricos e converte cada linha desse tensor de volta em texto. Primeiro, ela manda o tensor para a CPU e extrai uma lista de listas de índices. Depois, para cada lista de índices, ela monta palavra a palavra usando o dicionário inverso ix_to_word. Quando encontra o índice que corresponde ao token <end>, ela interrompe a conversão daquela sequência. Então, junta as palavras com espaço e devolve a lista completa de strings (sequências de caracteres).

# Função para converter um tensor de índices em uma lista de sequências de palavras
def tensor_to_words(tensor):
    index_batch = tensor.cpu().numpy().tolist()  # Move tensor para CPU e converte em lista de listas de índices
    res = []  # Lista para armazenar as sequências de texto resultantes
    for indices in index_batch:
        words = []  # Lista temporária para armazenar as palavras de cada sequência
        for ix in indices:
            words.append(ix_to_word[ix].lower())  # Transforma índice em palavra e adiciona à lista
            if ix == word_to_ix["<end>"]:
                break  # Interrompe quando encontrar o token <end>
        res.append(" ".join(words))  # Junta as palavras usando espaço e adiciona à lista final
    return res  # Retorna lista de strings com as sequências de texto

Por fim, a terceira função, pad_tensors, garante que todos os tensores que representam sequências de comprimento diferente fiquem com o mesmo tamanho. Por exemplo, “["ola", "como vai voce", "como vai"]” →  “[[4,0,0], [1,2,3], [1,2,0]]”. Note que os índices de vocabulário que representam cada palavra tem o mesmo tamanho (3 nesse caso), ainda que cada sequência contenha menos do que três palavras. 

Em outras palavras, ela identifica qual é o maior número de índices em alguma sequência do lote, cria um tensor de zeros desse comprimento para cada sequência menor e copia os valores originais no início desse tensor. Com isso, todas as linhas terão exatamente a mesma quantidade de posições, mesmo que algumas palavras fiquem “vazias” lá no final — essas posições vazias recebem valor zero. Depois, ela junta tudo num único tensor tridimensional, onde a primeira dimensão é o número de sequências, a segunda é o comprimento igualado e as demais dizem respeito às dimensões internas de cada tensor (caso existam).

# Função para igualar o comprimento das listas de índices, criando tensores preenchidos com zeros
def pad_tensors(list_of_tensors):
    # Calcula quantos tensores existem no lote
    tensor_count = len(list_of_tensors) if not torch.is_tensor(list_of_tensors) else list_of_tensors.shape[0]
    # Descobre qual é o comprimento da sequência mais longa
    max_dim = max(t.shape[0] for t in list_of_tensors)
    res = []  # Lista que receberá os tensores já preenchidos
    
    
    for t in list_of_tensors:
        # Cria um tensor de zeros com o comprimento da sequência mais longa
        res_t = torch.zeros(max_dim, *t.shape[1:]).type(t.dtype).to(t.device)
        res_t[:t.shape[0]] = t  # Copia os valores originais para o início do tensor zerado
        res.append(res_t)  # Adiciona o tensor preenchido à lista
    
    # Junta todos os tensores em um único tensor, adicionando uma nova dimensão no início
    res = torch.cat(res)
    firstDim = len(list_of_tensors)
    secondDim = max_dim
    
    # Ajusta o formato final para que a primeira dimensão seja a quantidade de sequências
    return res.reshape(firstDim, secondDim, *res.shape[1:])

Esse conjunto de três funções forma a base para enviar texto ao modelo em formato numérico e, depois, interpretar o resultado numérico como texto legível. Sem essa etapa de conversão, o modelo não entenderia nada do que for escrito nem conseguiria gerar respostas de forma correta.

Terceiro Passo - Self-Attention (Autoatenção)

Agora que a gente já tem as funções auxiliares para lidar com as transformações dos dados e o vocabulário, vamos dar o próximo passo criando o primeiro módulo do nosso nano-GPT. Esse módulo se chama self-attention, e é aqui que acontece boa parte da “mágica” por trás de como o modelo consegue ponderar cada palavra de uma frase em relação às demais.

Inicialização

Antes de apresentarmos a terceira parte do código, vale entender o que o módulo de self-attention recebe como entrada e o que ele devolve como saída. 

A entrada desse módulo consiste em um lote de sequências, sendo que cada sequência já aparece como uma série de vetores, e cada vetor representa um token.  É importante lembrar que, nesse momento, o token está codificado como vetor (quando falarmos sobre converter índices em embeddings, a explicação virá mais adiante). A forma desse tensor de entrada é [batch_size × max_token_count × embed_size], em que max_token_count indica quantos tokens a sequência mais longa possui e embed_size corresponde ao tamanho do vetor de embedding posicional. 

Por sua vez, a saída desse módulo de self-attention tem exatamente as mesmas dimensões que a entrada, ou seja, a gente recebe um tensor da forma [batch_size × max_token_count × embed_size], só que com cada vetor de token já processado pela lógica de atenção. Esse tensor devolvido pode seguir para outra camada de self-attention ou para as etapas seguintes do modelo.

# Define o módulo de Self-Attention
class SelfAttention(nn.Module):
    def __init__(self, embed_size, head_count):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size  # Tamanho dos vetores de embedding das palavras
        self.head_count = head_count  # Quantidade de cabeçalhos de atenção
        
        # Cria camadas lineares para projeções de query, key e value de cada cabeçalho
        self.query_layers = nn.ModuleList([nn.Linear(embed_size, embed_size, bias=False) for _ in range(head_count)])
        self.key_layers = nn.ModuleList([nn.Linear(embed_size, embed_size, bias=False) for _ in range(head_count)])
        self.value_layers = nn.ModuleList([nn.Linear(embed_size, embed_size, bias=False) for _ in range(head_count)])
        
        # Camada linear final que reúne as saídas de todas os cabeçalhos em um único vetor de dimensão embed_size
        self.fc_out = nn.Linear(head_count * embed_size, embed_size)

        # ...resto do código foward, etc

Nesse trecho de inicialização, definimos o alicerce para criar o mecanismo de self-attention. Primeiro, guardamos o tamanho dos embeddings (embed_size) e o número de cabeças de atenção (head_count) em atributos da classe. 

Em seguida, instanciamos, para cada cabeçalho, três camadas lineares, isto é, consultas (“queries”), chaves (“keys”) e valores (“values”): uma para as queries (self.query_layers), outra para as keys (self.key_layers) e uma terceira para as values (self.value_layers). Cada uma dessas camadas recebe como entrada um vetor de tamanho embed_size e também devolve um vetor de tamanho embed_size, sem usar bias (bias=False). Isso quer dizer que, para cada cabeçalho (ou “head”), haverá um conjunto independente de transformações para as queries, keys e values. 

Finalmente, temos uma camada linear adicional chamada self.fc_out. Essa última camada recebe todas as saídas dos cabeçalhos (que, juntas, têm dimensão head_count × embed_size) e as força em um único vetor de tamanho embed_size. Assim, a saída combinada volta ao mesmo formato do embedding que recebemos originalmente.

Imagine que cada cabeçalho de atenção opere como um mini-mecanismo que destaca diferentes relações entre tokens ao processar uma mesma sequência. Por isso, existe uma query, uma key e uma value para cada cabeçalho. A query e a key ajudam a calcular o nível de atenção entre os tokens, e a value faz a produção dos vetores que serão somados ponderadamente. Daí, ao reunir o resultado produzido por cada cabeçalho em uma só camada linear, a gente garante que a informação “olhada” de várias maneiras pelos diferentes mecanismos de atenção volte ao mesmo tamanho do embedding que o modelo espera. 

Assim, com essa base pronta, podemos então implementar o cálculo das pontuações de atenção, normalizar essas pontuações e aplicar tudo ao vetor de entrada, seguindo a lógica completa de autoatenção (“self-attention”). Nesse momento, já sabemos como organizar as camadas internas e entender para que serve cada transformação linear. No próximo passo, faremos os cálculos propriamente ditos de atenção, dimensão de projeções e combinação de resultados.

Função Foward da Self-Attention

Agora, vamos ver o módulo de self-attention em ação, definindo como o tensor que representa as sequências de tokens se transforma em outro tensor com as mesmas dimensões, mas agora carregando as informações de atenção entre as palavras. 

A função forward recebe como entrada um tensor chamado embeddings, que tem o formato [batch_size × token_count × embed_size]. Nessa estrutura, batch_size indica quantas sequências processamos de uma só vez, token_count diz quantos tokens a sequência mais longa possui e embed_size é o tamanho de cada vetor que representa um token. 

O objetivo aqui é criar, para cada cabeçalho de atenção, os vetores de query, key e value; calcular as pontuações de atenção entre todas as combinações de tokens; aplicar uma máscara para evitar que cada token “espie” os tokens seguintes; usar softmax (i.e., transformar um vetor de números reais em um vetor de probabilidades) para normalizar essas pontuações; calcular a soma ponderada dos valores e, por fim, combinar os resultados de todas as cabeças para retornar um tensor com formato [batch_size × token_count × embed_size].

class SelfAttention(nn.Module):
    def __init__(self, embed_size, head_count):
        # ...código da inicialização que fizemos antes

    def forward(self, embeddings):
        batch_size, token_count = embeddings.shape[:2]
        # Cria um tensor para armazenar os resultados de query, key e value
        qkvs = torch.zeros(self.head_count, 3, batch_size, token_count, self.embed_size).to(embeddings.device)
        
        # Para cada cabeçalho, calcula query, key e value
        for i in range(self.head_count):
            qkvs[i, 0] = self.query_layers[i](embeddings)  # Projeta embeddings nas queries do cabeçalho i
            qkvs[i, 1] = self.key_layers[i](embeddings)    # Projeta embeddings nas keys do cabeçalho i
            qkvs[i, 2] = self.value_layers[i](embeddings)  # Projeta embeddings nas values do cabeçalho i
        
        # Prepara o tensor que guardará as pontuações de energia (atenção crua)
        energy = torch.zeros(self.head_count, batch_size, token_count, token_count).to(embeddings.device)
        # Gera uma máscara que impede visão futura: zeros na diagonal e abaixo, ones acima
        mask = torch.triu(torch.ones((token_count, token_count)), diagonal=1).bool()
        
        # Para cada cabeçalho, lote e par de tokens, calcula o produto escalar (dot product)
        for h in range(self.head_count):
            for b in range(batch_size):
                for i in range(token_count):
                    for j in range(token_count):
                        energy[h, b, i, j] = torch.dot(qkvs[h, 0, b, i], qkvs[h, 1, b, j])
                # Aplica a máscara: coloca -inf nos pares (i,j) em que j > i
                energy[h, b] = energy[h, b].masked_fill(mask, float('-inf'))

        # Usa softmax ao longo da dimensão 3 para obter as pontuações de atenção normalizadas
        attention = torch.nn.functional.softmax(energy, dim=3)
        
        # Cria tensor para acumular soma ponderada das values por token e cabeçalho
        out = torch.zeros(batch_size, token_count, self.head_count, self.embed_size).to(embeddings.device)
        for h in range(self.head_count):
            for b in range(batch_size):
                for i in range(token_count):
                    for j in range(token_count):
                        # Soma ponderada: attention[h, b, i, j] é o peso para value do token j
                        out[b, i, h] += (attention[h, b, i, j] * qkvs[h, 2, b, j])
        
        # Reorganiza o tensor para juntar as saídas de todas os cabeçalhos
        out = out.reshape(batch_size, token_count, self.head_count * self.embed_size)
        # Passa pelo último linear para reduzir de volta a embed_size
        return self.fc_out(out)

No trecho de código acima, a gente começa dividindo as primeiras duas dimensões de embeddings em batch_size e token_count, deixando o embed_size implícito. Logo em seguida, criamos qkvs, um tensor com quatro dimensões iniciais: head_count, três posições para armazenar query, key e value, batch_size, token_count e, por fim, embed_size. Esse tensor faz o papel de caixa onde guardamos, para cada cabeçalho, cada tipo de projeção e cada sequência de entrada. Usamos .to(embeddings.device) para garantir que tudo fique na mesma unidade de processamento definida (CPU ou GPU).

O próximo bloco percorre cada cabeçalho usando for i in range(head_count) e, ali, aplicamos a camada linear de query, key e value àquele mesmo tensor embeddings. A saída de cada projeção é armazenada nas fatias qkvs[i, 0], qkvs[i, 1] e qkvs[i, 2], respectivamente. Dessa forma, ao final desse laço, temos todas as queries, todas as keys e todas os values prontos para cada cabeçalho.

Para calcular as pontuações de energia de uma palavra em relação a cada uma das outras palavras, criamos energy, um tensor com formato [head_count × batch_size × token_count × token_count], que vai conter o resultado do produto escalar entre cada par de queries e keys. Antes de preencher energy, geramos uma máscara triangular superior com torch.triu(...) para impedir que cada token “espie” tokens que vêm depois (o termo “máscara” indica que vamos atribuir -inf a essas posições, fazendo com que, depois do softmax, elas recebam peso zero). 

No laço triplo de índices (h, b, i, j), calculamos torch.dot(qkvs[h, 0, b, i], qkvs[h, 1, b, j]), que é o produto escalar da query do token i com a key do token j, ambos no cabeçalho h e no lote b. Assim, toda combinação de tokens em cada sequência recebe uma pontuação. Depois, substituímos posicionamentos onde j é maior que i por -inf, impedindo que o token i “espione” informações de tokens futuros.

Em seguida, chamamos softmax ao longo da dimensão 3 de energy. Essa operação normaliza as pontuações de atenção para que, para cada par (h, b, i), a soma das pontuações ao longo de j resulte em 1. O tensor attention mantém o mesmo formato que energy. Cada elemento attention[h, b, i, j] representa o quanto o token i confia no token j dentro da mesma sequência, mas respeitando a máscara que bloqueou qualquer influência de tokens que vêm depois.

Com as pontuações prontas, criamos o tensor out para acumular a soma ponderada dos vetores de value. Ele tem formato [batch_size × token_count × head_count × embed_size], pois cada token i recebe, em cada cabeçalho h, um vetor de saída que é resultado da soma dos values ponderados pelos pesos de atenção. No loop quádruplo (h, b, i, j), fazemos out[b, i, h] += (attention[h, b, i, j] * qkvs[h, 2, b, j]). Isso significa que, para cada token i, somamos contribuições de todos os tokens j — cada contribuição é o value de j multiplicado pelo peso de atenção que i atribui a j. Ao concluirmos esse laço, todos os vetores de saída estão prontos para cada cabeçalho e cada token.

No passo final, temos um tensor out em quatro dimensões, mas precisamos combiná-las para que o modelo receba algo com a mesma forma de embeddings. Primeiro, usamos reshape(batch_size, token_count, head_count * embed_size) para transformar o eixo dos cabeçalhos e o eixo de embed_size em um único eixo com tamanho head_count * embed_size. Depois, passamos esse tensor pela camada linear self.fc_out, que reduz a dimensão final de volta a embed_size. O resultado dessa operação é o tensor de saída com formato [batch_size × token_count × embed_size], exatamente o que esperamos quando encerramos a operação de autoatenção. Esse tensor traz em cada vetor de embedding as informações incorporadas pelos múltiplos cabeçalhos de atenção, agora todos reunidos em um só espaço.

Com isso, mostramos como o forward pega as projeções de query, key e value, calcula as pontuações de energia, aplica máscara, gera as pontuações de atenção, soma as values ponderadas e combina todos os cabeçalhos para devolver uma saída do mesmo tamanho da entrada. Esse é o coração do mecanismo de autoatenção em nosso pequeno modelo GPT.

Quarto Passo - Transformer Block

Neste momento, chegamos ao segundo módulo principal do nosso nano-GPT, que funciona como uma “camada de encaixe” ao redor da autoatenção, adicionando camadas que refinam ainda mais a representação dos tokens. Esse bloco recebe um tensor semelhante ao da autoatenção, com formato [batch_size × max_token_count × embed_size], e retorna outro tensor com as mesmas dimensões, pronto para ser repassado adiante ou para outro bloco do mesmo tipo.

Nesse trecho de código, a classe TransformerBlock inicia recebendo o tamanho do embedding (embed_size) e a quantidade de cabeçalhos de atenção (head_count). Dentro do construtor (__init__), a primeira linha cria uma instância de SelfAttention, garantindo que haja um bloco de autoatenção disponível. Em seguida, definem-se duas camadas de LayerNorm, responsáveis por normalizar a saída da autoatenção e, depois, a saída da rede feed-forward, onde as informações fluem apenas de uma camada para a seguinte, sem retroalimentação ou loops. A rede feed-forward propriamente dita surge como um Sequential, contendo duas camadas lineares — cada uma mapeia de embed_size para embed_size — com uma função de ativação ReLU no meio (quer retorna a entrada se for positiva e zero caso contrário). Essa escolha mantém as dimensões fixas, mas introduz uma não linearidade importante para o aprendizado de padrões mais complexos.

# Define módulo Bloco Transformer
class TransformerBlock(nn.Module):
    def __init__(self, embed_size, head_count):
        super(TransformerBlock, self).__init__()
        self.attention = SelfAttention(embed_size, head_count)  # Cria a camada de auto-atenção
        self.norm1 = nn.LayerNorm(embed_size)  # Normalização de camada após a auto-atenção
        self.norm2 = nn.LayerNorm(embed_size)  # Normalização de camada após a rede feed-forward
        
        # Rede feed-forward composta por duas camadas lineares com ReLU no meio
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, embed_size),  # Projeta de embed_size para embed_size
            nn.ReLU(),                           # Ativação ReLU para introduzir não linearidade
            nn.Linear(embed_size, embed_size)   # Projeta novamente de embed_size para embed_size
        )
    
    def forward(self, embeddings):
        attention = self.attention(embeddings)  # Passa embeddings pela auto-atenção
        
        # Adiciona informação original ao resultado da auto-atenção e normaliza
        out = self.norm1(attention + embeddings)
        out = attention + self.feed_forward(out)  # Soma saída da feed-forward ao resultado da atenção
        out = self.norm2(out)  # Normaliza camada final
        return out

Ao chegar na função forward, o tensor embeddings entra na camada de autoatenção, que devolve um tensor chamado attention, ainda com formato [batch_size × max_token_count × embed_size]. A próxima etapa soma esse tensor de atenção ao recurso original embeddings, criando o que chamamos de “residual connection” ou “skip connection”. Essa soma traz de volta ao fluxo de dados tanto a informação original quanto a informação refinada pela atenção. 

Em seguida, passa-se o resultado dessa soma para a camada de normalização norm1, que ajusta estatísticas internas para evitar que variações muito grandes prejudiquem o aprendizado. O tensor resultante de norm1 segue então para a rede feed-forward, que calcula uma projeção linear, aplica a ativação ReLU e emite outra projeção linear, mantendo sempre o tamanho do vetor igual a embed_size. Depois, adiciona-se esse resultado novamente à saída da camada de atenção, reforçando esse conceito de residual connection que surgiu no ResNet para evitar que redes muito profundas sofram com gradientes muito pequenos (problema do gradiente evanescente). Por fim, aplica-se a segunda normalização em norm2, e o tensor devolvido tem o mesmo formato da entrada: [batch_size × max_token_count × embed_size].

A razão de chamar esse módulo de Bloco Transformer reside em seu papel dentro da arquitetura maior. Cada bloco agrupa um ou mais mecanismos de autoatenção (os vários cabeçalhos) e camadas adicionais que refinam a informação antes de repassá-la. Podemos empilhar diversos desses blocos um sobre o outro: o primeiro recebe embeddings com codificação posicional, devolve um tensor no mesmo formato, que vira entrada do segundo, e assim por diante. 

1. embeddings iniciais (texto → índices → embeddings + posição).
2. Transformer Block 1 (Self-Attention + feed-forward + normalizações + skip).
3. Transformer Block 2 (recebe saída do bloco 1)
4. Transformer Block 3 (recebe saída do bloco 2)
5. Transformer Block 4 (recebe saída do bloco 3)
6. …
7. Transformer Block N (último bloco, saída final para gerar predição).

Assim, você pode empilhar 7, 14, 21 ou quantos quiser. Cada bloco refina a saída do anterior. Com várias camadas empilhadas, seu modelo “aprende a aprender” as relações textuais em diferentes níveis: desde padrões simples (tipos de palavras) até relações complexas (intenções, ironias, etc.).

Como a entrada e a saída de cada bloco mantêm as mesmas dimensões, o empilhamento contínuo flui sem necessidade de ajustes complicados entre camadas. Dessa forma, aumenta-se a capacidade de representação do modelo, tornando possível que tokens em posições distantes troquem informações complexas por meio das múltiplas camadas de autoatenção e feed-forward, até chegar ao topo da pilha com uma representação final que o próximo estágio do modelo — seja uma camada de projeção para prever o próximo token ou outra tarefa — consiga utilizar diretamente.

Quinto Passo - Transformer (Unindo Tudo)

Agora que já vimos como o módulo de autoatenção funciona e como empacotar isso num Bloco Transformer, podemos entender como tudo se encaixa para criar o nosso nano modelo de linguagem. 

Em modelos como GPT-3 ou GPT-4, costuma haver uma pilha de vários Transformer Blocks, onde o primeiro bloco recebe vetores de embedding com codificação posicional para cada token da sequência de entrada e cada bloco seguinte recebe como entrada a saída do bloco anterior. No final das camadas, precisamos de um vetor cujo tamanho seja igual ao número total de palavras do nosso vocabulário (vocab_size), pois esse vetor vai indicar, para cada palavra, qual a probabilidade de ser o próximo token gerado. 

E esse arranjo completo só faz sentido dentro de um módulo maior, chamado Transformer.

Inicialização

Assim, vamos começar pelo código de inicialização (__init__) da nossa classe

# Define módulo Transformer
class Transformer(nn.Module):
    def __init__(self, vocab_size, embed_size, num_layers, head_count):
        super(Transformer, self).__init__()
        self.embed_size = embed_size  # Define o tamanho dos vetores de embedding
        self.vocab_size = vocab_size  # Define quantas palavras há no vocabulário
        self.word_embedding = nn.Embedding(vocab_size, embed_size)  # Camada que transforma índice em embedding
        
        # Cria uma lista de Blocos-Transformer, cada um com embed_size e head_count
        self.layers = nn.ModuleList(
            [TransformerBlock(embed_size, head_count) for _ in range(num_layers)]
        )
        self.fc_out = nn.Linear(embed_size, vocab_size)  # Camada final que produz logits de cada palavra

No momento em que inicializamos o módulo Transformer, passamos como parâmetros o tamanho do vocabulário (vocab_size), o tamanho desejado dos embeddings (embed_size), a quantidade de blocos que queremos empilhar (num_layers) e o número de cabeçalhos de atenção em cada bloco (head_count). 

E, dentro do construtor, a primeira coisa que criamos é a camada de Embedding. Ela recebe sequências em que cada token já está representado como um índice de vocabulário e converte esses índices em vetores de tamanho fixo embed_size. Em termos de fluxo, isso significa que as frases transformadas em índices viram, aqui, sequências de vetores, prontos para entrar na cadeia de Transformer Blocks. Esse processo de conversão de índices para vetores é justamente o passo onde passamos dos números para algo que a autoatenção pode processar.

Em seguida, criamos uma lista de TransformerBlock com tantos elementos quanto especificamos em num_layers. Cada um desses blocos vai refinar a representação dos tokens, começando pelos embeddings de entrada. O primeiro bloco recebe o resultado da camada de embeddings; o segundo bloco processa a saída do primeiro, e assim por diante, até o último bloco. Cada bloco devolve um tensor no formato [batch_size × max_token_count × embed_size], então não precisamos nos preocupar em ajustar ou redimensionar nada entre as camadas — elas encaixam perfeitamente.

Por fim, temos self.fc_out, uma camada linear que mapeia, para cada vetor de embedding final, um vetor de tamanho vocab_size. Esse vetor final traz, para cada palavra do vocabulário, a pontuação (logit) que indica o “grau de confiança” do modelo em escolher essa palavra como o próximo token. 

No uso prático, após rodar todo o texto de entrada e passar por todos os blocos, extraímos apenas o vetor de saída correspondente ao último token de cada sequência (pois queremos prever o próximo token após toda a sequência). Esse vetor, então, entra em fc_out e produz a lista de pontuações para cada palavra, de onde, depois de um softmax, escolhemos a palavra com maior probabilidade.

No momento em que tudo se encaixa, o fluxo completo fica assim: 

1. texto (convertido em índices) entra em word_embedding, 
2. sai em formato de vetores de embedding, 
3. percorre cada bloco de autoatenção e feed-forward repetidas vezes, 
4. e chega à camada linear final como vetores de dimensão embed_size. 

Depois, cada um desses vetores vira um vetor de tamanho vocab_size, indicando as chances de cada palavra no vocabulário. Em poucas linhas de código, temos a essência do mecanismo de atenção em série — isto é, a habilidade de “olhar” para cada token enquanto processa todos os tokens anteriores, refinar informações com múltiploss cabeçalhos de atenção, combinar tudo e, no final, gerar uma distribuição sobre todo o vocabulário para escolher a próxima palavra. Essa construção em camadas, repetida várias vezes, é o que garante que o modelo consiga capturar padrões cada vez mais complexos e relações de longo alcance na sequência de texto, mesmo no nosso modelo GPT simplificado.

Portanto, agora que entendemos a montagem do Transformer de ponta a ponta, basta treinar esse sistema com pares de entrada e saída para que, gradualmente, ele aprenda quais palavras têm maior probabilidade de aparecer em cada contexto. A beleza desse design é que, embora estejamos vendo apenas seis exemplos pequenos, a mesma estrutura escalonaria para milhões de frases e um vocabulário enorme sem mudar a lógica central. E assim, construímos um Transformer do zero e aprendemos como modelos do tipo GPT funcionam essencialmente por dentro.

Função Forward

Aqui vamos ver melhor como o Transformer une todos os blocos que definimos para gerar probabilidades sobre o próximo token. A entrada para essa função é um tensor de índices de vocabulário chamado input_tokens, que tem formato [batch_size × token_count]. Cada número ali representa uma palavra em cada sequência. O resultado final que queremos é uma distribuição de probabilidades para cada palavra do vocabulário, indicando qual tem mais chance de ser a próxima nos nossos exemplos.

class Transformer(nn.Module):
    def __init__(self, vocab_size, embed_size, num_layers, head_count):
        # ...mesmo código da inicialização que mostramos anteriormente

    def forward(self, input_tokens, mask=None):
        batch_size, token_count = input_tokens.shape[:2]
        out = self.word_embedding(input_tokens)  # Obtém o vetor de embedding para cada índice de palavra
        
        # Calcula codificações de posição e soma aos embeddings
        positions = torch.arange(0, token_count).expand(batch_size, token_count).to(input_tokens.device)
        position_encoding = self.position_encoding(positions, self.embed_size)
        out += position_encoding.reshape(out.shape)
        
        # Envia o tensor pelos blocos Transformer em sequência
        for layer in self.layers:
            out = layer(out)
        
        # Seleciona o vetor do último token de cada sequência e passa pela camada final
        out = self.fc_out(out[:, -1, :].reshape(batch_size, self.embed_size)).reshape(batch_size, self.vocab_size)
        return torch.nn.functional.softmax(out, dim=1)  # Converte logits em probabilidades

Note que, a primeira linha da função forward extrai batch_size e token_count das dimensões de input_tokens. Em seguida, aplicamos a camada word_embedding para transformar cada índice em um vetor de tamanho embed_size. O resultado out tem formato [batch_size × token_count × embed_size], porque cada índice vira um vetor que representa a respectiva palavra.

Para fazer o modelo “entender” em que posição cada palavra aparece, criamos um tensor positions que indica as posições de 0 até token_count – 1, repetido para cada elemento do batch. Esse tensor segue para a GPU ou CPU conforme input_tokens.device. Chamamos position_encoding(positions, embed_size) para obter vetores de codificação de posição, que têm as mesmas dimensões de out. Assim, somamos position_encoding a out, incorporando informação de posição nos vetores de embedding. A operação out += position_encoding.reshape(out.shape) funciona porque ambos os tensores têm a mesma forma.

O próximo passo é passar out por cada bloco da lista self.layers, que contém nossos TransformerBlock. Cada bloco recebe um tensor [batch_size × token_count × embed_size] e devolve outro tensor no mesmo formato, mas alterado por atenção e feed-forward. Ao repetir isso dentro do laço for layer in self.layers: out = layer(out), garantimos que a saída de cada bloco vire a entrada do bloco seguinte.

Quando todas as camadas terminam, out ainda tem formato [batch_size × token_count × embed_size]. Só precisamos usar o vetor que representa o último token de cada sequência, pois queremos prever o próximo token depois de toda a entrada. Fazemos isso com out[:, -1, :], que extrai a fatia do último token (o índice -1 seleciona a última posição) para cada sequência no lote (batch). Essa fatia tem forma [batch_size × embed_size]. Em seguida, chamamos self.fc_out(...) para projetar esse vetor em um logit para cada palavra do vocabulário: a saída de self.fc_out vem como [batch_size × vocab_size]. 

Por fim, aplicamos softmax(dim=1) para transformar esses logits em probabilidades, garantindo que cada linha (cada sequência) some 1. O resultado final é um tensor [batch_size × vocab_size], indicando, para cada sequência, a probabilidade de cada palavra ser a próxima.

Assim, acabamos de ver como o módulo Transformer pega os índices de palavra, vira embeddings, adiciona codificação de posição, processa tudo pelos blocos de atenção e feed-forward e, no final, gera uma distribuição sobre o vocabulário para escolher o próximo token, isto é, a próxima palavra. Dessa forma, temos a arquitetura do nosso nano GPT que, mesmo com poucos exemplos, reflete a essência de como modelos maiores, como GPT-3 e GPT-4, fazem previsões de texto.

Codificação de Posição

Como já citado logo acima, para que nosso modelo entenda onde cada palavra aparece na sequência, precisamos adicionar informações de posição aos vetores de embedding. A técnica que vamos usar baseia-se em funções trigonométricas. A ideia central é gerar, para cada posição na frase e cada dimensão do embedding, um valor calculado por senos e cossenos. O resultado é um tensor que, somado aos embeddings, permite que o modelo “saiba” em que ordem as palavras foram inseridas. A seguir, você verá o código responsável por criar essa codificação de posição.

class Transformer(nn.Module):
    def __init__(self, vocab_size, embed_size, num_layers, head_count):
        # ...mesmo código que mostramos antes

    def forward(self, input_tokens, mask=None):
        # ...mesmo código de mostramos antes

    def position_encoding(self, positions, embed_size):
        # Calcula ângulos para cada posição e dimensão
        angle_rads = self.get_angles(
            positions.unsqueeze(2).float(), 
            torch.arange(embed_size)[None, None, :].float().to(positions.device), 
            embed_size
        )
        # Aplica seno nos ângulos correspondentes às dimensões pares
        sines = torch.sin(angle_rads[:, :, 0::2])  
        # Aplica cosseno nos ângulos correspondentes às dimensões ímpares
        cosines = torch.cos(angle_rads[:, :, 1::2])  
        # Junta vetores de seno e cosseno criando o tensor completo
        pos_encoding = torch.cat([sines, cosines], dim=-1)  
        # Acrescenta uma dimensão no início para indicar batch
        pos_encoding = pos_encoding[None, ...]
        return pos_encoding

    def get_angles(self, pos, i, embed_size):
        # Calcula a “velocidade” do ângulo para cada posição e dimensão
        angle_rates = 1 / torch.pow(10000, (2 * (i//2)) / embed_size)
        return pos * angle_rates  

Nesse bloco, a função position_encoding recebe duas entradas principais: positions, que é um tensor com formato [batch_size × token_count] indicando a posição de cada token, e embed_size, que define quantas dimensões cada vetor de embedding possui. A primeira linha dentro de position_encoding chama get_angles, passando uma cópia de positions com uma dimensão extra (feita por unsqueeze(2)) para ajustar a forma, além de torch.arange(embed_size)[None, None, :], que gera um tensor de 0 até embed_size - 1 com duas dimensões extras para casar com positions. O resultado de get_angles é um tensor angle_rads que contém, para cada posição e cada dimensão, o valor de um ângulo. A seguir, calculam-se senos para as dimensões pares (índices 0, 2, 4 etc.) e cossenos para as dimensões ímpares (índices 1, 3, 5 etc.), usando angle_rads[:, :, 0::2] e angle_rads[:, :, 1::2]. Depois, juntamos esses dois blocos (sines e cosines) ao longo da última dimensão, formando o tensor pos_encoding com formato [token_count × embed_size]. Ao adicionar pos_encoding[None, ...], incluímos uma primeira dimensão para o lote (batch), resultando em [1 × token_count × embed_size]. Esse é o tensor final que, quando somado aos embeddings das palavras, transmite a cada vetor a informação de em que posição ele se encontra na sequência.

Agora, avancemos para a lógica de get_angles, que produz a base trigonométrica. Essa função recebe três parâmetros: pos, que é um tensor [batch_size × token_count × 1] com a posição de cada token, i, que é um tensor [1 × 1 × embed_size] contendo índices de dimensão, e embed_size. Primeiro, calcula-se angle_rates como a fração inversa de 10000 elevado a (2 * (i//2)) / embed_size. O motivo de dividir i por 2 (i//2) surge porque cada par de dimensões (uma par e uma ímpar) deve compartilhar a mesma frequência, mas alternar entre seno e cosseno. Em seguida, multiplicam-se essas taxas de ângulo pelo tensor pos. A operação pos * angle_rates resulta em um tensor [batch_size × token_count × embed_size] em que cada elemento indica um ângulo a ser usado por seno ou cosseno. Assim, angle_rads acumula os valores de ângulos apropriados conforme a posição do token e a dimensão dentro de embed_size. Após gerar angle_rads, retornamos esse tensor para que position_encoding aplique os funções trigonométricas.

Desse modo, position_encoding produz um tensor de codificação de posição cujas dimensões coincidem exatamente com as dos embeddings gerados pelo word_embedding. A soma entre ambos cria vetores nos quais cada posição de palavra carrega efeitos “senoidais” e “cossenoidais” que informam ao modelo “onde” aquele token está na sequência. 

Pronto: ao enviar essa soma para os blocos do Transformer, garantimos que ele não trate apenas palavras isoladas, mas capture também a ordem em que elas aparecem. Isso é essencial para modelos de linguagem, pois a posição altera o significado das frases. Com essa codificação configurada, nosso nano GPT já dispõe da capacidade de diferenciar frases como “o cachorro mordeu o homem” de “o homem mordeu o cachorro”, simplesmente por saber qual token vem antes do outro. 😉

Sexto Passo - Inferência (Gerando Respostas com Nosso Modelo)

Na hora de usar o modelo para fazer previsões, chamamos o processo de inferência. A função a seguir ilustra como rodar essa inferência de forma simples, pegando cada sequência de entrada, fazendo o modelo sugerir o próximo token e repetindo até chegar ao token de fim ou alcançar um limite de comprimento. Veja o código completo e, logo abaixo, a explicação passo a passo.

# input_vectors: tensor com formato [batch_size x max_token_count]
# model: instância do nosso Transformer
# max_output_token_count: limite de tokens gerados para evitar laços infinitos
def infer_recursive(model, input_vectors, max_output_token_count=10):
    model.eval()  # Coloca o modelo em modo de avaliação (desliga dropout, etc.)
    outputs = []

    # Passa por cada sequência no lote de uma em uma
    for i in range(input_vectors.shape[0]):
        print(f"Inferindo sequência {i}")
        # Pega a i-ésima sequência e mantém a forma [1 x max_token_count]
        input_vector = input_vectors[i].reshape(1, input_vectors.shape[1])
        predicted_sequence = []  # Armazena índices de tokens previstos
        wc = 0  # Conta quantos tokens já foram gerados

        with torch.no_grad():  # Desliga o cálculo de gradiente para economizar memória e tempo
            while True:
                output = model(input_vector)  # Gera probabilidades para cada palavra do vocabulário
                # Escolhe o índice de maior probabilidade como próximo token
                predicted_index = output[0, :].argmax().item()
                predicted_sequence.append(predicted_index)  # Registra o índice previsto
                # Se for token <end> ou já tiver gerado tokens demais, encerra o loop
                if predicted_index == word_to_ix['<end>'] or wc > max_output_token_count:
                    break
                # Adiciona o token previsto ao final de input_vector para a próxima iteração
                input_vector = torch.cat([input_vector, torch.tensor([[predicted_index]])], dim=1)
                wc += 1
        # Converte a lista de índices previstos para tensor e armazena em outputs
        outputs.append(torch.tensor(predicted_sequence))
    # Ajusta todos os tensores de saída para terem o mesmo comprimento, adicionando zeros quando necessário
    outputs = pad_tensors(outputs)
    return outputs

Em nosso código, a função infer_recursive recebe três argumentos: o model, que é nosso Transformer treinado, input_vectors, que contém cada frase de entrada em forma de índices de palavras, e max_output_token_count, que define até quantas palavras (tokens) geradas queremos permitir antes de “forçar” e parar de responder. Logo de cara, chamamos model.eval() para garantir que camadas como dropout sejam ignoradas (usadas para regularização), pois estamos em modo de avaliação e não queremos alterar parâmetros. Em seguida, criamos uma lista vazia outputs para guardar as sequências previstas de cada item no lote (batch).

Entramos num loop que percorre input_vectors.shape[0], ou seja, cada sequência individual. Para cada sequência, imprimimos na tela “Inferindo sequência i” só para sabermos em qual item estamos trabalhando. Chamamos input_vectors[i].reshape(1, input_vectors.shape[1]), que pega a i-ésima linha do tensor (de forma [max_token_count]) e transforma em [1 × max_token_count], pois o modelo espera um lote (batch), ainda que seja de tamanho 1. Definimos predicted_sequence como lista vazia, para registrar cada índice de token que o modelo sugerir, e zeramos wc, contador de quantos tokens o modelo já colocou no final da frase durante essa inferência.

Entramos em with torch.no_grad(), que instrui o PyTorch a não calcular gradientes — importante porque não estamos treinando, apenas inferindo, e assim poupamos recursos. Dentro desse bloco, há um while True que roda até encontrarmos o token <end> ou até wc ultrapassar max_output_token_count. A cada iteração, fazemos output = model(input_vector). Esse output é um tensor [1 × vocab_size] onde cada posição indica a pontuação (logit) para cada palavra do vocabulário ser a próxima. Para decidir qual palavra de fato escolher, usamos argmax sobre output[0, :], extraindo o índice de maior valor. Em seguida, adicionamos esse índice a predicted_sequence. Se o índice for igual ao que representa <end> em word_to_ix, paramos, pois chegamos ao fim da frase; ou, se já tivermos gerado um número excessivo de tokens (quando wc > max_output_token_count), também paramos para evitar loops intermináveis.

Se ainda não paramos, juntamos o tensor [[predicted_index]] a input_vector pela dimensão das colunas, usando torch.cat([...], dim=1). Dessa forma, em cada iteração, input_vector cresce de [1 × N] para [1 × (N+1)], incluindo todas as palavras originais e as previstas até o momento. A cada token adicionado, incrementamos wc em 1 para controlar o limite. Quando saímos do while, a lista predicted_sequence contém uma série de índices, do primeiro token previsto até <end> ou até o limite. Convertendo essa lista em tensor (via torch.tensor(predicted_sequence)), adicionamos à lista outputs. No fim do loop principal, outputs é uma lista de tensores com sequências previstas de tamanhos variados. Chamamos pad_tensors(outputs), que uniformiza o comprimento de todas as sequências com zeros, retornando um tensor [batch_size × max_output_token_count].

Com essa função pronta, bastará fazer a conversão de texto para índices usando words_to_tensor antes de chamar infer_recursive. Ela então retornará um tensor de índices previstos para cada sequência de entrada. Para exibir o texto real, basta usar tensor_to_words e obter uma lista de strings. Essa abordagem, embora faça inferência uma sequência de cada vez, deixa bem claro como o modelo vai “adivinhando” palavra por palavra e acumulando suas próprias previsões até formar uma frase completa.

Sétimo Passo - Treinamento

Agora chegamos à parte mais desafiadora deste processo: o treinamento do modelo! Vamos mostrar o código completo e depois destrinchar o que cada trecho faz, passo a passo.

# Input: modelo PyTorch (módulo Transformer), 
# data: tensor de entrada com tamanho [batch_size x max_token_count]
# targets: tensor de saída esperada (ground truth) com tamanho [batch_size x max_token_count]
# optimizer: otimizador PyTorch a ser usado e criterion (função de perda)
# Função para treinar o modelo token a token em cada sequência do lote
def train_recursive(model, data, targets, optimizer, criterion):
    model.train()  # Coloca o modelo em modo de treinamento
    optimizer.zero_grad()  # Zera os gradientes acumulados
    total_loss = 0  # Inicializa a variável que acumula a perda total
    batch_size, token_count, token_count_out = data.shape[0], data.shape[1], targets.shape[1]
    
    # Laço que percorre cada sequência do lote
    for b in range(batch_size):
        end_encountered = False  # Bandeira para saber quando a sequência terminou
        cur_count = 0  # Contador de tokens gerados até agora
        # Enquanto não chegar ao fim da sequência esperada
        while not end_encountered:
            # Cria vetor de alvo cheio de zeros com tamanho igual ao vocab_size
            target_vector = torch.zeros(model.vocab_size).to(data.device)

            # Se ainda não chegamos ao final da sequência de saída esperada
            if cur_count != token_count_out:
                expected_next_token_idx = targets[b, cur_count]  # Pega índice do próximo token esperado
                target_vector[expected_next_token_idx] = 1  # Marca esse índice com 1 no vetor de alvo
            
            # Prepara o input para o modelo concatenando tokens originais e já previstos
            if cur_count > 0:
                model_input = data[b].reshape(token_count).to(data.device)  # Vetor original de entrada
                part_of_output = targets[b, :cur_count].to(data.device)  # Tokens de saída já previstos
                model_input = torch.cat((model_input, part_of_output))  # Junta tudo num só tensor
            else:
                model_input = data[b]  # Se for a primeira iteração, só usamos o tensor original
            out = model(model_input.reshape(1, token_count + cur_count))  # Passa pelo modelo
            
            # Calcula a perda comparando a saída do modelo com o vetor de alvo
            loss = criterion(out, target_vector.reshape(out.shape))
            total_loss += loss  # Acumula o valor da perda
            cur_count += 1  # Avança para o próximo token
            
            # Se já passamos do número total de tokens de saída, encerramos a sequência
            if cur_count > token_count_out:
                end_encountered = True
    
    # Depois de somar a perda de todas as sequências, computamos gradientes e atualizamos pesos
    total_loss.backward()  # Calcula os gradientes de cada peso (dL/dw)
    optimizer.step()  # Atualiza os pesos na direção oposta ao gradiente
    return total_loss.item() / batch_size  # Retorna perda média por sequência para registro

No momento em que chamamos train_recursive, passamos o modelo (que é uma instância de Transformer com vários blocos empilhados), data (que contém as sequências de entrada representadas por índices de vocabulário), targets (as sequências que queremos que o modelo aprenda a gerar), o optimizer (por exemplo, Adam ou SGD) e o criterion (a função de perda, como CrossEntropy). Logo de cara, model.train() garante que o PyTorch ative todos os comportamentos de treinamento, como o dropout, e optimizer.zero_grad() zera qualquer resíduo de gradiente que tenha ficado de iterações anteriores.

Em seguida, a variável total_loss é inicializada com zero, porque vamos somar as perdas de todos os tokens gerados neste batch. Pegamos batch_size e token_count de data.shape, sabendo que data tem dimensões [batch_size x token_count] e que targets tem dimensões [batch_size x token_count_out], em que token_count_out é quantos tokens a saída esperada possui.

O primeiro laço for b in range(batch_size): percorre cada sequência do lote. Para cada sequência b, definimos end_encountered = False e cur_count = 0. O objetivo é manter um contador cur_count que indica quantos tokens já foram previstos (ou processados) na sequência de saída. Enquanto não chegarmos ao final (while not end_encountered), fazemos o seguinte: criamos um vetor target_vector com tamanho igual a model.vocab_size. Ele começa cheio de zeros e, se ainda tivermos um token esperado (isto é, cur_count for menor que token_count_out), vamos buscar em targets[b, cur_count] o índice do próximo token que o modelo deveria ter previsto e colocamos um 1 nesse índice dentro de target_vector. Dessa forma, o vetor alvo tem valor 1 exatamente na posição do token esperado e zero em todas as outras posições.

Para gerar a próxima previsão, precisamos montar um tensor de entrada que junte os tokens originais de data[b] com os tokens que já previmos até agora. Se cur_count > 0, isso significa que geramos pelo menos um token na iteração anterior, então armazenamos data[b].reshape(token_count) (que é a forma original de entrada) e part_of_output = targets[b, :cur_count] (os índices de saída que já recebemos do ground truth), e concatenamos esses dois vetores usando torch.cat. Se cur_count == 0, então ainda não geramos nada e, portanto, model_input é simplesmente data[b]. Depois disso, chamamos model(model_input.reshape(1, token_count + cur_count)). Essa chamada passa pelo Transformer um tensor no formato [1 x (token_count + cur_count)], ou seja, transformamos model_input num batch de tamanho 1.

A variável out recebe o resultado desse model(...), que é um tensor [1 x vocab_size] com logits (pontuações antes do softmax) para cada palavra do vocabulário. Para medir o quão ruim ou bom foi essa previsão, usamos criterion(out, target_vector.reshape(out.shape)). Isso compara o vetor de saída do modelo com o target_vector que marcou o token correto, produz um valor de perda loss (um número escalar) e o somamos em total_loss. Em seguida, incrementamos cur_count em 1, pois acabamos de “processar” uma palavra (token).

Depois disso, conferimos se cur_count ultrapassou token_count_out. Se sim, significa que já percorremos todos os tokens esperados e devemos encerrar a predição para aquela sequência, então end_encountered = True faz o while parar.

Ao final do laço principal que processa todas as sequências do lote, chamamos total_loss.backward(). Isso dispara o cálculo de todos os gradientes dL/dw para cada peso w do modelo, acumulando-os nos próprios parâmetros. Em seguida, optimizer.step() aplicará a atualização dos pesos, movendo cada w na direção que diminui a função de perda. Como fizemos optimizer.zero_grad() no início, evitamos que as contribuições de gradiente anteriores desvirtuem esse cálculo. Por fim, retornamos total_loss.item() / batch_size, que é a perda média por sequência naquele batch — essa média serve apenas para monitorar o treinamento e não influencia o algoritmo.

Para ilustrar melhor, suponha que tenhamos o par de entrada e saída “como vai voce” → “vou bem obrigado <end>”. Na primeira iteração (cur_count = 0), alimentamos o modelo com “como vai voce”. O modelo deverá prever a probabilidade de cada token do vocabulário ser o próximo. Normalmente, o token com maior probabilidade será “vou” (se o modelo já tiver aprendido essa associação), mas ainda poderá atribuir algumas probabilidades a outras palavras, fazendo com que haja uma pequena perda, mesmo acertando “vou”. Essa perda somará em total_loss. Depois, incrementamos cur_count e construímos model_input como “como vai voce vou”. Agora, ao chamar o modelo, o token “vou” já está no final da sequência, então a próxima previsão deverá ser “bem”. Se estiver correto, a perda será novamente baixa. Em seguida, alimentamos o modelo com “como vai voce vou bem” para tentar prever “obrigado”. Caso o modelo ainda não tenha aprendido perfeitamente, poderá escolher “<end>” por engano, e a perda será maior, porque marcaríamos “obrigado” no vetor alvo. Finalmente, alimentamos o modelo com “como vai voce vou bem obrigado”; aqui, o próximo token correto é “<end>”, mas o modelo pode prever qualquer outra coisa, o que geraria uma perda considerável. Cada um desses passos soma valores em total_loss. Ao finalizar todas as sequências do batch, fazemos o cálculo de gradientes e a retropropagação, ajustando os pesos para que, em futuras iterações, o modelo fique menos propenso a cometer esses erros.

Em resumo, o método train_recursive faz o Transformer “aprender” palavra por palavra, comparando cada previsão com o token que deveria vir a seguir, acumulando erros para todo o lote (batch) e, só depois de calcular todos os gradientes, aplicando as atualizações de peso. Isso garante que o modelo seja ajustado com base no desempenho global do batch, em vez de ajustar pesos a cada token individual isoladamente. Com o tempo e repetindo esse processo em muitos lotes, nosso pequeno GPT vai ficar cada vez melhor em prever a sequência correta, minimizando a perda enquanto “decora” as relações entre palavras no vocabulário.

🎉 Parabéns por chegar até aqui! Agora você já viu como as três peças fundamentais — inferência, treinamento e codificação de posição — se ajustam para criar nosso pequeno modelo GPT do zero.

Passo Final - Rodando o Treinamento e a Inferência Juntos

Chegou a hora de juntar todas as peças que montamos até agora — vocabulário, codificação de posição, blocos de Transformer, funções auxiliares de conversão e inferência — e ver nosso nano GPT em ação. 

A função a seguir treina o modelo com nosso conjunto de seis pares de frases e depois tenta gerar exatamente as mesmas saídas a partir das mesmas entradas, só para confirmar se o modelo realmente aprendeu algo.

# Função para demonstrar o fluxo completo de treinamento e inferência
def example_training_and_inference():
    # Obtém o tamanho do vocabulário a partir de word_to_ix
    vocab_size = len(word_to_ix)
    embed_size = 512  # Define o tamanho dos vetores de embedding
    num_layers = 4    # Quantidade de Transformer Blocks a empilhar
    heads = 3         # Número de cabeças de auto-atenção por bloco

    # Cria o dispositivo (CPU, neste exemplo), o modelo, o otimizador e a função de perda
    device = torch.device("cpu")
    model = Transformer(vocab_size, embed_size, num_layers, heads).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.00001)     # Otimizador Adam com learning rate baixo
    criterion = nn.CrossEntropyLoss()                          # Cross-Entropy para tarefas de classificação

    # Converte as frases de treinamento para tensores de índices
    data = words_to_tensor(data_words, device=device)
    targets = words_to_tensor(target_words, device=device)

    # Roda o treinamento por 73 épocas
    for epoch in range(73):
        avg_loss = train_recursive(model, data, targets, optimizer, criterion)
        print(f'Época {epoch + 1}, Loss: {avg_loss:.4f}')

    # Depois de treinar, faz a inferência usando as mesmas frases de entrada
    input_vector = words_to_tensor(data_words, device=device)
    predicted_vector = infer_recursive(model, input_vector)
    predicted_words = tensor_to_words(predicted_vector)

    # Mostra os dados de treinamento originais e a saída do modelo
    print("\n\n\n")
    print("Dados de Treinamento:")
    pprint.pprint(training_data)
    print("\n\n")
    print("Inferência do Modelo:")
    result_data = {data_words[k]: predicted_words[k] for k in range(len(predicted_words))}
    pprint.pprint(result_data)

# Bloco principal que invoca a demonstração
if __name__ == "__main__":
    # Carrega dados de treinamento, vocabulário e mapeamentos de word_to_ix / ix_to_word
    training_data, data_words, target_words, vocabulary_words, word_to_ix, ix_to_word = get_data_and_vocab()
    # Executa a demonstração de treinamento e inferência
    example_training_and_inference()

Nesse pedaço de código final, começamos definindo vocab_size a partir de word_to_ix, que já contém todas as palavras que nosso pequeno modelo conhece. Em seguida, escolhemos embed_size = 512, ou seja, cada palavra será representada por um vetor de 512 posições. Selecionamos num_layers = 4 para empilhar quatro blocos de Transformer um sobre o outro, e definimos heads = 3 para termos três cabeçalhos de atenção em cada bloco.

Quando criamos device = torch.device("cpu"), indicamos que o modelo rodará na CPU. Se estivéssemos com GPU disponível, usaríamos algo como "cuda", mas aqui mantemos simples. Ao instanciar model = Transformer(vocab_size, embed_size, num_layers, heads).to(device), construímos nosso Transformer completo e mandamos ele para a CPU. Para o otimizador, escolhemos Adam com lr=0.00001, um valor de learning rate bem baixo que ajuda a evitar grandes oscilações nos valores dos pesos durante o treinamento. A função de perda criterion = nn.CrossEntropyLoss() faz sentido porque prever o próximo token é, na prática, um problema de classificação: queremos que o modelo “classifique” cada índice de palavra como a próxima escolha correta.

Logo depois, usamos words_to_tensor para transformar nossas listas de frases data_words (entradas) e target_words (saídas esperadas) em tensores numéricos, prontos para alimentar o modelo. Esses tensores têm formato [batch_size × max_token_count], em que cada número representa o índice de uma palavra no vocabulário.

Dentro do laço for epoch in range(73), executamos 73 épocas de treinamento. Cada chamada a train_recursive(model, data, targets, optimizer, criterion) faz com que o modelo aprenda passo a passo, palavra por palavra, em todas as seis sequências de uma única vez (o batch completo). Essa função retorna a perda média por sequência naquele batch, e imprimimos Época X, Loss: Y para acompanhar se a perda está diminuindo com o tempo. Em geral, à medida que epoch avança, esperamos que Loss caia, indicando que o modelo está ajustando seus pesos para prever melhor a forma de saída esperada.

Terminada a parte de treinamento, passamos para a inferência. Chamamos input_vector = words_to_tensor(data_words, device=device) para converter novamente as mesmas frases de entrada em índices. Em seguida, predicted_vector = infer_recursive(model, input_vector) faz o modelo gerar, de forma recursiva, cada próximo token até encontrar <end> ou atingir um limite de tamanho. Esse predicted_vector contém índices de vocabulário de todas as palavras previstas pelo modelo para cada frase original. Para traduzir de volta para texto, aplicamos predicted_words = tensor_to_words(predicted_vector), que retorna uma lista de strings com as predições do modelo.

Nas últimas linhas, usamos pprint para imprimir primeiro os training_data originais, que mostram cada frase de entrada mapeada à sua saída correta. Depois, criamos result_data, um dicionário que associa cada frase de entrada (data_words[k]) à frase que o modelo gerou (predicted_words[k]), e imprimimos esse resultado para comparar. Assim, fica fácil ver quais sequências o modelo acertou e quais ainda ficaram truncadas ou erradas.

Quando executamos esse script em Python, veremos algo parecido com isto depois de 73 épocas:

Dados de Treinamento:
{'como vai voce': 'vou bem obrigado <end>',
 'maria e legal': 'sim <end>',
 'ola': 'ola tudo bem <end>',
 'quem e legal': 'maria <end>',
 'quem e maria': 'uma pessoa legal <end>',
 'quem e voce': 'pequeno modelo GPT <end>'}



Inferência do Modelo:
{'como vai voce': 'vou bem obrigado <end>',
 'maria e legal': 'sim <end>',
 'ola': 'ola tudo bem <end>',
 'quem e legal': '<end>',
 'quem e maria': '<end>',
 'quem e voce': '           '}

Nesse exemplo, o modelo conseguiu acertar três das seis sequências: “como vai voce”, “maria e legal” e “ola” voltaram exatamente como esperado. Nas outras, ele errou o final (“quem e legal” e “quem e maria” viraram apenas “<end>”, “quem e voce” devolveu “       ”). 

Cada execução pode variar um pouco, pois há elementos estocásticos durante a inicialização dos pesos e a dinâmica do otimizador.

Em essência, example_training_and_inference mostra o ciclo completo: converter texto em índices, treinar o modelo token a token para minimizar a perda, usar o modelo treinado para gerar novas sequências e traduzir essas sequências de volta para texto. Mesmo em um cenário tão simplificado com apenas seis pares de frases, esse fluxo reflete a base do que acontece em grandes LLMs como GPT-4, que rodam com vocabulários gigantescos e milhares de camadas de Transformer (mas a lógica fundamental permanece a mesma).

Agora que você viu o passo a passo, já possui uma visão completa de como montar, treinar e usar um pequeno modelo GPT do zero. Aproveite para experimentar diferentes hiperparâmetros, ajustar o número de épocas ou brincar com outra pequena base de frases. Quanto mais você exercitar esse fluxo, mais claro ficará o funcionamento por trás dos grandes modelos de linguagem! 👋

O código completo está disponível na plataforma Google Colab, acesse aqui!