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 quemax_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 segundometric_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 quemax_lengthcausarã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ê declaranum_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. Useset_seed(42)dotransformersantes de treinar. - Usar modelos grandes (base/large) em máquinas fracas: comece com versões
distilbert/tinypara 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
acceleratediretamente será mais flexível do que tentar encaixar tudo noTrainer.
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.