How to Use Transformers in TensorFlow | Towards Data Science

Transformers are, without a doubt, one of the biggest advances in NLP in the past decade. They have (quite fittingly) transformed the landscape of language-based ML.
Despite this, there are no built-in implementations of transformer models in the core TensorFlow or PyTorch frameworks. To use them, you either need to apply for the relevant Ph.D. program, and we’ll see you in three years — or you pip install transformers.
Although this is simplifying the process a little — in reality, it really is incredibly easy to get up and running with some of the most cutting-edge models out there (think BERT and GPT-2).
When using Huggingface’s transformers library, we have the option of implementing it via TensorFlow or PyTorch. We will be covering everything you need to know to get started with the TensorFlow flavor in this article.
What is HuggingFace?  - Finding Models  - Visualing AttentionTensorFlow  1. Pre-Processing  2. Tokenizer and Model  3. Encoding Inputs  4. Full Model Architecture  5. Metrics, Loss, and Optimizer  6. TrainingResults
If you prefer video, I’ve covered the full build process here too:
HuggingFace is a company building and maintaining the hugely popular Transformers library. We can easily hit the ground running with the majority of the big, most cutting-edge transformer models available today through this library.
To download and begin working with some of the biggest models out there, including (but not limited to) — BERT, RoBERTa, GPT, GPT-2, XLNet, and HuggingFace’s own DistilBERT and DistilGPT-2 — it takes no more than three lines of code, which look like this:
from transformers import TFAutoModel, AutoTokenizermodel = TFAutoModel.from_pretrained("<model-name>")tokenizer = AutoTokenizer.from_pretrained("<model-name>")
Not only can we access all of these models with incredible ease, but we can even take advantage of prebuilt transformers for question and answering, sentiment analysis, text summarization, and much more.
It’s a hugely popular library, built by an incredibly talented team, democratizing some of the most powerful models at the cutting-edge of NLP.
HuggingFace has built an incredible ecosystem that provides an insanely large number of ready-to-use transformers, the full list of which we can find here. Both community-built and HuggingFace-built models are available.
We can filter for models via the Tags dropdown. We can search based on the framework they are built with, such as PyTorch/TensorFlow — it’s use-case such as classification, QnA, etc.
We will be using TensorFlow, and we can see a list of the most popular models using this filter. Uncased/cased refers to whether the model will identify a difference between lowercase and uppercase characters — which can be important in understanding text sentiment.
The bert-base-cased model fits our use-case, and we will be implementing this model (later) using the instructions provided on the model page.
As a brief but interesting side-note, HuggingFace, alongside IBM Research and Harvard NLP, have built an incredibly fascinating tool for visualizing the attention in several transformer models with exBERT lite.
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 dataInitialize the HuggingFace tokenizer and modelEncode input data to get input IDs and attention masksBuild the full model architecture (integrating the HuggingFace model)Setup optimizer, metrics, and lossTraining
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.
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:
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!
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:
Once we have our sequences, we can encode them withtokenizer.encode_plus:
Where SEQ_LEN=50, the output of the input IDs will look like this:
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:
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.
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.
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 variablesmodel.fit({"input_ids": input_ids, "attention_mask": mask},          labels=labels)(2) We place them into a tf.data.Dataset object beforehandmodel.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.
Our model summary shows the two input layers, BERT, and our final classification layers. We have a total of 108M+ parameters, of which just 100K are trainable because we froze the BERT parameters.
Here we are returning to the standard TensorFlow build process.
First, we use the optimizer we all know and love. Next, we use category cross-entry and categorical accuracy for our loss and single metric. Because we have one-hot encoded our outputs, we use Categorical.
Alternatively, we can avoid one-hot encoding our outputs and use losses.SparseCategoricalCrossentropy(from_logits=True) with metrics.SparseCategoricalAccuracy('accuracy').
Finally, we compile our model with the .compile method — we are now ready to begin training!
Note: If training BERT layers too, try Adam optimizer with weight decay — which can help reduce overfitting and improve generalization [1]. I would recommend this article for understanding why.
import tensorflow_addons as tfa  # Adam with weight decayoptimizer = tfa.optimizers.AdamW(0.005, learning_rate=0.01)
We train as per usual using the fit method. If not using a tf.data.Dataset object we must explicitly state our multiple inputs using a dictionary. All this requires is:
Training will take a long time with most GPUs —for those of you on CPU only, good luck!
If you are stuck on CPU, try out Google Colab — it’s a free, cloud-based notebook service provided by Google. Colab includes a GPU as standard — albeit not a particularly powerful one (but it is free).
Here is a full version of the code:
After several hours training on my measly 940MX we get a validation accuracy of 94.9% — with more training epochs, the model looks like will easily improve on this too:
These are pretty great results for such a simple output network. Further fine-tuning, the addition of CNNs, LSTMs, or other more expressive networks may improve our results even further.
Alternatively, (although I found this to be detrimental) we can even use BERTs pre-pooled output tensors by swapping out last_hidden_state with pooler_output — but that is for another time.
I hope you’ve enjoyed this article on integrating TF2 and HuggingFace’s transformers library. If you’d like more, I post programming/ML tutorials on YouTube here!
If you have any questions, let me know via Twitter or in the comments below.
Thanks for reading!
[1] I. Loshchilov, F. Hutter, Decoupled Weight Decay Regularization (2019), ICLR
🤖 NLP With Transformers Course
*All images are by the author except where stated otherwise