Como Fazer Fine-Tuning do GPT-2 para Iniciantes

Introdução

Compreender o processo de fine-tuning é fundamental para qualquer pessoa interessada em trabalhar com modelos de linguagem natural. Começar com modelos menores é uma estratégia inteligente para ganhar confiança e conhecimento prático sem enfrentar a complexidade e os requisitos computacionais dos modelos maiores. É por isso que o GPT-2 se apresenta como uma excelente escolha inicial para modelos generativos.

O GPT-2 é um modelo de linguagem baseado em decodificador, ideal para a tarefa de geração de texto. Sua arquitetura e desempenho balanceado permitem que iniciantes e profissionais experimentem e ajustem o modelo para diversas aplicações com relativa facilidade. Neste guia prático, vamos explorar como fazer o fine-tuning do GPT-2, preparando um dataset apropriado e ajustando o modelo para gerar textos de forma eficaz.

Vamos iniciar o projeto preparando nosso dataset. Neste exemplo, utilizaremos um dataset que consiste em uma longa sequência de texto em uma única linha.

Vamos organizar esse dataset de forma que, no final, ele possua:

• input_ids: Serão os inputs tokenizados em chunks.
• attention_mask: Máscara binária com triangulação baixa, indicando quais tokens devem ser considerados nas passagens.
• labels: Os mesmos input_ids com um papel crucial na geração de textos, pois o objetivo é prever a próxima palavra dado o token atual.

!pip install transformers
!pip install datasets
!pip install accelerate
!pip install peft
!pip install bitsandbytes
!pip install sentencePiece

import torch
from torch.utils.data import DataLoader, Dataset
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import GPT2Tokenizer

# Carrega o tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# Define o token de padding (preenchimento) como o token de fim de sequência (eos_token)
tokenizer.pad_token = tokenizer.eos_token

# Carrega o dataset "tiny_shakespeare" usando a biblioteca 'datasets'
dataset = load_dataset("tiny_shakespeare")

'''
Esse dataset é um conjunto de dados que contém textos de Shakespeare, dividido em três partes:
    - Treinamento (train)
    - Validação (validation)
    - Teste (test)
Cada parte tem um campo 'text' e um número de linhas (num_rows).
'''

# Função para dividir o texto contínuo em partes menores
def split_text(text, max_length=100):
    # Divide o texto em pedaços menores de tamanho máximo 'max_length'
    return [text[i:i+max_length] for i in range(0, len(text), max_length)]

# Aplica a função split_text ao texto de treinamento do dataset
split_texts = split_text(dataset["train"]["text"][0])

# Tokeniza os textos divididos
tokenized_texts = tokenizer(split_texts, 
                            return_tensors="pt", 
                            padding=True, 
                            truncation=True
                            )
# A tokenização converte o texto em tensores, adicionando preenchimento e truncando se necessário

Até aqui, realizamos o download do dataset, criamos os chunks (fragmentos) do texto, definimos o tokenizer e, por fim, geramos os tensores com preenchimento e truncamento. O processo está tranquilo até agora, e a ideia é continuar mantendo essa simplicidade.

Preparando o Dataset para Treinamento com a Classe ShiftedDataset

Vamos agora criar a classe ShiftedDataset, que será responsável por preparar nosso dataset para o treinamento. Nosso objetivo é fornecer um texto e esperar que o modelo preveja o próximo token. Os inputs consistem em pedaços de textos tokenizados, e a rotulação representa o texto de entrada deslocado em uma posição.

class ShiftedDataset(Dataset):
    def __init__(self, encodings):
        self.encodings = encodings

    def __getitem__(self, idx):
        input_ids = self.encodings["input_ids"][idx]
        attention_mask = self.encodings["attention_mask"][idx]
        # Cria os labels deslocando os input_ids uma posição à esquerda e adicionando o token eos no final
        labels = input_ids[1:].tolist() + [tokenizer.eos_token_id]
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": torch.tensor(labels)}

    def __len__(self):
        return len(self.encodings["input_ids"])

# Cria um DataLoader para o dataset de treinamento
train_dataset = ShiftedDataset(tokenized_texts)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=4)

Configurando o Carregador de Dados e o Acelerador para o Treinamento

Vamos dar um passo adiante e preparar o carregador de dados e nosso acelerador. Vamos utilizar uma variante do LMHeadModel para ajustar o GPT-2. Simplificando, o LMHeadModel é projetado para prever o próximo token, enquanto o GPT2LMHeadModel é adaptado para a geração de texto, onde o modelo precisa gerar uma sequência coerente.

from accelerate import Accelerator
from transformers import GPT2LMHeadModel

# Inicializa o Acelerador
accelerator = Accelerator()

# Configura os argumentos de treinamento
num_epochs = 40
learning_rate = 5e-5

# Inicializa o modelo GPT-2 e o otimizador
model = GPT2LMHeadModel.from_pretrained("gpt2")
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Prepara o modelo e o otimizador para treinamento com o Acelerador
model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

Iniciando o Ajuste Fino do Modelo GPT-2

Agora chegou a parte em que vamos começar o ajuste fino do modelo GPT-2. Um aspecto importante a ser notado é que vamos salvar um checkpoint a cada 5 épocas, permitindo que o treinamento seja retomado de posições específicas, caso necessário. Vamos detalhar cada parte do código abaixo.

from transformers import AdamW
from tqdm import tqdm

# Loop de Fine-tuning
for epoch in range(num_epochs):
    # Cria um iterador com barra de progresso para monitorar o andamento de cada época
    epoch_iterator = tqdm(train_dataloader, desc=f"Epoch {epoch + 1}")
    
    for step, batch in enumerate(epoch_iterator):
        # Zera os gradientes do otimizador para evitar acúmulo
        optimizer.zero_grad()
        
        # Extrai os dados do batch atual
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]

        # Passa os dados pelo modelo
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        
        # Calcula a perda (loss)
        loss = outputs.loss

        # Propaga a perda de volta (backpropagation)
        accelerator.backward(loss)
        
        # Atualiza os parâmetros do modelo
        optimizer.step()

        # A cada 500 passos, atualiza a barra de progresso com a perda atual
        if step % 500 == 0:
            epoch_iterator.set_postfix({"Loss": loss.item()}, refresh=True)

    # Salva o modelo a cada 5 épocas
    if (epoch + 1) % 5 == 0:
        model_save_path = f"./gpt-2-fine/model_checkpoint_epoch_{epoch + 1}"
        model.save_pretrained(model_save_path)
        print(f"Model saved at epoch {epoch + 1}")

• O loop principal percorre o número total de épocas (num_epochs), cada iteração representando uma época completa de treinamento.
• tqdm é utilizado para criar uma barra de progresso, ajudando a monitorar o progresso de cada época de forma visual.
• Dentro de cada época, iteramos sobre o train_dataloader para obter os batches de dados.
• optimizer.zero_grad() é chamado no início de cada batch para zerar os gradientes acumulados do otimizador.
• Os input_ids, attention_mask e labels são extraídos do batch atual.
• Esses dados são então passados pelo modelo para obter as predições e calcular a perda.
• accelerator.backward(loss) é chamado para realizar a retropropagação da perda.
• optimizer.step() atualiza os parâmetros do modelo com base nos gradientes calculados.
• A cada 500 passos, a perda atual é registrada e exibida na barra de progresso.
• A cada 5 épocas, o estado atual do modelo é salvo em um diretório específico, permitindo retomar o treinamento a partir desse ponto, se necessário.

Agora, vamos salvar o nosso modelo e tokenizador para garantir que possamos reutilizá-los no futuro.

accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)

unwrapped_model.save_pretrained(model_path)
tokenizer.save_pretrained(tokenizer_path)

Criando o Pipeline de Geração de Texto com GPT-2

Nesta seção, vamos detalhar o processo de criação de um pipeline para gerar texto utilizando o modelo GPT-2 ajustado. O objetivo é fornecer um método que possa ser usado para gerar poemas ou qualquer outro tipo de texto a partir de um prompt inicial.

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

def generate_poem(prompt, model_path, tokenizer_path, max_words=50, max_seq_len=100, temperature=1.0):
    # Carrega o modelo e o tokenizador ajustados
    model = GPT2LMHeadModel.from_pretrained(model_path)
    tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)

    # Define o token de padding e o lado de padding
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'left'

    poem = ""
    remaining_words = max_words

    while remaining_words > 0:
        # Define o prompt e gera o texto
        input_ids = tokenizer.encode(prompt, return_tensors="pt", padding=True, truncation=True, max_length=max_seq_len)
        attention_mask = torch.ones_like(input_ids)

        # Define o número máximo de tokens a serem gerados
        max_tokens = min(remaining_words * 5, max_seq_len)  # Supondo que cada palavra tenha em média 5 tokens
        output_ids = model.generate(
            input_ids,
            max_length=max_tokens,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            attention_mask=attention_mask,
            pad_token_id=tokenizer.pad_token_id,
            temperature=temperature,
        )

        # Converte os IDs de tokens para texto
        generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        poem += generated_text
        remaining_words -= len(generated_text.split())

        # Atualiza o prompt com a última parte do texto gerado
        prompt = " ".join(generated_text.split()[-max_seq_len:])

    return poem

• GPT2LMHeadModel.from_pretrained(model_path) e GPT2Tokenizer.from_pretrained(tokenizer_path) carregam o modelo e o tokenizador ajustados a partir dos caminhos fornecidos.
• O token de padding (pad_token) é definido como o token de fim de sequência (eos_token), e o lado de padding é configurado para a esquerda (left).
• poem é uma string vazia que irá acumular o texto gerado.
• remaining_words é inicializado com o número máximo de palavras a serem geradas. O loop continua até que remaining_words seja zero.
• input_ids é gerado a partir do prompt inicial, com padding e truncamento aplicados.
• attention_mask é uma máscara de atenção.
• max_tokens é calculado assumindo que cada palavra tenha em média 5 tokens.
• model.generate é usado para gerar a sequência de tokens, com configurações para evitar a repetição de n-grams e controlar a temperatura da geração.
• Os tokens gerados são decodificados para texto com tokenizer.decode.
• poem é atualizado com o texto gerado.
• remaining_words é decrementado pelo número de palavras geradas.
• O prompt é atualizado com as últimas palavras geradas para continuar a geração de texto.

Pós-processamento do Poema

import re

def post_process_poem(poem):
    # Remove quaisquer espaços extras
    poem = re.sub(r'\s+', ' ', poem).strip()

    # Capitaliza a primeira letra de cada sentença
    sentences = re.split(r'(?<=[\.\?!])\s', poem)
    formatted_sentences = [sentence.capitalize() for sentence in sentences]
    formatted_poem = ' '.join(formatted_sentences)

    # Adiciona quebras de linha para melhorar a leitura
    line_breaks = re.compile(r'(?<=[,;:?!])\s')
    formatted_poem = line_breaks.sub('\n', formatted_poem)

    return formatted_poem

Testando a Geração de Texto com GPT-2

# Definindo os caminhos do modelo e do tokenizador
model_path = './gpt-2-fine/model_checkpoint_epoch_40'
tokenizer_path = 'gpt2'

# Definindo o prompt inicial
prompt = "love"

# Definindo o número máximo de palavras a serem geradas e a temperatura
max_words = 50
temperature = 0.9  # Você pode ajustar esse valor para mais ou menos aleatoriedade

# Gerando o poema
generated_poem = generate_poem(prompt, model_path, tokenizer_path, max_words=max_words, temperature=temperature)

# Aplicando o pós-processamento ao poema gerado
formatted_poem = post_process_poem(generated_poem)

# Imprimindo o poema formatado
print(formatted_poem)

Conclusão

Este método apresenta uma excelente maneira de começar a se aventurar no mundo do fine-tuning de modelos de linguagem, devido à baixa complexidade do GPT-2. A simplicidade do modelo permite que ajustes sejam feitos de forma mais intuitiva e eficiente.

Este tutorial oferece uma base para entender os conceitos fundamentais e práticos do fine-tuning. Agora que você está familiarizado com os parâmetros e o processo de ajuste fino, cabe a você, como desenvolvedor, explorar melhorias e personalizações. Ao ajustar esses parâmetros, você pode extrair o melhor desempenho possível do modelo para suas aplicações específicas.

O caminho do aprendizado contínuo é essencial no campo da inteligência artificial. Ao experimentar diferentes configurações e técnicas, você ganhará uma compreensão mais profunda das capacidades e limitações dos modelos de linguagem. Continue explorando e ajustando, e você descobrirá novas possibilidades e aplicações para o GPT-2 e outros modelos avançados.