In this article, well explore the process of fine-tuning language models for text classification. Well do so in three levels: first, by manually adding a classification head in PyTorch* and training the model so you can see the full process; second, by using the Hugging Face* Transformers library to streamline the process; and third, by leveraging PyTorch Lightning* and accelerators to optimize training performance. By the end of this guide, you’ll have a well-rounded understanding of the fine-tuning workflow.
Introduction
The concept of using fine-tuning in Natural Language Processing (NLP) was borrowed from Computer Vision (CV). The CV models were first trained on large datasets such as ImageNet to teach them the basic features of images such as edges or colors. These pretrained models were then fine-tuned on a downstream task such as classifying birds with a relatively small number of labeled examples.
Fine-tuned models typically achieved a higher accuracy than supervised models trained from scratch on the same amount of labeled data.
Despite the popularity and success of transfer learning in CV, for many years it wasnt clear what the analogous pretraining process was for NLP. Consequently, NLP applications required large amounts of labeled data to achieve high performance.
How is Fine-tuning Different from Pretraining?
With pretraining , language models gain a general understanding of languages. During this process, they learn language patterns but typically are not capable of following instructions or answering questions. In the case of GPT models, this self-supervised learning includes predicting the next word (unidirectional) based on their training data, which is often webpages. In the case of BERT (Bidirectional Encoder Representations from Transformers), learning involves predicting randomly masked words (bidirectional) and sentence-order prediction. But how can we adapt language models for our own data or our own tasks?
Fine-tuning continues training a pretrained model to increase its performance on specific tasks. For instance, through instruction fine-tuning you can teach a model to behave more like a chatbot. This is the process for specializing a general purpose model like OpenAI* GPT-4 into an application like ChatGPT* or GitHub* Copilot. By fine-tuning your own language model, you can increase the reliability, performance, and privacy of your model while reducing the associated inference costs compared to subscription-based services, especially if you have a large volume of data or frequent requests.
Fine-Tuning a Language Model for Text Classification
Preprocessing and Preparing DataLoader
Feel free to skip this section if youre comfortable with preprocessing data. Throughout we assume that we have our labeled data saved in train, validation, and test csv files each with a text and a label column. For training, the labels should be numeric, so if thats not the case, youcan use a label_to_id dictionary such as and do a mapping to get the desired format.
For concreteness, we will use BERT as the base model and set the number of classification labels to 4. After running the code below, you are encouraged to swap BERT for DistilBERT which reduces the size of the BERT model by 40%, speeding inference by 60%, while retaining 97% of BERTs language understanding capabilities.
A Quick Look at BERT
BERT was introduced by Google in 2018 and has since revolutionized the field of NLP. Unlike traditional models that process text in a unidirectional manner, BERT is designed to understand the context of a word in a sentence by looking at both its left and right surroundings. This bidirectional approach allows BERT to capture the nuances of language more effectively.
Key Features of BERT
Pretraining: BERT is pretrained on a massive corpus of text,including the entire Wikipedia and BookCorpus. The pretraininginvolves two tasks: Masked Language Modeling (MLM) and NextSentence Prediction (NSP).
Architecture: BERT_BASE has 12 layers (transformer blocks), 768hidden units, and 12 attention heads, totaling 110 millionparameters.
You can run this tutorial on Intel® Tiber AI Cloud, using an Intel® Xeon® CPU instance. This platform provides ample computing resources for smooth execution of our code. For a complete implementation of a fine-tuner in PyTorch Lightning check out our GitHub repository: https://github.com/intel/polite-guard/tree/main/fine-tuner
import os
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer
import pandas as pd
# Parameters
model_ckpt = "bert-base-uncased"
num_labels = 4
batch_size = 32
num_workers = 6
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
# Custom Dataset class
class TextDataset(Dataset):
def __init__(self, dataframe, tokenizer, max_length=512):
self.data = dataframe
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
row = self.data.iloc[idx]
text = row["text"] # Replace "text" with your actual column name for text
label = row["label"] # Replace "label" with your actual column name for labels
# Tokenize the input text
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return {
"input_ids": encoding["input_ids"].squeeze(0), # Remove batch dimension with squeeze
"attention_mask": encoding["attention_mask"].squeeze(0),
"label": torch.tensor(label, dtype=torch.long),
}
# Avoid tokenizer parallelism issues
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Load CSV files
try:
train_df = pd.read_csv("train.csv")
val_df = pd.read_csv("val.csv")
test_df = pd.read_csv("test.csv")
except FileNotFoundError as e:
print(f"Error: {e}. Make sure the CSV files are present in the correct path.")
exit()
# Validate the presence of required columns
required_columns = ["text", "label"]
for df_name, df in zip(["train", "val", "test"], [train_df, val_df, test_df]):
for col in required_columns:
if col not in df.columns:
raise ValueError(f"Column '{col}' is missing in {df_name} dataset.")
# Create Dataset objects
train_dataset = TextDataset(train_df, tokenizer)
val_dataset = TextDataset(val_df, tokenizer)
test_dataset = TextDataset(test_df, tokenizer)
# Create DataLoaders
train_loader = DataLoader(
train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers
)
val_loader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers
)
test_loader = DataLoader(
test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers
)
The Classification Token
The token is typically added at the beginning of the input sequence in transformer-based models such as BERT and its variants. During fine-tuning, the model learns to assign meaningful information to the token, which aggregates the input sequences context. The last hidden state corresponding to the token is then used as a representation of the entire input, which can be passed through a classifier layer for downstream tasks like sentiment analysis, topic categorization, or any task requiring a decision based on the entire sequence. This mechanism allows the model to focus on both the global understanding of the text and task-specific features for accurate predictions.
Unlike traditional models that may rely on static embeddings (like word2vec), transformers generate contextualized embeddings , so that the meaning of a token depends on the tokens around it. The token, as it passes through the layers, becomes increasingly aware of the entire sequences meaning, which makes it a good summary representation for downstream tasks. For some tasks, especially those requiring finer-grained understanding, other strategies might be employed. For instance, for document classification, where every word contributes equally, some models use mean pooling over all token embeddings.
Level 1: PyTorch
In this section, we manually add a classification head to the base model and do the fine-tuning. We achieve this using the class which converts the tokens (or rather token encodings) to embeddings and then then feeds them through the encoder stack to return the hidden states. While is helpful for understanding the idea behind what were doing, to fine-tune for text classification its better practice to work with instead, as we discuss below.
import torch
from torch import nn
from transformers import AutoModel
# Load the base model with AutoModel and add a classifier
class CustomModel(nn.Module):
def __init__(self, model_ckpt, num_labels):
super(CustomModel, self).__init__()
self.model = AutoModel.from_pretrained(model_ckpt) # Base transformer model
self.classifier = nn.Linear(self.model.config.hidden_size, num_labels) # Classification head
def forward(self, input_ids, attention_mask):
# Forward pass through the transformer model
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
# Use the [CLS] token (0-th token in the sequence) for classification
cls_output = outputs.last_hidden_state[:, 0, :] # Shape: (batch_size, hidden_size)
# Pass through the classifier head
logits = self.classifier(cls_output)
return logits
# Initialize the model
model_ckpt = 'bert-base-uncased' # For example; replace with your checkpoint if needed
num_labels = 2 # Replace with the number of labels in your classification task
model = CustomModel(model_ckpt, num_labels)
# Loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
# Training function
def train(model, optimizer, train_loader, loss_fn):
model.train()
total_loss = 0
for batch in train_loader:
optimizer.zero_grad()
# Unpack the batch data
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["label"]
# Forward pass
outputs = model(input_ids, attention_mask)
# Compute loss
loss = loss_fn(outputs, labels)
loss.backward()
# Update the model parameters
optimizer.step()
total_loss += loss.item()
print(f"Train loss: {total_loss / len(train_loader):.2f}")
# Evaluation function
def evaluate(model, test_loader, loss_fn):
model.eval() # Set model to evaluation mode
total_loss = 0
total_acc = 0
total_samples = 0
with torch.no_grad(): # No gradient computation needed during evaluation
for batch in test_loader:
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["label"]
# Forward pass
outputs = model(input_ids, attention_mask)
# Compute loss
loss = loss_fn(outputs, labels)
total_loss += loss.item()
# Compute accuracy
predictions = torch.argmax(outputs, dim=1)
total_acc += torch.sum(predictions == labels).item()
total_samples += labels.size(0)
# Calculate average loss and accuracy
avg_loss = total_loss / len(test_loader)
avg_acc = total_acc / total_samples * 100
print(f"Test loss: {avg_loss:.2f}, Test acc: {avg_acc:.2f}%")
Finally, we can train, evaluate, and save the model.
num_epochs = 3
for epoch in range(num_epochs):
train(model, optimizer, train_loader, loss_fn)
evaluate(model, test_loader, loss_fn)
torch.save(model.state_dict(), "./fine-tuned-model.pt")
Level 2: Hugging Face Transformers
Now, we use the convenience of class that will add the classification head to the base model automatically. Compare this against what we did with the class in the previous section!
Also note that the Trainer class from Hugging Faces Transformerslibrary can directly handle objects without needing a, as it automatically handles batching and shuffling foryou.
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments
model = AutoModelForSequenceClassification.from_pretrained(
model_ckpt,
num_labels=num_labels
)
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
warmup_steps=500,
weight_decay=0.01,
logging_dir="./logs",
logging_steps=10, # Log every 10 steps
evaluation_strategy="steps",
save_steps=500, # Save model checkpoint every 500 steps
load_best_model_at_end=True, # Load the best model at the end of training
metric_for_best_model="accuracy",
)
# Train the model
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset,
)
trainer.train()
trainer.evaluate(test_dataset)
Level 3: PyTorch Lightning
Lightning is, in the words of its documentation, the deep learning framework with batteries included for professional AI researchers and machine learning engineers who need maximal flexibility while super-charging performance at scale.
As we shall see, with a bit of additional organizational code, the Lightning automates the following:
Epoch and batch iteration
, , calls
Calling of , enabling and disabling grads duringevaluation
Checkpoint Saving and Loading
Logging
Accelerator, Multi-GPU, and TPU Support (No calls required.)
Mixed-precision training
You can accelerate training with Intel Gaudi processors, which allow you to conduct more deep learning training at a lower expense. You can try an Intel Gaudi instance for free on Intel® Tiber AI Cloud.
import torch.optim
import torchmetrics
import lightning as L
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger
from transformers import AutoModelForSequenceClassification
# A LightningModule is a torch.nn.Module with added functionality.
# It wraps around a regular PyTorch model.
class LightningModel(L.LightningModule):
def __init__(self, model, learning_rate=5e-5):
super().__init__()
self.learning_rate = learning_rate
self.model = model
self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels)
self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels)
def forward(self, input_ids, attention_mask, labels):
return self.model(input_ids, attention_mask=attention_mask, labels=labels)
def _shared_step(self, batch, batch_idx):
outputs = self(
batch["input_ids"],
attention_mask=batch["attention_mask"],
labels=batch["label"],
)
return outputs
def training_step(self, batch, batch_idx):
outputs = self._shared_step(batch, batch_idx)
self.log("train_loss", outputs["loss"])
return outputs["loss"]
def validation_step(self, batch, batch_idx):
outputs = self._shared_step(batch, batch_idx)
self.log("val_loss", outputs["loss"], prog_bar=True)
logits = outputs["logits"]
self.val_acc(logits, batch["label"])
self.log("val_acc", self.val_acc, prog_bar=True)
def test_step(self, batch, batch_idx):
outputs = self._shared_step(batch, batch_idx)
logits = outputs["logits"]
self.test_acc(logits, batch["label"])
self.log("accuracy", self.test_acc, prog_bar=True)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
return optimizer
# Instantiate the Hugging Face model
model = AutoModelForSequenceClassification.from_pretrained(
model_ckpt,
num_labels=num_labels
)
lightning_model = LightningModel(model)
# Define callbacks and logger
callbacks = [
ModelCheckpoint(save_top_k=1, mode="max", monitor="val_acc") # Save top 1 model
]
logger = TensorBoardLogger(save_dir="./logs", name="fine-tuned-model")
# Initialize the Lightning trainer
trainer = L.Trainer(
max_epochs=3,
callbacks=callbacks,
accelerator="hpu",
precision="bf16-mixed", # By default, HPU training uses 32-bit precision. To enable mixed precision, set the precision flag.
devices="auto",
logger=logger,
log_every_n_steps=10,
)
# Fit and evaluate the model
trainer.fit(lightning_model, train_dataloaders=train_loader, val_dataloaders=val_loader)
trainer.test(lightning_model, train_dataloaders=train_loader, ckpt_path="best")
trainer.test(lightning_model, val_dataloaders=val_loader, ckpt_path="best")
trainer.test(lightning_model, test_dataloaders=test_loader, ckpt_path="best")
While the Transformers Trainer class supports distributed training, it doesnt offer the same level of integration and flexibility as Lightning when it comes to advanced features like custom callbacks, logging, and seamless scaling across multiple GPUs or nodes.
Practical Advice
Now that youre familiar with the fine-tuning process, you mightwonder how you can apply it to your specific task. Heres somepractical advice:
Collect real data for your task, or generate synthetic data. See, for instance, Synthetic Data Generation with Language Models: A Practical Guide.
Fine-tune a relatively small model.
Evaluate your language model on your test set, and on a benchmark if available for your task.
Increase the training dataset size, base model size, and, if necessary, task complexity.
Keep in mind that the standard or conventional fine-tuning of languagemodels as described in this writing can be expensive. Rather thanupdating all the weights and biases, we could update only the lastlayer as follows:
# Freeze all layers
for param in model.parameters():
param.requires_grad = False
# Unfreeze the pre_classifier layer
for param in model.pre_classifier.parameters():
param.requires_grad = True
# Unfreeze the classifier layer
for param in model.classifier.parameters():
param.requires_grad = True
In future articles, we shall discuss more efficient fine-tuning techniques, so stay tuned!
For more AI development how-to content, visit Intel® AI DevelopmentResources.
Finally, if you found this article valuable, please consider giving it a and sharing it with your network.
** Acknowledgments**
The author thanks Jack Erickson for providing detailed feedback on an earlier draft of this work.