TensorFlow and Transformers
By: James Briggs
Tags: Machine Learning, Getting Started, Editors Pick, Data Science, Programming

TensorFlow

TensorFlow support in the transformers library came later than that for PyTorch, meaning the majority of articles you read on the topic will show you how to integrate HuggingFace and PyTorch — but not TensorFlow.

Of-course, the steps are slightly different — but at a high-level, the process is the same:

Pre-process the data Initialize the HuggingFace tokenizer and model Encode input data to get input IDs and attention masks Build the full model architecture (integrating the HuggingFace model) Setup optimizer, metrics, and loss Training

We will cover each of these steps — but focusing primarily on steps 2–4.

1. Pre-Processing

First, we need to prepare our data for our transformer model. What is done here will depend very much on the data and use-case. We will be taking some clean data and processing it with BERT as a classifier for sentiment analysis.

We can use the IMDB movie review dataset, which provides us with sentiment ratings from 0 (terrible) to 4 (amazing).

You can get the dataset here — or via the Kaggle API:

Because there are a lot of sentence fragments, these can easily pollute the validation set with near-matches to that in the training set. So, I removed them using drop_duplicates , keeping the first record of each unique SentenceId (the full review, meaning we drop all review segments).

All we need to do at this point is one-hot encode our sentiment labels, like so:

One-hot encoding. Before (top) and after (bottom).

2. Tokenizer and Model

We will be using the HuggingFace transformers library to source our transformer models. A smaller transformer model available to us is DistilBERT — a smaller version of BERT with ~40% of the parameters while maintaining ~95% of the accuracy.

DistilBERT is a good option for anyone working with less compute. Just switch out bert-base-cased for distilbert-base-cased below. We initialize the BERT tokenizer and model like so:

It really is as easy as that!

3. Encoding Inputs

We have our input data and tokenizer ready, so now we can encode our input data into two arrays — (1) the input IDs and (2) the attention mask.

Input IDs are simply a set of integers that represent a word, “hello” could be 0, “world” might be 1. But, BERT uses a predefined set of mappings — hence why we loaded our tokenizer using the .from_pretrained method.

There are several special tokens that are used by BERT. These are:

Table of all tokens used by BERT models.

Once we have our sequences, we can encode them with tokenizer.encode_plus :

Where SEQ_LEN=50 , the output of the input IDs will look like this:

Input ID array. This sentence only contained a single word as we can see the [SEP] token (101), a word (5650), and the [CLS] token (102). Following this are 47 [PAD] tokens.

Next up is the attention mask. This mask is simply an array of 0s and 1s where each 1 represents a valid word/input ID, and a 0 represents padding. The encode_plus method outputs both input IDs and the attention mask tensors inside a dictionary:

The encode_plus output. We are given a dictionary with output encodings for a single sentence. Included are the input IDs (first) and the attention mask (second). Note that the attention mask tells us to focus on the first three tokens only, ignoring the remaining padding tokens.

BERT's attention layers consume this mask and apply attention operations to word embedding that corresponds to a 1 while ignoring those matching up with a 0.

The effect of this is that we only apply attention to real words, while additional padding tokens are ignored.

4. Full Model Architecture

We have our encoded inputs IDs and attention masks, and the initialized BERT model — now, we need to add the additional layers required for inputting the input ID and attention mask arrays and the layers required for classifying the BERT output into sentiment ratings.

The flow of information through the BERT classifier model. We have two inputs, input_ids and attention_mask, which feed into BERT. BERT outputs two tensors — of which we use the last_hidden_state tensor and discard the pooler_output tensor.

For every BERT-based transformer model, we need two input layers that match our sequence length. We encoded our inputs to a length of 50 tokens — so we use an input shape of (50,) here:

Both our input IDs and attention mask arrays contain integers only, so we specify dtype='int32' . More importantly, the layer names must match the key-value pairs in our data inputs, for which we have two options:

(1) We keep input and output arrays as seperate variables model.fit({"input_ids": input_ids, "attention_mask": mask},

labels=labels)

(2) We place them into a (2) We place them into a tf.data.Dataset object beforehand model.fit(dataset)

We won’t cover option (2) any further than this here — but if you’re working with a lot of data, it’s a much better option (which you can read about here).

Next up are our classification layers. These will take the output from our BERT model and produce one of our three sentiment labels — there are a lot of ways to do this, but we will keep it simple:

Here we pull the outputs from distilbert and use a MaxPooling layer to convert the tensor from 3D to 2D — alternatively, use a 3D network (like convolutional or recurrent neural nets) followed by MaxPooling.

When we put all of this together we get the following:

If you are wanting to train the transformer parameters further, the final line is not necessary! We choose not to as BERT is already an incredibly well built and fine-tuned model. It would take a very long time to train, so for the likely minuscule performance increase — there’s little justification.