Data ScienceNLPCode available

Finetune Transformers Model for your own NLP Tasks 🤗

ByMichele De Filippo

15 Feb 2024

When it comes to NLP, Hugging Face should be arrow in your quiver. I’m sure you’ve heard of it somewhere or seen it numerous times on our blogs, it’s more than just an emoji 🤗. Hugging Face is an AI company based in New York that provides an open-source platform for building, training, and deploying machine learning models. Transformers, the focus of this article, is its most popular library providing access to thousands of pre-trained models to perform tasks on text, image, and audio.

Before getting into the details of finetuning, there are 2 steps that need to be done: choosing a framework and a base model.

Choose a Deep Learning Framework

PyTorch, TensorFlow, and Keras are the three most popular deep learning frameworks. While all of them can finetune models, we need to choose the right one for your project.

TensorFlow was developed by the Google Brain team and made open source in 2015. It offers both high-level and low-level APIs and supports Python, JavaScript, C++, C#, Haskell, Julia, MATLAB and so on. TensorFlow can train models on both CPUs and GPUs and allow for parallel processing of data across multiple hardware.

PyTorch was developed by Meta’s AI team and made open source in 2016. It is widely adopted in academia for its simple syntax and Pythonic style. It utilizes a dynamic computation graph to speed up the execution, whereas TensorFlow adopts a static one that needs to be optimized to improve speed.

Keras was released in 2015, as “an API designed for human beings, not machines.” It’s not the fastest or the most flexible one, but the biggest advantage is that it has an extremely simple and user-friendly API. If it’s your first time interacting with deep learning frameworks, you should go for it.

In the table below we compare three frameworks point by point to provide a clearer picture.

	TensorFlow	PyTorch	Keras
API Level	Both high and low	Low	High
Architecture	Complex	Complex	Easy to understand
Dataset	Large datasets, high performance	Large datasets, high performance	Small datasets only
Speed	Fast	Fast	Slow
Deploy	Easy	Hard	Easy

TensorFlow is by far in the leading position in the industry mainly for three reasons. First, it’s an end-to-end framework, meaning it's ready for deployment, whereas with PyTorch and Keras, there's still a need to set up the backend. Second, it has comprehensive documentation, tons of well-trained models and tutorials. Third, it offers better visualization for the debugging process by using a module called TensorBoard.

Select a Base Model

Hugging Face offers a wide range of base models. When choosing a model, we need to consider factors like the size of your dataset, computational resources, type of tasks, etc. BERT is the original model trained on a large corpus of unlabeled text that outperformed state-of-the-art in 2018. Following a similar design of BERT with minor modifications, we have RoBERTa and XLNet which have better performance, and DistilBERT which is faster in terms of inference.

	BERT	RoBERTa	DistilBERT	XLNet
Pretrain Dataset	16 GB BERT data and 3.3 billion words	144 GB larger than BERT	same as BERT	97GB larger than BERT
Model Parameters	110 million	340 million	66 million	110 million
Performance	baseline	2-20% improvement	5% degradation	2-15% improvement

For most general-purpose tasks, DistilBERT is a good starting point. It preserves 95% performance of the original BERT with only half of the model size, therefore greatly speeding up the inference. However, if you are looking for the best performance, RoBERTa is the choice.

Start the Finetuning

For the demonstration, we'll be using TensorFlow. While it's simple to install TensorFlow on Linux, the process is more complicated for other systems. Google Colab usually offers the latest stable version of TensorFlow, along with a free GPU, which is a perfect starting point for us to experiment and get familiar with framework.

To switch version of TensorFlow, we can use

%tensorflow_version 2.x # for TensorFlow 2.x

Load and Preprocess the Dataset

For this example, we will be finetuning for the task of spam detection. The dataset has two columns, label and text. The label is either “ham” or “spam”.

Glimpse of Dataset 👀

After loading the dataset from the txt file, we first need to convert the label to numerical values. Here we use one-hot encoding, creating two extra columns, “spam” and “ham”. The value of the corresponding column is assigned 0 or 1 depending on the label. After defining all variables, we split the data into train and test. We only finetune the model with the train set and evaluate on the test set. Note that we should try to provide a balanced dataset, or it will affect the performance.

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.read_csv('dataset.txt', sep='\t', names=["label", "text"])
X = list(df['text'])
y = list(df['label'])
y = list(pd.get_dummies(y, drop_first=True)['spam']) # one-hot encoding 

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=0)

Tokenization

We've selected DistilBERT as our base model. To ensure consistency, use the same tokenizer for pretrain and finetune.

# !pip install transformers
|from transformers import DistilBertTokenizerFast

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(X_train, truncation=True, padding=True)
test_encodings = tokenizer(X_test, truncation=True, padding=True)

For the above code, we directly import the DistilBERT tokenizer. However, sometimes we are unsure about what tokenizer to use, or we may want to change it later. In this case, we can use AutoTokenizer instead. It will automatically detect the right tokenizer to use based on the model name or model checkpoint you specify.

from transformers import AutoTokenizer

# Specify model name
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(X_train, truncation=True, padding=True)
test_encodings = tokenizer(X_test, truncation=True, padding=True)

Convert to TensorFlow Dataset

After tokenization, the data is just a dictionary of tokenized sequences. In this step, we convert it to a TensorFlow dataset, an iterable object. Other than keeping the consistency of data and framework, we do it for two other reasons. First, Tensorflow is good with large datasets. Using its data format can increase the speed of data operations like multi-thread loading with GPU. Second, it provides lots of built-in functions, like batch, shuffle and repeat.

import tensorflow as tf

train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(train_encodings),
    y_train
))

test_dataset = tf.data.Dataset.from_tensor_slices((
    dict(test_encodings),
    y_test
))

In TensorFlow dataset, the first element is a dictionary of input features, the second element is the corresponding label.

train_dataset

<_TensorSliceDataset
    element_spec={
        'input_ids': TensorSpec(
            shape=(238,),
            dtype=tf.int32,
            name=None
        ),
        'attention_mask': TensorSpec(
            shape=(238,),
            dtype=tf.int32,
            name=None
        )
    },
    TensorSpec(
        shape=(),
        dtype=tf.int32,
        name=None
    )
>

Set Training Arguments and Load the Model

In this section, we set up the configuration for training and initialized an instance of the pre-trained DistilBERT model.

from transformers import TFDistilBertForSequenceClassification, TFTrainer, TFTrainingArguments

training_args = TFTrainingArguments(
    output_dir='./results',  # model output directory
    num_train_epochs=2,  # times the model go through entire dataset
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,
    warmup_steps=500,  # for stabilizing the intial training 
    weight_decay=0.01, # prevent overfitting by penalizing large weights
    eval_steps=100
)

with training_args.strategy.scope():
    model = TFDistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

Train and Evaluate

Using TFTrainer, we combine the training configs with model and data. TFTrainer automatically detects and utilizes available GPUs without the need for explicit specification. This part is where the finetuning actually happens. Depending on the size of model, dataset, training configuration and hardware, the time will vary. For our dataset of 1,115 rows, it takes around 4 minutes to finish the entire process.

# construct the trainer

trainer = TFTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset
)

# start training
trainer.train()

The evaluate() function gives us a loss, measuring how different the model's predictions are from the actual target values.

trainer.evaluate(test_dataset)

# output
{'eval_loss': 0.019664422103336878}

In our case, it’s quite low, but that doesn’t necessarily mean a good performance. We need to look at other metrics, like accuracy, precision, recall, or F1 score.

To explicitly get the predictions, we need to use the predict() function.

output=trainer.predict(test_dataset)[1]

from sklearn.metrics import confusion_matrix

cm=confusion_matrix(y_test,output)
cm

# output
array([[955,   0],
       [  0, 160]])

Under the directory results, we can find two types of statistics about the model.

Checkpoints are snapshots of the model's weights at different time points. They can be used to resume training next time or evaluate the model without retraining.

# Load the base model
model = TFDistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

# Load model from checkpoint
checkpoint_path = './results/checkpoint'  # specify the path to your checkpoint directory
model = TFDistilBertForSequenceClassification.from_pretrained(checkpoint_path)

Runs contain logs of the training process. We can visualize it with TensorBoard for debugging. Read more about how to get started with Tensorboard here.

# Load the TensorBoard notebook extension
%load_ext tensorboard
%tensorboard --logdir=results/runs

An Example of TensorBoard

Now that you've fine-tuned your model, you can share it to other people by uploading it to Hugging Face repositories (same as GitHub repositories). Follow this link for a detailed guide on how to upload and share your model.