Hands-On: Fine-Tuning a Transformer Model

import os
import torch
import numpy as np
from datasets import load_dataset, concatenate_datasets
from transformers import (
    AutoTokenizer, 
    AutoModelForSequenceClassification, 
    TrainingArguments, 
    Trainer
)
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
from collections import Counter

# Environment setup
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

def compute_metrics(pred):
    """
    Compute evaluation metrics
    
    Args:
        pred: Prediction object from Trainer
    
    Returns:
        Dictionary of metrics
    """
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    
    # Calculate metrics with zero_division parameter
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, preds, average='binary', zero_division=0
    )
    acc = accuracy_score(labels, preds)
    
    # Print confusion matrix to diagnose - handle case of single label
    cm = confusion_matrix(labels, preds, labels=[0, 1])
    print(f"Confusion matrix:\n{cm}")
    
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

def prepare_dataset(dataset, tokenizer):
    """
    Prepare dataset for training
    
    Args:
        dataset: Original dataset
        tokenizer: Tokenizer to use
    
    Returns:
        Processed dataset
    """
    def tokenize_function(examples):
        return tokenizer(
            examples["text"], 
            padding="max_length",  
            truncation=True, 
            max_length=128
        )
    
    # Tokenize dataset
    tokenized_dataset = dataset.map(
        tokenize_function, 
        batched=True
    )
    
    # Set format for PyTorch - make sure we keep the label column
    tokenized_dataset = tokenized_dataset.remove_columns(
        [col for col in tokenized_dataset.column_names if col not in ["input_ids", "attention_mask", "label"]]
    )
    
    # Rename label column to labels (required by the Trainer)
    tokenized_dataset = tokenized_dataset.rename_column("label", "labels")
    
    # Set format to PyTorch tensors
    tokenized_dataset.set_format("torch")
    
    return tokenized_dataset

def main():
    try:
        # Model and tokenizer initialization
        model_name = "distilbert-base-uncased"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Explicitly set up model for binary classification
        model = AutoModelForSequenceClassification.from_pretrained(
            model_name, 
            num_labels=2,
            id2label={0: "NEGATIVE", 1: "POSITIVE"},
            label2id={"NEGATIVE": 0, "POSITIVE": 1}
        )

        # Load dataset explicitly by class to ensure we get both positive and negative
        print("Loading dataset...")
        # Load negative examples (0) and positive examples (1) separately
        negative_examples = load_dataset("imdb", split='train+test').filter(lambda example: example["label"] == 0).select(range(3000))
        positive_examples = load_dataset("imdb", split='train+test').filter(lambda example: example["label"] == 1).select(range(3000))
        
        # Verify the datasets
        print(f"Negative examples loaded: {len(negative_examples)}")
        print(f"Positive examples loaded: {len(positive_examples)}")
        
        # Combine the datasets
        dataset = concatenate_datasets([negative_examples, positive_examples])
        
        # Shuffle the dataset
        dataset = dataset.shuffle(seed=42)
        
        # Print initial distribution to verify
        print("Combined dataset size:", len(dataset))
        label_counts = Counter(dataset["label"])
        print(f"Combined label distribution: Negative (0): {label_counts[0]}, Positive (1): {label_counts[1]}")
        
        # Prepare dataset
        print("Preparing dataset...")
        tokenized_dataset = prepare_dataset(dataset, tokenizer)
        
        # Split into train and validation
        dataset_split = tokenized_dataset.train_test_split(test_size=0.2, seed=42)
        train_dataset = dataset_split['train']
        val_dataset = dataset_split['test']
        
        # Check label distribution for debugging
        train_labels = train_dataset["labels"].tolist()
        val_labels = val_dataset["labels"].tolist()
        
        print(f"Training set size: {len(train_dataset)}")
        print(f"Validation set size: {len(val_dataset)}")
        print(f"Training label distribution: 0 (Negative): {train_labels.count(0)}, 1 (Positive): {train_labels.count(1)}")
        print(f"Validation label distribution: 0 (Negative): {val_labels.count(0)}, 1 (Positive): {val_labels.count(1)}")

        # Training arguments
        training_args = TrainingArguments(
            output_dir="./results",
            eval_strategy="epoch",
            save_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",
            save_total_limit=2,
            load_best_model_at_end=True,
            metric_for_best_model="f1",
            fp16=torch.cuda.is_available(),
            logging_steps=100,
        )

        # Initialize Trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            compute_metrics=compute_metrics,
        )
        
        # Fine-tune the model
        print("Starting model training...")
        trainer.train()
        
        # Evaluate the model
        print("Evaluating model...")
        results = trainer.evaluate()
        print("Evaluation Results:", results)
        
        # Test on a few examples
        print("\nTesting on a few examples:")
        test_texts = [
            "This movie was absolutely terrible. I hated every minute of it.",
            "What a masterpiece! One of the best films I've ever seen.",
            "The acting was mediocre, but the story was compelling.",
            "I was disappointed with the ending, but overall it was good."
        ]
        
        # Tokenize and prepare test examples
        test_encodings = tokenizer(test_texts, truncation=True, padding=True, return_tensors="pt")
        
        # Get predictions
        with torch.no_grad():
            outputs = model(**test_encodings)
            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
            predicted_classes = torch.argmax(predictions, dim=-1)
        
        # Print results
        for i, text in enumerate(test_texts):
            pred_class = "POSITIVE" if predicted_classes[i] == 1 else "NEGATIVE"
            confidence = predictions[i][predicted_classes[i]].item()
            print(f"Text: {text}")
            print(f"Prediction: {pred_class}, Confidence: {confidence:.4f}\n")
        
        # Save the fine-tuned model and tokenizer
        print("Saving model...")
        trainer.save_model("./fine_tuned_model")
        tokenizer.save_pretrained("./fine_tuned_model")
        print("Model and tokenizer saved successfully.")

    except Exception as e:
        print(f"An error occurred: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()