A Beginner‘s Guide to Text Classification with BERT

In recent years, transfer learning with large pre-trained language models like BERT has revolutionized the field of natural language processing (NLP). Models like BERT can be fine-tuned on a wide variety of downstream NLP tasks, often achieving state-of-the-art performance with minimal task-specific architecture changes.

In this tutorial, we‘ll walk through how to use a pre-trained BERT model to build a powerful text classifier, even if you‘re new to NLP and deep learning. By the end, you‘ll be able to train your own BERT-based model that can classify text into predefined categories with high accuracy. Let‘s jump in!

What is BERT?

BERT, which stands for Bidirectional Encoder Representations from Transformers, is a large language model developed by researchers at Google. It has been pre-trained on a massive amount of unlabeled text data (Wikipedia + BookCorpus) using two novel unsupervised learning tasks:

1. Masked Language Modeling (MLM): Some percentage of words in each input sequence are randomly masked, and the model learns to predict the original masked words based on the surrounding context. This trains the model to develop a bidirectional understanding of language.

2. Next Sentence Prediction (NSP): The model is given pairs of sentences as input and learns to predict if the second sentence follows the first in the original text. This trains the model to understand relationships between sentences.

The pre-training allows BERT to develop a deep and nuanced understanding of language structure that can then be leveraged for a variety of downstream NLP tasks. The original BERT model comes in two sizes:

• BERT-Base: 12 transformer layers, 768 hidden size, 12 attention heads, 110M parameters
• BERT-Large: 24 transformer layers, 1024 hidden size, 16 attention heads, 340M parameters

For most applications BERT-Base is sufficient, while BERT-Large may give an additional performance boost in some cases.

Why use BERT for Text Classification?

Text classification is a common NLP task that involves assigning predefined categories to text documents. Some example applications include:

• Sentiment analysis (positive vs. negative)
• Spam detection (spam vs. not spam)
• Topic labeling (politics, sports, technology, etc.)
• Intent detection for chatbots (booking a flight, checking account balance, etc.)

Historically, building an accurate text classifier required training a model from scratch on a large labeled dataset, which was time-consuming and required substantial domain expertise.

However, by leveraging a pre-trained BERT model, it‘s possible to build highly accurate text classifiers even with a relatively small amount of labeled training data. This is because the general language understanding that BERT developed during pre-training provides a strong foundation for learning more specific tasks like text classification.

Concretely, to use BERT for text classification, you:

1. Feed tokenized text data into the pre-trained BERT model to generate embeddings
2. Add a simple classification layer on top of the BERT embeddings
3. Train this end-to-end model on a labeled text classification dataset, fine-tuning the pre-trained BERT parameters in the process

This fine-tuning process is relatively quick, typically requiring just a few epochs to converge. And the resulting model is often much more accurate than one trained from scratch on the same labeled data.

With this context in mind, let‘s walk through the actual steps to build a BERT classifier.

Step 1: Load a Pre-trained BERT Model

The first step is to load a pre-trained BERT model and its associated tokenizer. The easiest way to do this is using the transformers library by Hugging Face, which provides a simple API for working with many pre-trained language models.

First install the library:

!pip install transformers

Then load the BERT model and tokenizer:

from transformers import BertTokenizer, TFBertForSequenceClassification

model_name = ‘bert-base-uncased‘
tokenizer = BertTokenizer.from_pretrained(model_name)
model = TFBertForSequenceClassification.from_pretrained(model_name, num_labels=2)

Here we‘re using the ‘bert-base-uncased‘ model, which is the smaller BERT variant (12-layers, 768-hidden, 12-heads, 110M parameters). The ‘uncased‘ part means it was pre-trained on lowercase text.

The num_labels argument specifies the number of classes in our classifier. Here we‘re assuming a binary classification task (e.g. positive vs. negative sentiment), but you can change this for multi-class problems.

Step 2: Preprocess and Tokenize Text Data

Next, we need to preprocess our text data and convert it into the format expected by BERT. This involves:

1. Tokenization: Breaking text into tokens (words, subwords, or characters) that are part of BERT‘s vocabulary
2. Truncating or padding the sequence to a fixed length
3. Adding special tokens like [CLS] and [SEP] to mark the start and end of sequences

The BERT tokenizer loaded in step 1 can handle all of this for us:

def preprocess(text, max_len=256):
    tokens = tokenizer.encode_plus(text, max_length=max_len, truncation=True,
                                   padding=‘max_length‘, add_special_tokens=True, 
                                   return_tensors=‘tf‘)
    return tokens[‘input_ids‘], tokens[‘attention_mask‘]

input_ids, attention_mask = preprocess("This movie was great!")

The encode_plus method performs the tokenization and formatting. Here the key arguments are:

• max_length: The maximum sequence length. Sequences longer than this will be truncated.
• padding: Shorter sequences will be padded to this length.
• return_tensors: Specifies we want TensorFlow tensors as output.

The function returns two tensors:

• input_ids: The token IDs for the input sequence.
• attention_mask: Indicates which tokens are real (1) vs padding (0).

Step 3: Create a tf.data.Dataset

For training, we need to create a tf.data.Dataset object that generates batches of (input_ids, attention_mask, labels) tuples. Assuming we have lists of input texts and corresponding labels:

texts = [
    "This movie was great!",
    "Absolutely terrible acting.",
    "The plot was a bit slow but overall I enjoyed it.",
    ...
]
labels = [1, 0, 1, ...]

def create_dataset(texts, labels):
    input_ids, attention_mask = zip(*[preprocess(text) for text in texts])
    labels = tf.constant(labels)
    input_ids = tf.stack(input_ids)
    attention_mask = tf.stack(attention_mask)
    dataset = tf.data.Dataset.from_tensor_slices((input_ids, attention_mask, labels))
    dataset = dataset.map(lambda x, y, z: ({"input_ids":x, "attention_mask":y}, z))
    return dataset

dataset = create_dataset(texts, labels)
dataset = dataset.shuffle(1000).batch(32)

The key steps are:

1. Preprocess all the input texts, converting them to (input_ids, attention_mask) tuples.
2. Stack the input_ids, attention_mask, and labels into separate tensors.
3. Create a tf.data.Dataset from these tensors using from_tensor_slices.
4. Use a map function to convert the Dataset elements to the format expected by our BERT model – a tuple of ({"input_ids":…, "attention_mask":…}, labels).
5. Shuffle and batch the dataset.

Step 4: Fine-tune BERT

We‘re now ready to fine-tune our pre-trained BERT model on our labeled text data. With the Keras API, this is as simple as calling fit() on our model:

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer, loss=loss, metrics=[‘accuracy‘])

history = model.fit(dataset, epochs=3)

Some key things to note:

• We‘re using a very small learning rate (1e-5). This is important because we don‘t want to make large updates to the pre-trained BERT parameters, which are already quite good. We‘re just fine-tuning them slightly for our specific task.

• We‘re using the SparseCategoricalCrossentropy loss, which is appropriate for a classification problem where the labels are integers representing the class. The from_logits=True argument specifies that the model outputs logits rather than probabilities.

• We‘re training for just 3 epochs. Because we‘re starting from a pre-trained model, we often don‘t need many epochs to converge. Training for too long risks overfitting to the training data.

Step 5: Evaluate and Predict

After fine-tuning, we can evaluate our model‘s performance on a held-out test set:

test_loss, test_acc = model.evaluate(test_dataset, verbose=2)
print(‘Test accuracy:‘, test_acc)

And we can use it to make predictions on new text:

def predict(text):
    input_ids, attention_mask = preprocess(text)
    outputs = model({"input_ids":[input_ids], "attention_mask":[attention_mask]})
    logits = outputs.logits
    probs = tf.nn.softmax(logits, axis=1).numpy()[0]
    return probs

probs = predict("This movie was amazing!")
print(f"Positive probability: {probs[1]:.2f}")

Tips and Tricks

Here are a few tips to get the most out of BERT for text classification:

1. Use a validation set during training to monitor for overfitting. If validation loss starts increasing while training loss is still decreasing, you may be starting to overfit.

2. Experiment with different learning rates and numbers of epochs. The optimal values can vary depending on the size and complexity of your dataset.

3. If your dataset is very small, consider using a smaller BERT variant like DistilBERT or ALBERT. These have fewer parameters and may be less prone to overfitting on small data.

4. For multi-class problems, make sure to set num_labels correctly when loading the pre-trained model.

5. Consider adding additional dense layers between the BERT outputs and the final classification layer. This can give the model more flexibility to adapt to your specific task.

Conclusion

In this tutorial, we‘ve seen how to use a pre-trained BERT model to build a powerful text classifier with just a few lines of code. By leveraging the deep language understanding that BERT develops during pre-training, we can often achieve excellent results even with relatively small labeled datasets.

Of course, BERT is not the only game in town. Newer models like RoBERTa, XLNet, and ELECTRA have been shown to outperform BERT on many tasks. However, the general approach of fine-tuning a pre-trained model remains the same.

I hope this guide has given you a solid foundation for using BERT in your own NLP projects. The complete code for this tutorial is available on [GitHub](). If you have any questions or suggestions, feel free to reach out on [Twitter]() or [LinkedIn](). Happy coding!