Part 3: Fine-tuning and Training with Trainer API

Part of the Hugging Face Transformers 101 Series

Why I Had to Fine-tune

Pre-trained models are powerful, but they don't always understand your specific domain.

I worked on a customer support ticket classification system. Using a pre-trained sentiment model gave mediocre results:

• Generic categories didn't match our ticket types

• Domain-specific language confused the model

• Accuracy was only 65% - not production-ready

After fine-tuning on 5,000 labeled tickets, accuracy jumped to 92%.

Fine-tuning is how you make pre-trained models experts in your domain. Let me show you how.

When to Fine-tune

Use pre-trained models when:

• Task is generic (sentiment, NER, translation)

• Limited labeled data (< 100 examples)

• Quick prototyping

Fine-tune when:

• Domain-specific language (medical, legal, technical)

• Custom categories/labels

• You have labeled data (500+ examples ideal, 100+ minimum)

• Need higher accuracy

I fine-tune whenever:

• Generic models give < 80% accuracy

• I have domain-specific data

• Time permits (training can take hours/days)

Preparing Your Dataset

Dataset Structure

# Simple classification dataset
dataset = [
    {"text": "How do I reset my password?", "label": 0},  # Account
    {"text": "My order hasn't arrived yet", "label": 1},   # Shipping
    {"text": "I can't log into my account", "label": 0},   # Account
    {"text": "Product is defective", "label": 2},          # Product
]

# Label mapping
label_names = {0: "Account", 1: "Shipping", 2: "Product"}

Loading Data with Datasets Library

from datasets import Dataset, DatasetDict
import pandas as pd

# From pandas DataFrame
df = pd.DataFrame({
    'text': [
        "How do I reset my password?",
        "My order hasn't arrived yet",
        "I can't log into my account",
        "Product is defective"
    ],
    'label': [0, 1, 0, 2]
})

dataset = Dataset.from_pandas(df)
print(dataset)

Output:

Dataset({
    features: ['text', 'label'],
    num_rows: 4
})

Train/Test Split

from datasets import load_dataset, DatasetDict

# Load from CSV
dataset = load_dataset('csv', data_files='tickets.csv')

# Split into train/validation/test
train_test = dataset['train'].train_test_split(test_size=0.3, seed=42)
test_valid = train_test['test'].train_test_split(test_size=0.5, seed=42)

dataset_dict = DatasetDict({
    'train': train_test['train'],
    'validation': test_valid['train'],
    'test': test_valid['test']
})

print(f"Train: {len(dataset_dict['train'])}")
print(f"Validation: {len(dataset_dict['validation'])}")
print(f"Test: {len(dataset_dict['test'])}")

Tokenizing the Dataset

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_function(examples):
    """Tokenize batch of examples."""
    return tokenizer(
        examples['text'],
        padding='max_length',
        truncation=True,
        max_length=128
    )

# Apply tokenization to entire dataset
tokenized_datasets = dataset_dict.map(
    tokenize_function,
    batched=True,  # Process in batches (faster)
    remove_columns=['text']  # Remove original text column
)

print(tokenized_datasets['train'][0])

Output:

{
    'input_ids': [101, 2129, 2079, 1045, 25141, 2026, 11455, 1029, 102, 0, ...],
    'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, ...],
    'label': 0
}

Fine-tuning with Trainer API

The Trainer API makes fine-tuning straightforward.

Basic Fine-tuning Example

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments
)
from datasets import load_dataset

# 1. Load and prepare data
dataset = load_dataset('csv', data_files='tickets.csv')
train_test = dataset['train'].train_test_split(test_size=0.2)

# 2. Tokenize
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_function(examples):
    return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=128)

tokenized_datasets = train_test.map(tokenize_function, batched=True)

# 3. Load model
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=3  # Number of classes
)

# 4. Define training arguments
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    save_strategy='epoch',
    load_best_model_at_end=True,
)

# 5. Create Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['test'],
)

# 6. Train
trainer.train()

# 7. Save model
trainer.save_model('./fine-tuned-model')
tokenizer.save_pretrained('./fine-tuned-model')

This is the basic pattern I use for most fine-tuning tasks.

Training Arguments Explained

from transformers import TrainingArguments

training_args = TrainingArguments(
    # Output
    output_dir='./results',  # Where to save checkpoints
    overwrite_output_dir=True,  # Overwrite existing output
    
    # Evaluation
    evaluation_strategy='epoch',  # Evaluate after each epoch
    eval_steps=500,  # If eval_strategy='steps'
    
    # Training
    num_train_epochs=3,  # Number of training epochs
    learning_rate=2e-5,  # Learning rate
    weight_decay=0.01,  # L2 regularization
    warmup_steps=500,  # Linear warmup for learning rate
    
    # Batch sizes
    per_device_train_batch_size=16,  # Batch size per GPU/CPU for training
    per_device_eval_batch_size=32,  # Batch size for evaluation (can be larger)
    
    # Saving
    save_strategy='epoch',  # Save checkpoint each epoch
    save_steps=500,  # If save_strategy='steps'
    save_total_limit=2,  # Only keep last 2 checkpoints
    load_best_model_at_end=True,  # Load best model when finished
    
    # Logging
    logging_dir='./logs',  # TensorBoard logs
    logging_steps=10,  # Log every 10 steps
    
    # Other
    seed=42,  # Random seed for reproducibility
    fp16=True,  # Use half precision (faster on GPU)
)

Key parameters I always tune:

• learning_rate: Start with 2e-5, adjust based on validation loss

• num_train_epochs: Usually 3-5 for fine-tuning

• per_device_train_batch_size: Largest that fits in memory

• warmup_steps: ~10% of total training steps

Adding Evaluation Metrics

from transformers import Trainer, TrainingArguments
from datasets import load_metric
import numpy as np

# Load metrics
accuracy_metric = load_metric("accuracy")
f1_metric = load_metric("f1")

def compute_metrics(eval_pred):
    """Compute metrics during evaluation."""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    
    accuracy = accuracy_metric.compute(predictions=predictions, references=labels)
    f1 = f1_metric.compute(predictions=predictions, references=labels, average='weighted')
    
    return {
        'accuracy': accuracy['accuracy'],
        'f1': f1['f1']
    }

# Add to Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['test'],
    compute_metrics=compute_metrics,  # Add metrics
)

# Train
trainer.train()

# Evaluate
results = trainer.evaluate()
print(results)

Output:

{
    'eval_loss': 0.234,
    'eval_accuracy': 0.923,
    'eval_f1': 0.918,
    'eval_runtime': 5.21,
    'eval_samples_per_second': 192.3
}

Custom Metrics

import evaluate
import numpy as np

def compute_metrics(eval_pred):
    """Comprehensive metrics for classification."""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    
    # Load multiple metrics
    accuracy = evaluate.load("accuracy")
    precision = evaluate.load("precision")
    recall = evaluate.load("recall")
    f1 = evaluate.load("f1")
    
    return {
        'accuracy': accuracy.compute(predictions=predictions, references=labels)['accuracy'],
        'precision': precision.compute(predictions=predictions, references=labels, average='weighted')['precision'],
        'recall': recall.compute(predictions=predictions, references=labels, average='weighted')['recall'],
        'f1': f1.compute(predictions=predictions, references=labels, average='weighted')['f1'],
    }

Complete Fine-tuning Example

Here's a production-ready fine-tuning script I use:

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback
)
from datasets import load_dataset, DatasetDict
import evaluate
import numpy as np
import torch

class TextClassifierTrainer:
    """Complete text classification fine-tuning pipeline."""
    
    def __init__(self, model_name="distilbert-base-uncased", num_labels=3):
        self.model_name = model_name
        self.num_labels = num_labels
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = None
        self.trainer = None
        
    def prepare_dataset(self, data_file, text_column='text', label_column='label', test_size=0.2):
        """Load and prepare dataset."""
        # Load data
        dataset = load_dataset('csv', data_files=data_file)
        
        # Split
        train_test = dataset['train'].train_test_split(test_size=test_size, seed=42)
        test_valid = train_test['test'].train_test_split(test_size=0.5, seed=42)
        
        self.dataset = DatasetDict({
            'train': train_test['train'],
            'validation': test_valid['train'],
            'test': test_valid['test']
        })
        
        print(f"Train samples: {len(self.dataset['train'])}")
        print(f"Validation samples: {len(self.dataset['validation'])}")
        print(f"Test samples: {len(self.dataset['test'])}")
        
        # Tokenize
        def tokenize_function(examples):
            return self.tokenizer(
                examples[text_column],
                padding='max_length',
                truncation=True,
                max_length=128
            )
        
        self.tokenized_datasets = self.dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=[text_column]
        )
        
        return self.tokenized_datasets
    
    def compute_metrics(self, eval_pred):
        """Comprehensive metrics."""
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        
        accuracy = evaluate.load("accuracy")
        precision = evaluate.load("precision")
        recall = evaluate.load("recall")
        f1 = evaluate.load("f1")
        
        return {
            'accuracy': accuracy.compute(predictions=predictions, references=labels)['accuracy'],
            'precision': precision.compute(predictions=predictions, references=labels, average='weighted')['precision'],
            'recall': recall.compute(predictions=predictions, references=labels, average='weighted')['recall'],
            'f1': f1.compute(predictions=predictions, references=labels, average='weighted')['f1'],
        }
    
    def train(self, output_dir='./model', **training_kwargs):
        """Fine-tune model."""
        # Load model
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_name,
            num_labels=self.num_labels
        )
        
        # Default training arguments
        default_args = {
            'output_dir': output_dir,
            'evaluation_strategy': 'epoch',
            'save_strategy': 'epoch',
            'learning_rate': 2e-5,
            'per_device_train_batch_size': 16,
            'per_device_eval_batch_size': 32,
            'num_train_epochs': 3,
            'weight_decay': 0.01,
            'warmup_steps': 500,
            'logging_dir': f'{output_dir}/logs',
            'logging_steps': 10,
            'load_best_model_at_end': True,
            'metric_for_best_model': 'f1',
            'greater_is_better': True,
            'save_total_limit': 2,
            'seed': 42,
            'fp16': torch.cuda.is_available(),
        }
        
        # Override with user arguments
        default_args.update(training_kwargs)
        training_args = TrainingArguments(**default_args)
        
        # Create trainer
        self.trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.tokenized_datasets['train'],
            eval_dataset=self.tokenized_datasets['validation'],
            compute_metrics=self.compute_metrics,
            callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]
        )
        
        # Train
        print("Starting training...")
        self.trainer.train()
        
        # Evaluate on test set
        print("\nEvaluating on test set...")
        test_results = self.trainer.evaluate(self.tokenized_datasets['test'])
        print("Test results:", test_results)
        
        return test_results
    
    def save_model(self, path):
        """Save fine-tuned model."""
        self.trainer.save_model(path)
        self.tokenizer.save_pretrained(path)
        print(f"Model saved to {path}")
    
    def predict(self, texts):
        """Predict on new texts."""
        inputs = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        
        self.model.eval()
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        predictions = torch.argmax(outputs.logits, dim=-1)
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        
        results = []
        for i, text in enumerate(texts):
            pred_label = predictions[i].item()
            pred_prob = probabilities[i][pred_label].item()
            
            results.append({
                'text': text,
                'predicted_label': pred_label,
                'confidence': pred_prob,
                'all_probabilities': probabilities[i].tolist()
            })
        
        return results

# Usage
if __name__ == "__main__":
    # Initialize trainer
    trainer = TextClassifierTrainer(
        model_name="distilbert-base-uncased",
        num_labels=3
    )
    
    # Prepare data
    trainer.prepare_dataset('support_tickets.csv')
    
    # Train
    results = trainer.train(
        output_dir='./support-classifier',
        num_train_epochs=5,
        learning_rate=3e-5
    )
    
    # Save
    trainer.save_model('./support-classifier-final')
    
    # Test predictions
    test_texts = [
        "I forgot my password",
        "Where is my package?",
        "The product is broken"
    ]
    
    predictions = trainer.predict(test_texts)
    for pred in predictions:
        print(f"\nText: {pred['text']}")
        print(f"Label: {pred['predicted_label']}")
        print(f"Confidence: {pred['confidence']:.4f}")

This is my go-to script for classification fine-tuning. Adjust for your use case.

Fine-tuning Other Tasks

Named Entity Recognition (NER)

from transformers import AutoModelForTokenClassification

# Load model for token classification
model = AutoModelForTokenClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=9  # B-PER, I-PER, B-ORG, I-ORG, B-LOC, I-LOC, B-MISC, I-MISC, O
)

# Dataset format
ner_dataset = [
    {
        "tokens": ["John", "lives", "in", "New", "York"],
        "ner_tags": [1, 0, 0, 3, 4]  # B-PER, O, O, B-LOC, I-LOC
    }
]

# Same training process with Trainer

Question Answering

from transformers import AutoModelForQuestionAnswering

model = AutoModelForQuestionAnswering.from_pretrained("bert-base-uncased")

# Dataset format
qa_dataset = [
    {
        "question": "What is AI?",
        "context": "AI stands for Artificial Intelligence.",
        "answers": {
            "text": ["Artificial Intelligence"],
            "answer_start": [16]
        }
    }
]

Text Generation (Causal LM)

from transformers import AutoModelForCausalLM, TextDataset, DataCollatorForLanguageModeling

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Prepare dataset
train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path="train.txt",
    block_size=128
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False  # Causal LM (not masked)
)

# Train with Trainer

Tracking Training Progress

TensorBoard Integration

# Training arguments with logging
training_args = TrainingArguments(
    output_dir='./results',
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy='epoch',
)

# Run training
trainer.train()

# View in TensorBoard
# tensorboard --logdir ./logs

Weights & Biases Integration

# Install: pip install wandb

import wandb

# Login
wandb.login()

# Training arguments
training_args = TrainingArguments(
    output_dir='./results',
    report_to='wandb',  # Enable W&B
    run_name='ticket-classifier-v1',
)

# Train (automatically logs to W&B)
trainer.train()

Custom Callbacks

from transformers import TrainerCallback

class CustomCallback(TrainerCallback):
    """Custom callback for training monitoring."""
    
    def on_epoch_end(self, args, state, control, **kwargs):
        """Called at the end of each epoch."""
        print(f"\nEpoch {state.epoch} complete!")
        print(f"Training loss: {state.log_history[-1].get('loss', 'N/A')}")
        
    def on_train_end(self, args, state, control, **kwargs):
        """Called at the end of training."""
        print("\nTraining complete!")

# Add to Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    callbacks=[CustomCallback()]
)

Hyperparameter Tuning

from transformers import Trainer, TrainingArguments

def model_init():
    """Initialize model for hyperparameter tuning."""
    return AutoModelForSequenceClassification.from_pretrained(
        "distilbert-base-uncased",
        num_labels=3
    )

# Define search space
def hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 5e-5, log=True),
        "num_train_epochs": trial.suggest_int("num_train_epochs", 2, 5),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [8, 16, 32]),
        "weight_decay": trial.suggest_float("weight_decay", 0.0, 0.1),
    }

trainer = Trainer(
    model_init=model_init,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    compute_metrics=compute_metrics,
)

# Run hyperparameter search
best_run = trainer.hyperparameter_search(
    direction="maximize",
    backend="optuna",
    hp_space=hp_space,
    n_trials=10,
)

print("Best hyperparameters:", best_run.hyperparameters)

Best Practices

From my fine-tuning experience:

1. Start with a small learning rate (2e-5 to 5e-5) - fine-tuning, not training from scratch.

2. Use early stopping - prevent overfitting on small datasets.

3. Monitor validation metrics - train accuracy can be misleading.

4. Use warmup - helps training stability.

5. Save checkpoints - training can fail, saves time.

6. Test on held-out data - separate test set for final evaluation.

7. Use appropriate batch sizes - larger = faster but more memory.

8. Try different base models - DistilBERT (fast), BERT (balanced), RoBERTa (accurate).

9. Data quality matters more than quantity - 1000 clean examples > 10000 noisy ones.

10. Document experiments - track what worked and what didn't.

Common Issues

Issue 1: Overfitting

Symptoms: Training accuracy much higher than validation

Solutions:

# 1. Add dropout
model.config.hidden_dropout_prob = 0.2
model.config.attention_probs_dropout_prob = 0.2

# 2. Use weight decay
training_args = TrainingArguments(weight_decay=0.01)

# 3. Early stopping
from transformers import EarlyStoppingCallback
trainer = Trainer(callbacks=[EarlyStoppingCallback(early_stopping_patience=2)])

# 4. More data or data augmentation

Issue 2: Out of Memory

Solutions:

# 1. Reduce batch size
training_args = TrainingArguments(per_device_train_batch_size=8)

# 2. Gradient accumulation
training_args = TrainingArguments(gradient_accumulation_steps=2)

# 3. Use smaller model
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

# 4. Mixed precision
training_args = TrainingArguments(fp16=True)

Issue 3: Poor Performance

Solutions:

# 1. More epochs
training_args = TrainingArguments(num_train_epochs=5)

# 2. Adjust learning rate
training_args = TrainingArguments(learning_rate=3e-5)

# 3. Try different model
model = AutoModelForSequenceClassification.from_pretrained("roberta-base")

# 4. Better data preprocessing
# Clean and balance your dataset

What's Next?

You now know how to fine-tune models on custom data. In Part 4, we'll explore advanced features like custom models, quantization, and parameter-efficient fine-tuning (PEFT/LoRA).

Next: Part 4 - Advanced Features and Techniques

---

Previous: Part 2 - Understanding Models, Tokenizers, and Preprocessing

This article is part of the Hugging Face Transformers 101 series. Check out the series overview for more content.