Fine-tuning pretrained NLP models with Huggingface’s Trainer | by Vincent Tan | Towards Data Science

Motivation: While working on a data science competition, I was fine-tuning a pre-trained model and realised how tedious it was to fine-tune a model using native PyTorch or Tensorflow. I experimented with Huggingface’s Trainer API and was surprised by how easy it was. As there are very few examples online on how to use Huggingface’s Trainer API, I hope to contribute a simple example of how Trainer could be used to fine-tune your pretrained model.
Before we start, here are some prerequisites to understand this article:
Intermediate understanding of PythonBasic understanding in training neural network modelsBasic understanding of transfer learning
Intermediate understanding of Python
Basic understanding in training neural network models
Basic understanding of transfer learning
To save your time, I will just provide you the code which can be used to train and predict your model with Trainer API. However, if you are interested in understanding how it works, feel free to read on further.
Step 1: Initialise pretrained model and tokenizer
In the code above, the data used is a IMDB movie sentiments dataset. The data allows us to train a model to detect the sentiment of the movie review- 1 being positive while 0 being negative. This is a NLP task of sequence classification, as we want to classify each review(sequence of text) into positive or negative.
There are many pretrained models which we can use to train our sentiment analysis model, let us use pretrained BERT as an example. There are many variants of pretrained BERT model, bert-base-uncased is just one of the variants. You can search for more pretrained model to use from Huggingface Models page.
model_name = "bert-base-uncased"tokenizer = BertTokenizer.from_pretrained(model_name)model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)
Since we are using a pretrained model, we need to ensure that the input data is in the same form as what the pretrained model was trained on. Thus, we would need to instantiate the tokenizer using the name of the model.
Now that the model and tokenizer have been initialised, we can proceed to preprocess the data.
Step 2: Preprocess text using pretrained tokenizer
X_train_tokenized = tokenizer(X_train, padding=True, truncation=True, max_length=512)X_val_tokenized = tokenizer(X_val, padding=True, truncation=True, max_length=512)
Let us preprocess the text using the tokenizer intialised earlier.
The input text that we are using for the tokenizer is a list of strings.
We have set padding=True, truncation=True, max_length=512 so that we can get same length inputs for the model- the long texts will be truncated to 512 tokens while the short texts will have extra tokens added to make it 512 tokens.
512 tokens is used because this is the maximum token length that the BERT model can take.
After tokenizing your text, you will get a python dictionary with 3 keys:
Input_idstoken_type_idsattention_mask
Input_ids
token_type_ids
attention_mask
Step 3: Creating a torch dataset
The Trainer API requires the model to be in a torch.utils.data.Dataset class. Hence, we would need to create a new class that inherits from the torch Dataset class.
In the inherited class, we need to have the __getitem__and __len__ method which allows Trainer to create batches of data and to obtain the length respectively.
class Dataset(torch.utils.data.Dataset):        def __init__(self, encodings, labels=None):                  self.encodings = encodings                self.labels = labels         def __getitem__(self, idx):        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}        if self.labels:            item["labels"] = torch.tensor(self.labels[idx])        return item     def __len__(self):        return len(self.encodings["input_ids"])train_dataset = Dataset(X_train_tokenized, y_train)val_dataset = Dataset(X_val_tokenized, y_val)
The purpose of setting the default labels parameter as None is so that we can reuse the class to make prediction on unseen data as these data do not have labels.
The __getitem__ method basically returns a dictionary of values for each text. By running the method, it creates a dictionary with input_ids , attention_mask and token_type_ids for each text when the data is batched during the training process.
The __len__method needs to return the length of the input data.
Step 4: Defining TrainingArguments and Trainer
def compute_metrics(p):        pred, labels = p    pred = np.argmax(pred, axis=1)    accuracy = accuracy_score(y_true=labels, y_pred=pred)    recall = recall_score(y_true=labels, y_pred=pred)    precision = precision_score(y_true=labels, y_pred=pred)    f1 = f1_score(y_true=labels, y_pred=pred)    return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1} # Define Trainerargs = TrainingArguments(    output_dir="output",    evaluation_strategy="steps",    eval_steps=500,    per_device_train_batch_size=8,    per_device_eval_batch_size=8,    num_train_epochs=3,    seed=0,    load_best_model_at_end=True,)trainer = Trainer(    model=model,    args=args,    train_dataset=train_dataset,    eval_dataset=val_dataset,    compute_metrics=compute_metrics,    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],) # Train pre-trained modeltrainer.train()
Here is where the magic of the Trainer function is. We can define the training parameters in the TrainingArguments and Trainer class as well as train the model with a single command.
We need to first define a function to calculate the metrics of the validation set. Since this is a binary classification problem, we can use accuracy, precision, recall and f1 score.
Next, we specify some training parameters, set the pretrained model, train data and evaluation data in the TrainingArgs and Trainer class.
After we have defined the parameters , simply run trainer.train() to train the model.
Step 5: Making prediction
# Tokenize test dataX_test_tokenized = tokenizer(X_test, padding=True, truncation=True, max_length=512) # Create torch datasettest_dataset = Dataset(X_test_tokenized) # Load trained modelmodel_path = "output/checkpoint-50000"model = BertForSequenceClassification.from_pretrained(model_path, num_labels=2) # Define test trainertest_trainer = Trainer(model) # Make predictionraw_pred, _, _ = test_trainer.predict(test_dataset) # Preprocess raw predictionsy_pred = np.argmax(raw_pred, axis=1)
After the model is trained, we repeat the same steps for the test data:
Tokenize test data with pretrained tokenizerCreate torch datasetLoad trained modelDefine Trainer
Tokenize test data with pretrained tokenizer
Create torch dataset
Load trained model
Define Trainer
To load the trained model from the previous steps, set the model_path to the path containing the trained model weights.
To make prediction, only a single command is needed as well test_trainer.predict(test_dataset) .
After making a prediction, you will only get the raw prediction. Additional preprocessing steps will be needed to get it to a usable format.
Since the task is just a simple sequence classification task, we can just obtain the argmax across axis 1. Note that other NLP tasks may require different ways to preprocess the raw predictions.
I hope this post will help to simplify your process of fine-tuning pretrained NLP models. Feel free to read more into official Huggingface documentation to understand the code better and learn about what other features it could do.
Edit 1 (23/6/21): Removed save_steps parameter from TrainingArgument as it is ignored when load_best_model_at_end is set to True. Thanks 杨程 for the feedback!