Classify text with BERT

View on TensorFlow.org Run in Google Colab View on GitHub Download notebook See TF Hub model

This tutorial contains complete code to fine-tune BERT to perform sentiment analysis on a dataset of plain-text IMDB movie reviews. In addition to training a model, you will learn how to preprocess text into an appropriate format.

In this notebook, you will:

  • Load the IMDB dataset
  • Load a BERT model from TensorFlow Hub
  • Build your own model by combining BERT with a classifier
  • Train your own model, fine-tuning BERT as part of that
  • Save your model and use it to classify sentences

If you're new to working with the IMDB dataset, please see Basic text classification for more details.

About BERT

BERT and other Transformer encoder architectures have been wildly successful on a variety of tasks in NLP (natural language processing). They compute vector-space representations of natural language that are suitable for use in deep learning models. The BERT family of models uses the Transformer encoder architecture to process each token of input text in the full context of all tokens before and after, hence the name: Bidirectional Encoder Representations from Transformers.

BERT models are usually pre-trained on a large corpus of text, then fine-tuned for specific tasks.

Setup

# A dependency of the preprocessing for BERT inputs
pip install -U "tensorflow-text==2.13.*"

You will use the AdamW optimizer from tensorflow/models.