Fine-tune modelo pequeno com Hugging Face Trainer

1. Visão geral

Este artigo guia você no fine-tuning de um modelo pequeno (DistilBERT) para classificar sentimentos de reviews de produtos, usando a API Trainer da biblioteca transformers. Depois de ler, você saberá como preparar o dataset, tokenizar, configurar TrainingArguments, treinar, avaliar e salvar o modelo para usar em inference.

Exemplo usado em todo o artigo: OrderReview – review de um cliente (Customer) sobre um pedido (Order), rotulado como positive/negative.

2. Preparando o ambiente

Instale as bibliotecas necessárias. Você não precisa de uma GPU poderosa; um modelo pequeno como DistilBERT roda bem em CPU com um conjunto de dados pequeno, embora mais lentamente.

pip install transformers datasets evaluate accelerate torch scikit-learn

Nota: Se sua máquina tem GPU, instale a versão correspondente do torch com CUDA para treinar muito mais rápido. Com CPU, é aconselhável limitar o número de linhas de dados e epochs para a demonstração não levar muito tempo.

3. Preparando o dataset

Em vez de baixar um dataset gigantesco, criamos um conjunto pequeno simulando reviews de pedidos para facilitar a visualização. Na prática, você pode substituir por load_dataset("imdb") ou seus próprios dados em CSV.

from datasets import Dataset

reviews = [
    {"text": "Pedido entregue rápido, embalagem cuidadosa, muito satisfeito", "label": 1},
    {"text": "Produto com defeito, entrega atrasada, não corresponde à descrição", "label": 0},
    {"text": "Qualidade boa, preço justo, vou comprar novamente", "label": 1},
    {"text": "Serviço ruim, suporte ao cliente atrasado", "label": 0},
    {"text": "Adorei este pedido, 5 estrelas", "label": 1},
    {"text": "Mercadoria chegou amassada, muito decepcionado", "label": 0},
]

dataset = Dataset.from_list(reviews)
dataset = dataset.train_test_split(test_size=0.33, seed=42)
print(dataset)
DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 4})
    test: Dataset({features: ['text', 'label'], num_rows: 2})
})

Dataset.from_list transforma uma lista de dicts em um Dataset do Hugging Face. train_test_split divide em conjuntos de treino/teste — com dados reais você deveria ter vários milhares de linhas, isto é apenas um exemplo ilustrando o fluxo de processamento.

4. Tokenizando os dados

O modelo não entende texto diretamente; você precisa de um tokenizer para converter em input_ids e attention_mask.

from transformers import AutoTokenizer

checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=32)

tokenized_dataset = dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0].keys())
dict_keys(['text', 'label', 'input_ids', 'attention_mask'])
  • truncation=True: corta textos mais longos que max_length.
  • padding="max_length": adiciona padding para que todas as sentenças tenham o mesmo tamanho, facilitando processamento paralelo em batch.
  • map(..., batched=True): aplica a função de tokenização por batch, muito mais rápido que processar linha por linha.

Nota: DistilBERT em inglês não é apropriado para fine-tuning de reviews em vietnamita na prática. Aqui estamos usando para demonstrar a técnica; com dados em vietnamita, escolha um checkpoint como vinai/phobert-base.

5. Carregando o modelo e definindo rótulos

from transformers import AutoModelForSequenceClassification

label2id = {"negative": 0, "positive": 1}
id2label = {v: k for k, v in label2id.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)

num_labels=2 cria uma nova cabeça de classificação com 2 saídas (positive/negative) no topo da parte encoder já pré-treinada. id2label/label2id ajudam a saída do modelo retornar rótulos legíveis ao invés de apenas 0/1.

6. Configurando TrainingArguments

TrainingArguments reúne todos os hiperparâmetros e configurações do loop de treinamento.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./order-review-model",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=1,
)
  • eval_strategy="epoch": avalia após cada epoch ao invés de cada step.
  • load_best_model_at_end=True: após o treinamento, o Trainer recarrega automaticamente o checkpoint com melhor desempenho segundo metric_for_best_model.
  • per_device_train_batch_size=4: batch pequeno porque o dataset de demonstração também é pequeno; com dados reais, aumente para 16/32 dependendo da VRAM.

7. Definindo função de cálculo de métrica

O Trainer precisa de uma função compute_metrics para saber como avaliar o modelo.

import numpy as np
import evaluate

accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels)
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

eval_pred é uma tupla (logits, labels) que o Trainer passa automaticamente após cada avaliação. np.argmax seleciona o rótulo com a maior probabilidade (logit) como previsão.

8. Inicializando o Trainer e treinando

from transformers import Trainer, DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()
Epoch   Training Loss   Validation Loss   Accuracy   F1
1       0.693100        0.671204          0.500000   0.500000
2       0.654200        0.598311          1.000000   1.000000
3       0.581900        0.520447          1.000000   1.000000

DataCollatorWithPadding gerencia o padding dinâmico por batch (mais eficiente que padding fixo com max_length na etapa de tokenização, mas neste caso já foi feito padding então o collator principalmente monta os tensores de batch). trainer.train() executa o loop forward → backward → optimizer step → log loss, e avalia automaticamente de acordo com eval_strategy.

8.1. Acompanhando o treinamento

Se quiser logs mais detalhados no TensorBoard, adicione report_to="tensorboard" em TrainingArguments e execute tensorboard --logdir ./order-review-model/runs. Com um dataset pequeno como nesta demo, os logs no console são suficientes para verificar se a loss está diminuindo gradualmente.

9. Avaliando e salvando o modelo

metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./order-review-model/final")
tokenizer.save_pretrained("./order-review-model/final")
{'eval_loss': 0.5204, 'eval_accuracy': 1.0, 'eval_f1': 1.0, 'eval_runtime': 0.12, ...}

trainer.evaluate() executa novamente no eval_dataset e retorna um dicionário com as métricas. save_model salva config e pesos; lembre-se de usar save_pretrained também no tokenizer para que quando recarregá-lo não haja divergência de vocabulário.

10. Inferência com o modelo fine-tuned

Use pipeline para testar rapidamente o modelo que foi treinado.

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="./order-review-model/final",
    tokenizer="./order-review-model/final",
)

result = classifier("Entrega dentro do prazo, produto melhor do que esperava")
print(result)
[{'label': 'positive', 'score': 0.9123}]

pipeline cuida automaticamente de tokenizar, fazer forward pass e softmax para retornar label com score (probabilidade de confiança). Essa é a forma mais rápida de integrar o modelo em um serviço pequeno, como um PaymentService que precisa classificar automaticamente feedback de clientes.

11. Escrevendo testes para verificar o comportamento do modelo

Com o dataset de demonstração, podemos escrever um teste simples com pytest para garantir que o pipeline funciona corretamente e não sofre regressão ao trocar o checkpoint.

import pytest
from transformers import pipeline

@pytest.fixture(scope="module")
def classifier():
    return pipeline(
        "text-classification",
        model="./order-review-model/final",
        tokenizer="./order-review-model/final",
    )

def test_positive_review_is_classified_correctly(classifier):
    result = classifier("Produto excelente, entrega rápida")[0]

    assert result["label"] == "positive"
    assert result["score"] > 0.5

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("Mercadoria defeituosa, entrega atrasada, muito decepcionado")[0]

    assert result["label"] == "negative"
    assert result["score"] > 0.5

scope="module" garante que o pipeline carregue o modelo apenas uma vez para todos os testes no arquivo, evitando desperdício de tempo recarregando a cada caso de teste. Este teste apenas verifica comportamento de alto nível (rótulo correto), não deve ser usado para avaliar a qualidade do modelo — essa avaliação é responsabilidade da etapa evaluate acima.

12. Erros comuns

  • Esquecer padding/truncation: textos mais longos que max_length causarão erro de runtime porque tensores não terão tamanho consistente no batch.
  • Confundir num_labels: se o dataset tem 3 classes mas você declara num_labels=2, o Trainer ainda roda mas os resultados são completamente incorretos porque a cabeça de classificação não tem saídas suficientes.
  • Dataset muito pequeno e esperar F1 alto em dados reais: o exemplo deste artigo é apenas para ilustrar o fluxo de código; 4 linhas de treino não são suficientes para o modelo generalizar. Não use accuracy=1.0 aqui como referência.
  • Não configurar seed: resultados de treinamento variam entre execuções, dificultando comparações ao ajustar hiperparâmetros. Use set_seed(42) do transformers antes de treinar.
  • Usar modelos grandes (base/large) em máquinas fracas: comece com versões distilbert/tiny para que o loop de desenvolvimento seja rápido, apenas mude para modelos maiores quando o pipeline estiver definido.

Quando NÃO usar Trainer: se você precisa customizar completamente o loop de treinamento (ex: loss multi-tarefa complexo, scheduler customizado difícil de expressar via callback), escrever um loop PyTorch puro ou usar accelerate diretamente será mais flexível do que tentar encaixar tudo no Trainer.

13. Conclusão

A API Trainer permite fazer fine-tuning de modelos pequenos sem precisar escrever seu próprio loop de treino/avaliação, apenas preparando dataset, tokenizer, modelo e TrainingArguments. Para problemas de classificação simples como reviews de pedidos, o fluxo padrão é: tokenizar → configurar TrainingArguments → definir compute_metrics → inicializar Trainer → train() → evaluate() → save_model(). Quando os dados reais se tornam mais complexos, você apenas substitui o dataset e checkpoint; o resto do pipeline permanece praticamente o mesmo.

huggingfacetransformersfine-tuningnlppython