Text Classification with BERT in PyTorch | by Ruben Winastwan | Towards Data Science

Back in 2018, Google developed a powerful Transformer-based machine learning model for NLP applications that outperforms previous language models in different benchmark datasets. And this model is called BERT.
In this post, we’re going to use a pre-trained BERT model from Hugging Face for a text classification task. As you might already know, the main goal of the model in a text classification task is to categorize a text into one of the predefined labels or tags.
Specifically, soon we’re going to use the pre-trained BERT model to classify whether the text of a news article can be categorized as sport, politics, business, entertainment, or tech category.
But before we dive into the implementation, let’s talk about the concept behind BERT briefly.
BERT is an acronym for Bidirectional Encoder Representations from Transformers. The name itself gives us several clues to what BERT is all about.
BERT architecture consists of several Transformer encoders stacked together. Each Transformer encoder encapsulates two sub-layers: a self-attention layer and a feed-forward layer.
There are two different BERT models:
BERT base, which is a BERT model consists of 12 layers of Transformer encoder, 12 attention heads, 768 hidden size, and 110M parameters.BERT large, which is a BERT model consists of 24 layers of Transformer encoder,16 attention heads, 1024 hidden size, and 340 parameters.
BERT base, which is a BERT model consists of 12 layers of Transformer encoder, 12 attention heads, 768 hidden size, and 110M parameters.
BERT large, which is a BERT model consists of 24 layers of Transformer encoder,16 attention heads, 1024 hidden size, and 340 parameters.
There are at least two reasons why BERT is a powerful language model:
It is pre-trained on unlabeled data extracted from BooksCorpus, which has 800M words, and from Wikipedia, which has 2,500M words.As the name suggests, it is pre-trained by utilizing the bidirectional nature of the encoder stacks. This means that BERT learns information from a sequence of words not only from left to right, but also from right to left.
It is pre-trained on unlabeled data extracted from BooksCorpus, which has 800M words, and from Wikipedia, which has 2,500M words.
As the name suggests, it is pre-trained by utilizing the bidirectional nature of the encoder stacks. This means that BERT learns information from a sequence of words not only from left to right, but also from right to left.
BERT model expects a sequence of tokens (words) as an input. In each sequence of tokens, there are two special tokens that BERT would expect as an input:
[CLS]: This is the first token of every sequence, which stands for classification token.[SEP]: This is the token that makes BERT know which token belongs to which sequence. This special token is mainly important for a next sentence prediction task or question-answering task. If we only have one sequence, then this token will be appended to the end of the sequence.
[CLS]: This is the first token of every sequence, which stands for classification token.
[SEP]: This is the token that makes BERT know which token belongs to which sequence. This special token is mainly important for a next sentence prediction task or question-answering task. If we only have one sequence, then this token will be appended to the end of the sequence.
To make it more clear, let’s say we have a text consisting of the following short sentence:
As a first step, we need to transform this sentence into a sequence of tokens (words) and this process is called tokenization.
Although we have tokenized our input sentence, we need to do one more step. We need to reformat that sequence of tokens by adding[CLS] and [SEP] tokens before using it as an input to our BERT model.
Luckily, we only need one line of code to transform our input sentence into a sequence of tokens that BERT expects as we have seen above. We will use BertTokenizer to do this and you can see how we do this later on.
It is also important to note that the maximum size of tokens that can be fed into BERT model is 512. If the tokens in a sequence are less than 512, we can use padding to fill the unused token slots with [PAD] token. If the tokens in a sequence are longer than 512, then we need to do a truncation.
And that’s all that BERT expects as input.
BERT model then will output an embedding vector of size 768 in each of the tokens. We can use these vectors as an input for different kinds of NLP applications, whether it is text classification, next sentence prediction, Named-Entity-Recognition (NER), or question-answering.
For a text classification task, we focus our attention on the embedding vector output from the special [CLS] token. This means that we’re going to use the embedding vector of size 768 from [CLS] token as an input for our classifier, which then will output a vector of size the number of classes in our classification task.
Below is the illustration of the input and output of the BERT model.
Now we’re going to jump into our main topic to classify text with BERT. In this post, we’re going to use the BBC News Classification dataset. If you want to follow along, you can download the dataset on Kaggle.
This dataset is already in CSV format and it has 2126 different texts, each labeled under one of 5 categories: entertainment, sport, tech, business, or politics.
Let’s take a look at what the dataset looks like.
import pandas as pd

datapath = 'bbc-text.csv'
df = pd.read_csv(datapath)
df.head()

df.groupby(['category']).size().plot.bar()

<AxesSubplot:xlabel='category'>
 

As you can see, the dataframe only has two columns, which is category that will be our label, and text which will be our input data for BERT.
As you might already know from the previous section, we need to transform our text into the format that BERT expects by adding [CLS] and [SEP] tokens. We can do this easily with BertTokenizer class from Hugging Face.
First, we need to install Transformers library via pip:
pip install transformers
To make it easier for us to understand the output that we get from BertTokenizer, let’s use a short text as an example.
Here is the explanation of BertTokenizer parameters above:
padding : to pad each sequence to the maximum length that you specify.
max_length : the maximum length of each sequence. In this example we use 10, but for our actual dataset we will use 512, which is the maximum length of a sequence allowed for BERT.
truncation : if True, then the tokens in each sequence that exceed the maximum length will be truncated.
return_tensors : the type of tensors that will be returned. Since we’re using Pytorch, then we use pt. If you use Tensorflow, then you need to use tf .
The outputs that you see from bert_input variable above are necessary for our BERT model later on. But what do those outputs mean?
The first row is input_ids , which is the id representation of each token. We can actually decode these input ids into the actual tokens as follows:
The first row is input_ids , which is the id representation of each token. We can actually decode these input ids into the actual tokens as follows:
As you can see, the BertTokenizer takes care of all of the necessary transformations of the input text such that it’s ready to be used as an input for our BERT model. It adds [CLS], [SEP], and [PAD] tokens automatically. Since we specified the maximum length to be 10, then there are only two [PAD] tokens at the end.
2. The second row is token_type_ids , which is a binary mask that identifies in which sequence a token belongs. If we only have a single sequence, then all of the token type ids will be 0. For a text classification task, token_type_ids is an optional input for our BERT model.
3. The third row is attention_mask , which is a binary mask that identifies whether a token is a real word or just padding. If the token contains [CLS], [SEP], or any real word, then the mask would be 1. Meanwhile, if the token is just padding or [PAD], then the mask would be 0.
As you might notice, we use a pre-trained BertTokenizer from bert-base-cased model. This pre-trained tokenizer works well if the text in your dataset is in English.
If you have datasets from different languages, you might want to use bert-base-multilingual-cased. Specifically, if your dataset is in German, Dutch, Chinese, Japanese, or Finnish, you might want to use a tokenizer pre-trained specifically in these languages. You can check the name of the corresponding pre-trained tokenizer here.
To sum up, below is the illustration of what BertTokenizer does to our input sentence.
Now that we know what kind of output that we will get from BertTokenizer , let’s build a Dataset class for our news dataset that will serve as a class to generate our news data.
In the above implementation, we define a variable called labels , which is a dictionary that maps the category in the dataframe into the id representation of our label. Notice that we also call BertTokenizer in the __init__ function above to transform our input texts into the format that BERT expects.
After defining dataset class, let’s split our dataframe into training, validation, and test set with the proportion of 80:10:10.
So far, we have built a dataset class to generate our data. Now let’s build the actual model using a pre-trained BERT base model which has 12 layers of Transformer encoder.
If your dataset is not in English, it would be best if you use bert-base-multilingual-cased model. If your data is in German, Dutch, Chinese, Japanese, or Finnish, you can use the model pre-trained specifically in these languages. You can check the name of the corresponding pre-trained model here.
As you can see from the code above, BERT model outputs two variables:
The first variable, which we named _ in the code above, contains the embedding vectors of all of the tokens in a sequence.
The second variable, which we named pooled_output, contains the embedding vector of [CLS] token. For a text classification task, it is enough to use this embedding as an input for our classifier.
We then pass the pooled_output variable into a linear layer with ReLU activation function. At the end of the linear layer, we have a vector of size 5, each corresponds to a category of our labels (sport, business, politics, entertainment, and tech).
Now it’s time for us to train the model. The training loop will be a standard PyTorch training loop.
We train the model for 5 epochs and we use Adam as the optimizer, while the learning rate is set to 1e-6. We also need to use categorical cross entropy as our loss function since we’re dealing with multi-class classification.
It is recommended that you use GPU to train the model since BERT base model contains 110 million parameters.
After 5 epochs with the above configuration, you’ll get the following output as an example:
Obviously you might not get similar loss and accuracy values as the screenshot above due to the randomness of training process. If you haven’t got a good result after 5 epochs, try to increase the epochs to, let’s say, 10 or adjust the learning rate.
Now that we have trained the model, we can use the test data to evaluate the model’s performance on unseen data. Below is the function to evaluate the performance of the model on the test set.
After running the code above, I got the accuracy of 0.994 from the test data. The accuracy that you’ll get will obviously slightly differ from mine due to the randomness during the training process.
Now you know the step on how we can leverage a pre-trained BERT model from Hugging Face for a text classification task. I hope this post helps you to get started with BERT.
One thing to remember is that we can use the embedding vectors from BERT to do not only a sentence or text classification task, but also the more advanced NLP applications such as question answering, next sentence prediction, or Named-Entity-Recognition (NER) tasks.
You can find all of the code snippets demonstrated in this post in this notebook.