Load the BERT tokenizer
Text classification is a common natural language processing task that assigns predefined categories to text. Multi-class classification refers to scenarios where text can be categorized into one of several possible classes, as opposed to just two (binary classification). Some examples of multi-class text classification include:
- Categorizing emails as Promotions, Social, Updates, or Forums
- Classifying customer support tickets as Bug Reports, Feature Requests, or General Inquiries
- Identifying the topic or department of a news article like Sports, Business, Technology, or Politics
Recently, a powerful new approach for text classification has emerged using large pre-trained language models like BERT (Bidirectional Encoder Representations from Transformers). BERT and similar models are trained on massive amounts of text data to build general language understanding. These models can then be fine-tuned on your specific text classification dataset, often achieving state-of-the-art performance.
In this post, we‘ll explore how to use BERT for multi-class text classification using the popular Keras deep learning library. We‘ll walk through the key steps from data preparation to model training and evaluation. Let‘s dive in!
What is BERT?
BERT is a recent natural language processing breakthrough that has crushed established benchmarks on a wide variety of language understanding tasks. The model is pre-trained on a large corpus of unlabeled text including the entire Wikipedia (2,500 million words!) and Book Corpus (800 million words).
BERT‘s key innovation is applying the bidirectional training of Transformer models to language modelling. This allows the BERT model to have a deeper sense of language context and flow compared to single-direction language models. The pre-trained BERT model can be fine-tuned with just one additional output layer to create state-of-the-art models for a wide range of NLP tasks.
Some of the key advantages of BERT for text classification include:
- Minimal task-specific architecture changes required – Just add a classification layer on top of the core BERT model
- Reduced data needs for supervised learning – BERT is pre-trained on huge unsupervised corpora so only a small labeled dataset is needed for fine-tuning
- Better generalization and less overfitting – The model learns general language patterns that are useful across tasks
- State-of-the-art performance – BERT outperforms previous methods on many language benchmarks and real-world tasks
Now that we have an understanding of what BERT is and why it‘s well-suited for text classification, let‘s look at how we can implement it in Keras.
Implementing BERT in Keras
The easiest way to use BERT with Keras is via the TensorFlow Hub library which provides a selection of pre-trained models including BERT. We simply have to use the hub.KerasLayer() function to load the BERT layer into our neural network.
Here is the basic code to build a BERT classifier in Keras:
import tensorflow as tf import tensorflow_hub as hubdef build_model(bert_layer, max_seq_length): input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_word_ids") input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_mask") segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="segment_ids")
pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])
clf_output = sequence_output[:, 0, :] net = tf.keras.layers.Dense(64, activation=‘relu‘)(clf_output) net = tf.keras.layers.Dropout(0.2)(net) net = tf.keras.layers.Dense(32, activation=‘relu‘)(net) net = tf.keras.layers.Dropout(0.2)(net) out = tf.keras.layers.Dense(3, activation=‘softmax‘)(net)
model = tf.keras.models.Model( inputs=[input_word_ids, input_mask, segment_ids], outputs=out) model.compile(tf.keras.optimizers.Adam(lr=1e-5), loss=‘categorical_crossentropy‘, metrics=[‘accuracy‘])
return model
module_url = "https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/2" bert_layer = hub.KerasLayer(module_url, trainable=True)
model = build_model(bert_layer, max_seq_length=128) model.summary()
The pooled_output is used for classification tasks on an entire sentence whereas sequence_output contains a representation for every token which is useful for token-level tasks. We extract the first token‘s embeddings of the sequence_output for our classification head.
Two dense layers with relu activations are used with some dropout regularization. Finally, a softmax output layer is used for multi-class probability. The model is compiled with the Adam optimizer and categorical cross-entropy loss.
Alternatively, you can also build the Keras model directly from the TFBertForSequenceClassification class provided by the Hugging Face Transformers library. This simplifies model creation even further.
The other key component for BERT models is the custom tokenizer used to preprocess text into the data format expected by BERT (token ids, masks, and segment ids). This involves:
- Tokenizing the text into tokens that match BERT‘s vocabulary
- Converting the string tokens into their corresponding integer IDs
- Adding the special [CLS] and [SEP] tokens required by BERT
- Padding or truncating sequences to a fixed length
- Creating attention masks and token type IDs
The Transformers library provides easy-to-use APIs for BERT tokenization. Here‘s a self-contained example of tokenizing text for BERT:
from transformers import BertTokenizertokenizer = BertTokenizer.from_pretrained(‘bert-base-uncased‘, do_lower_case=True)
text = "This is a sample sentence to be tokenized." tokenized_text = tokenizer.tokenize(text)
print (‘Tokenized text is: ‘, tokenized_text)
input_ids = tokenizer.convert_tokens_to_ids(tokenized_text)
print (‘Tokenized ID is: ‘, input_ids)
Once we have the BERT model architecture defined and text converted into the proper data format, we can move onto actually training our multi-class text classifier.
Code Example: Multi-Class Text Classification with BERT
Let‘s walk through an end-to-end example of building a BERT-based multi-class text classifier in Keras. We‘ll use the popular 20 Newsgroups dataset which contains ~20,000 newsgroup posts across 20 topic categories.
from sklearn.datasets import fetch_20newsgroupscategories = [‘alt.atheism‘, ‘comp.graphics‘, ‘sci.med‘, ‘soc.religion.christian‘] train_data = fetch_20newsgroups(subset=‘train‘, categories=categories, shuffle=True, random_state=42) test_data = fetch_20newsgroups(subset=‘test‘, categories=categories, shuffle=True, random_state=42)
Next, we‘ll tokenize our text data using the BERT tokenizer:
from transformers import BertTokenizertokenizer = BertTokenizer.from_pretrained(‘bert-base-uncased‘, do_lower_case=True)
train_input_ids = [] train_attention_masks = []
for sent in train_data.data: encoded_dict = tokenizer.encode_plus( sent,
add_special_tokens = True, max_length = 128,
pad_to_max_length = True, return_attention_mask = True,
return_tensors = ‘tf‘ )train_input_ids.append(encoded_dict[‘input_ids‘]) train_attention_masks.append(encoded_dict[‘attention_mask‘])train_input_ids = tf.concat(train_input_ids, axis=0)
train_attention_masks = tf.concat(train_attention_masks, axis=0)
train_labels = tf.keras.utils.to_categorical(train_data.target)val_input_ids = [] val_attention_masks = []
for sent in test_data.data:
encoded_dict = tokenizer.encode_plus(
sent,
add_special_tokens = True,
max_length = 128,
pad_to_max_length = True,
return_attention_mask = True,
return_tensors = ‘tf‘
)val_input_ids.append(encoded_dict[‘input_ids‘]) val_attention_masks.append(encoded_dict[‘attention_mask‘])val_input_ids = tf.concat(val_input_ids, axis=0)
val_attention_masks = tf.concat(val_attention_masks, axis=0)
val_labels = tf.keras.utils.to_categorical(test_data.target)For each text, we tokenize, add special tokens, pad/truncate to a fixed max length, and return the attention mask. The token IDs and attention masks are stored, concatenated, and converted to TensorFlow tensors. We also one-hot encode the multi-class labels.
Now, let‘s build our BERT classifier model:
import tensorflow as tf import tensorflow_hub as hubmax_seq_length = 128 num_classes = len(categories)
input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_word_ids") input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_mask") segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="segment_ids")
module_url = "https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/2" bert_layer = hub.KerasLayer(module_url, trainable=True)
_, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])
clf_output = sequence_output[:, 0, :] net = tf.keras.layers.Dense(64, activation=‘relu‘)(clf_output) net = tf.keras.layers.Dropout(0.2)(net) net = tf.keras.layers.Dense(32, activation=‘relu‘)(net) net = tf.keras.layers.Dropout(0.2)(net) out = tf.keras.layers.Dense(num_classes, activation=‘softmax‘)(net)
model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=out) model.compile(tf.keras.optimizers.Adam(lr=1e-5), loss=‘categorical_crossentropy‘, metrics=[‘accuracy‘])
The heart of the model is the BERT layer loaded from TensorFlow Hub which is set to be trainable. We feed in the required input token IDs, attention mask, and segment IDs. The first token‘s embeddings of the Transformer‘s sequence_output is extracted for our classification head – two hidden dense layers followed by a softmax layer for the final class probabilities.
Finally, we can train and evaluate the model:
history = model.fit([train_input_ids, train_attention_masks, segment_ids], train_labels, validation_data=([val_input_ids, val_attention_masks], val_labels), epochs=3, batch_size=32)_, accuracy = model.evaluate([val_input_ids, val_attention_masks], val_labels) print(‘Accuracy: ‘, accuracy)
After just 3 epochs of training, we achieve about 96.5% classification accuracy on the held-out test set! This demonstrates the power of transfer learning with BERT where we can achieve high performance with minimal training on a small dataset.
To make predictions on new raw text, we simply need to tokenize the text, convert to the proper data format, and pass it through the trained model:
text = "This post is about computer programming and coding"encoded_dict = tokenizer.encode_plus( text,
add_special_tokens = True, max_length = 128,
pad_to_max_length = True, return_attention_mask = True,
return_tensors = ‘tf‘
)input_ids = tf.expand_dims(encoded_dict[‘input_ids‘], 0)
attention_mask = tf.expand_dims(encoded_dict[‘attention_mask‘], 0)preds = model.predict([input_ids, attention_mask]) pred_label = categories[np.argmax(preds)] print(‘Predicted label: ‘, pred_label)
Tips and Best Practices
Here are a few tips and best practices to keep in mind when fine-tuning BERT for text classification:
Choose an appropriate pre-trained BERT model for your task. Larger models trained on more data generally perform better but are also slower and more resource-intensive. The ‘bert-base-uncased‘ model is a good starting point.
Experiment with different sequence lengths. While the maximum sequence length for BERT is 512, using shorter lengths can speed up training. Look at the distribution of your text lengths and choose a length that covers most of your data.
Pay attention to your learning rates. BERT models are usually fine-tuned with lower learning rates compared to training from scratch. Adam optimizer with learning rates in the range 1e-5 to 5e-5 work well.
Consider the number and size of dense layers. If you have a small dataset, using fewer hidden layers can help prevent overfitting. Increasing dropout can also help.
Monitor validation performance and use early stopping. BERT models can overfit quickly, especially on small datasets. Use your judgment to decide the best number of training epochs.
Fine-tune the entire model, not just the final layer. While it‘s possible to freeze the weights of the BERT layer and only train the new classifier, fine-tuning the entire model typically leads to better performance.
Conclusion
In this post, we explored how to use BERT with the Keras framework for multi-class text classification. BERT is a powerful approach that can significantly boost performance on many language understanding tasks with minimal task-specific architectural changes.
The key steps are:
- Preprocess text data using the BERT tokenizer
- Create train/val splits and format data for input to BERT
- Load a pre-trained BERT model and build a classifier architecture
- Fine-tune the entire model on your dataset
By following these steps and modifying the final classification layer, BERT can be easily adapted for a variety of text classification tasks including sentiment analysis, topic labeling, intent detection, and more. Give it a try on your own datasets!
