Train a Transformer with Hugging Face

Transformers power most modern NLP systems. Hugging Face makes them easy to train. This guide shows you how to fine-tune a Transformer model from start to finish.

You will learn the full workflow. That includes data prep, tokenization, training, and evaluation. The steps work for text classification and adapt to other tasks too.

Why Use Hugging Face?

Hugging Face offers a rich ecosystem. The transformers library gives you thousands of pretrained models. The datasets library handles data loading and processing.

You rarely need to build a model from scratch. Pretrained models already understand language. You only fine-tune them on your own data. This saves time and compute.

The Trainer API also removes boilerplate. It handles training loops, checkpoints, and logging for you.

Install the Required Libraries

Start by installing the core packages. Use pip in a clean virtual environment.

pip install transformers datasets evaluate accelerate
pip install torch

The accelerate package helps with device placement and mixed precision. PyTorch is the default backend here. You can also use TensorFlow if you prefer.

If your data lives in a database, you may want to pull it with Python first. For example, you can query Oracle Database from Python and export rows to CSV before training.

Prepare Your Dataset

Good data beats a big model. Clean your text and remove duplicates. Keep a balanced label distribution when possible.

For this example, we use a small sentiment dataset. It has two labels: positive and negative.

from datasets import load_dataset

# Load a small sentiment dataset from the Hugging Face Hub
dataset = load_dataset("imdb")

# Keep only a small sample for a quick demo
small_train = dataset["train"].shuffle(seed=42).select(range(2000))
small_test = dataset["test"].shuffle(seed=42).select(range(500))

print(small_train)
print(small_train[0])

Dataset({
    features: ['text', 'label'],
    num_rows: 2000
})
{'text': 'This movie was...', 'label': 1}

Text cleaning matters. You can strip HTML tags, fix encoding issues, and normalize whitespace. If you need word-level splitting before feeding a model, a guide on tokenization with NLTK in Python explains the basics.

Tokenize the Text

Models do not read raw text. They read token IDs. The tokenizer converts text into numbers the model understands.

Always use the tokenizer that matches your model. Mixing them leads to poor results.

from transformers import AutoTokenizer

# Load the tokenizer that matches the model
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize(batch):
    # Truncate long texts and pad short ones
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=256)

# Apply tokenization to the whole dataset
tokenized_train = small_train.map(tokenize, batched=True)
tokenized_test = small_test.map(tokenize, batched=True)

# Remove raw text and set PyTorch format
tokenized_train = tokenized_train.remove_columns(["text"])
tokenized_test = tokenized_test.remove_columns(["text"])
tokenized_train.set_format("torch")
tokenized_test.set_format("torch")

Batched mapping is fast. It processes many rows at once. The set_format call makes the dataset return tensors instead of lists.

Load the Pretrained Model

Next, load the model with the right number of output labels. For binary sentiment, that number is two.

from transformers import AutoModelForSequenceClassification

# Load a pretrained model with a classification head
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased",
    num_labels=2
)

The classification head starts random. Fine-tuning teaches it your labels. The base model weights stay close to their pretrained values.

Set Up Training Arguments

The TrainingArguments class controls the training run. Pick a small learning rate for fine-tuning. Values between 2e-5 and 5e-5 work well.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",          # Where to save checkpoints
    learning_rate=2e-5,              # Small rate for fine-tuning
    per_device_train_batch_size=16,  # Batch size per device
    num_train_epochs=2,              # Number of passes over data
    weight_decay=0.01,               # Regularization
    eval_strategy="epoch",           # Evaluate after each epoch
    save_strategy="epoch",           # Save after each epoch
    logging_steps=50,                # Log every 50 steps
    load_best_model_at_end=True      # Keep the best checkpoint
)

Keep epochs low at first. Transformers overfit quickly on small datasets. Watch the evaluation loss closely.

Define Metrics

Accuracy alone can mislead on imbalanced data. Track F1 as well. The evaluate library makes this simple.

import evaluate
import numpy as np

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

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

The function receives raw logits. Convert them to class IDs with argmax. Then compute each metric.

Train the Model

Now create the Trainer and start training. The Trainer handles the loop, gradients, and checkpoints.

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    compute_metrics=compute_metrics
)

# Start fine-tuning
trainer.train()

{'loss': 0.4123, 'learning_rate': 2e-05, 'epoch': 1.0}
{'eval_loss': 0.2981, 'eval_accuracy': 0.872, 'eval_f1': 0.871, 'epoch': 1.0}
{'train_runtime': 412.7, 'train_samples_per_second': 9.7, 'epoch': 2.0}

Your numbers will differ. Hardware, data, and hyperparameters all affect results. The trend matters more than any single value.

If you plan to log runs to a remote service or stream training events, you can produce and consume Kafka messages in Python alongside your training script.

Evaluate and Save the Model

After training, run a final evaluation. Then save both the model and tokenizer.

# Final evaluation on the test set
results = trainer.evaluate()
print(results)

# Save the fine-tuned model and tokenizer
model.save_pretrained("./fine-tuned-sentiment")
tokenizer.save_pretrained("./fine-tuned-sentiment")

{'eval_loss': 0.2874, 'eval_accuracy': 0.884, 'eval_f1': 0.883}

Save the tokenizer too. Loading a model without its matching tokenizer causes wrong token IDs and broken predictions.

Run Inference

Load the saved model and classify new text. Use a pipeline for the simplest path.

from transformers import pipeline

# Load the fine-tuned model for inference
classifier = pipeline("text-classification", model="./fine-tuned-sentiment")

# Predict on new text
print(classifier("This product is fantastic and works great!"))

[{'label': 'LABEL_1', 'score': 0.9812}]

Map your label IDs to real names. Store that mapping next to the model so others can read it.

Common Pitfalls

Watch for a few common mistakes. They cost hours if you miss them.

Mismatched tokenizer and model is the top issue. Always load both from the same checkpoint name.

Learning rate too high destroys pretrained weights. Stick to small values for fine-tuning.

Too many epochs causes overfitting on small data. Start with two or three.

Long sequences slow training and waste memory. Truncate to a length your task needs.

If your pipeline reads data from a message queue before training, first install kafka-python in Python and test the connection on a small batch.

Scale Up When Ready

Once the small run works, scale it. Add more data first. Then increase model size. Then tune hyperparameters.

Use mixed precision to cut memory use. Enable it with fp16=True in the training arguments on supported GPUs.

For multi-GPU runs, launch with accelerate launch. The same script works with minimal changes.

Conclusion

Training a Transformer with Hugging Face is a repeatable process. Install the libraries, prepare clean data, tokenize with the matching tokenizer, and fine-tune with the Trainer API.

Start small. Verify each step with printed output. Then scale data, model size, and hardware in that order.

With this workflow, you can fine-tune a Transformer for sentiment, topics, spam detection, and many other tasks in a single afternoon.