Loading NLP HuggingFace models into AllenNLP framework

How to take advantage of the transformers library in HuggingFace and extend its functionality using AllenNLP

If we would have to name two libraries in the NLP world that contain cutting-edge models architecture implementations we will probably name transformers by HuggingFace and AllenNLP from AllenAI.

I recently come across a scenario where I wanted to load a model constructed in transformers to the AllenNLP library to use its interpretability capabilities (Model interpretability — Making your model confesses: Saliency maps). In this post, I will go ever the steps to make it happen!

Introduction

Everyone loves HuggingFace and their transformers library. They have been on a mission to advance and democratize artificial intelligence through open source and open science. With ~50,000 pre-trained machine-learning (last time I checked) models and 5,000 datasets currently hosted on the platform, HuggingFace enables the community to build their own Machine Learning capabilities, share their own models, and more.

AllenNLP is a general deep learning framework for NLP, established by the world-famous Allen Institute for AI Lab. Its team envisions language-centered AI that equitably serves humanity. As a small team, researchers and engineers work closely together to publish impactful research and identify common investments that would benefit the field as a whole.

At this point, you may wonder: which one is the best one?

The transformers library

🤗 Transformers provides APIs to easily download and train state-of-the-art pretrained models. Using pretrained models can reduce your compute costs, consume cutting-edge model architectures faster, and save you time from training a model from scratch. The models can be used across different modalities such as text, images, audio and multimodal!

The library supports seamless integration between three of the most popular deep learning libraries: PyTorch, TensorFlow, and JAX.

AllenNLP library

AllenNLP is a general deep learning framework for NLP. It contains state-of-the-art reference models running on top of PyTorch. AllenNLP is a library that also seeks to implement abstractions that allow rapid model development and component reuse by detaching from the implementation details of each model. It is one of the key differentiators from others: its modular approach allows you to quickly test different variations of architectures by switching components.

Why choose?

Each library comes in handy at specific scenarios. In my own opinion, I find it extremely easy to get started with the transformers library. It’s easy to use, they have a clear interface to operate it and the community is huge.

AllenNLP, on the other hand, is great for composing new or existing architectures. Its modular approach allows you to quickly design architectures and try new combinations by just switching such components. They also have great tools for interpretability and really cutting-edge techniques get implemented faster in the library.

Considering that transformers and AllenNLP support PyTorch backend, we may consider loading models from one platform into the other. In fact, it is possible. In this post, I will should you how you can download a model from HuggingFace (or your own custom model built with transformers) and load it into AllenNLP.

Getting started

Getting a model from HuggingFace

First, let’s get a model from HuggingFace to port to AllenNLP. In this example, I will use [nlptown/bert-base-multilingual-uncased-sentiment](https://huggingface.co/nlptown/bert-base-multilingual-uncased-sentiment). This is a bert-base-multilingual-uncased model finetuned for sentiment analysis on product reviews (from Amazon) in six languages: English, Dutch, German, French, Spanish, and Italian. It predicts the sentiment of the review as a number of stars (between 1 and 5).

The model can be used directly as a sentiment analysis model for product reviews in any of the six languages, or further finetuned on related sentiment analysis tasks. To keep the example small, we won’t do any fine-tuning with our own data in this opportunity.

In transformers, we can load a model using the Auto version of tokenizers and models.

[embed]

That was easy, right? That’s all we need to get a model working in transformers and which is why I said this library makes getting starting so easy. Now, let's try to load this model into AllenNP.

Saving the components locally

It’s good to save the model locally so we can work with AllenNLP directly from our local file system. We can save a HuggingFace model into our local environment like this:

[embed]

Loading the model into AllenNLP

As we saw before, AllenNLP has a modular approach for constructing models. Let’s go over each of these components for the case of an NLP classification model based on transformers.

Tip: Don’t get scared for the level of granularity AllenNLP has compared to what is needed in transformers. AllenNLP will let you choose everything. At the end you will see a more compact way of loading a model, but we need to start simple.

The vocabulary

As with any other framework, a vocabulary in AllenNLP maps strings to integers. They are fit to a particular dataset, which is used to decide which tokens are in-vocabulary. Any token that is outside of the vocabulary is mapped to a particular token called out-of-vocabulary.

An important distinction is that in AllenNLP, vocabularies can have different namespaces so you can have separate indices for the same token. For instance, you can have one vocabulary for your words as inputs and another one for your words as outputs (in a text generation setting for instance).

[embed]

The tokenizer

The tokenization work is divided into 2 parts in AllenNLP, which allows having a more modular approach:

• Tokenizer: A Tokenizer splits chunks of text into tokens. Typically, this either splits text into word tokens or character tokens. Its job is to split sequences of text into sequences of discreet words or tokens. It goes from text into sequences of text.
• Indexer: its job is to take a sequence of tokens and translate them into word indexes in according to the vocabulary. It goes from sequences of text into sequences of indexes in the vocabulary.

[embed]

The embedder

The embedder's job is to provide vectors for each word index. Most NLP models utilize this kind of dense representation instead of indices. It basically takes a word’s index and returns its vector representation.

AllenNLP supports providing multiple embedders for the same or different inputs. In our case, we are going to provide embeddings for inputs provided in a field called “tokens”. (we will see what a field is later).

[embed]

Note that the type of our Embedder is BasicTextEmbedder. However, AllenNLP has different types of embedder depending on what you want to do.

The encoder

Sometimes your model doesn’t take your embeddings or word vectors directly as inputs. Instead, they consume a single multidimensional vector with some transformation. This is for instance the case of the BERT architecture.

Encoders have this job. There are multiple types of encoders, including sequence-to-sequence, sequence-to-vector, etc. The model we are going to use here (which is based on BERT) uses a sequence-to-vector encoder or Seq2VecEncoder.

[embed]

Details: A Seq2VecEncoder is a module that takes as input a sequence of vectors and returns a single vector. The input shape would be (batch_size, sequence_length, input_dim) and return a (batch_size, output_dim) tensor. In the BERT architecture, there is a pooling layer at the end of the BERT model. This returns an embedding for the [CLS] token, after passing it through a non-linear tanh activation; the non-linear layer is also part of the BERT model.

The model

So far we have taken into account the top of the BERT architecture. As you already realized, AllenNLP gives you a lot more control over which components you use and how they connect with each other. With great control comes great responsibility, said Uncle Ben.

Now it’s time to load the model itself. Since our model is a classifier, I will create a model of type BasicClassifier .

[embed]

Note here that I’m passing all the elements we constructed before, plus:

• dropout , which is the dropout probability of the final classifier. We used the default value in BERT architecture.
• The number of labels of the final classifier: In this case, 5 categories corresponding to the 5 stars in the rating system.

However, if you pay attention, we didn’t indicate model_name as we usually do in the previous cases. It should give you the intuition that no weights have been loaded for this specific classifier.

For those cases where we want to fine-tune the model to the specific task at hand, then, this is expected cause these are the weights you want to learn. However, in our case, our model is already trained in a downstream task and we do have the weights of the classifier. Let’s load them:

[embed]

Note that we switch the model to evaluation mode. This is typically done to freeze the dropout layers when you are doing inference.

Note to the reader: We are using here a private member to assign the weights, _classification_layer . Take into account that private method are not guarantee to continue working.

The DataReader

If you take a closer look at how the model has been constructed, then you will notice that tokenizer and token_indexer have not been indicated. Why is that? In AllenNLP those tasks are kept separated from the model (in the same way they are separated in transformers) in a module called DataReader.

You can define a data reader like follows:

[embed]

Note that in the same way that you have to apply multiple embedders, the same applies to token indexers. In this case, I’m specifying only one for the key “tokens”.

Putting all the pieces together

Now it’s time to run the model and see if it works! Let’s start with simple text:

text_sample = "this is a great read everyone should have"

AllenNLP allows us to put everything in a single object for convenience. This is similar to the transformer’s pipeline object. They are of great convenience for inference. Here we are placing our model and dataset reader in a single TextClassifierPredictor .

from allennlp.predictors import TextClassifierPredictor

predictor = TextClassifierPredictor(model, dataset_reader)

Let’s call predict!

We managed to import all the elements from the HuggingFace transformer into AllenNLP! The great about it is that you can now take each of them for what you feel they are better at.

Bonus track: Using JSONNET

AllenNLP allows a declarative way of loading models. The entire architecture can be specified using the language JSONNET. This allows even faster iteration for trying new combinations of different architectures as it only takes changing the relevant part in the declarative JSON.

The following JSONNET (which is indicated using the type Params but you can do exactly the same saving the JSON structure in a JSONNET file) is the very same equivalent to the model we used before.

[embed]

Notice how now predictor has been instantiated from the method Predictor.from_params() which will go over all the specifications and instantiate all the required elements for you. This turns out to be a very compact way to define the model’s architectures!

Conclusion

If you are familiar with HuggingFace but not so much with AllenNLP, I encouraged you to have a look at their website and see all the cool projects they work on and the capabilities the library has. Now, you don’t have excuses to use them! If you want a real case where this is handy check this post:

[embed]Model interpretability — Making your model confesses: Saliency maps How to interpret NLP models using Saliency Maps: an example with transformers.santiagof.medium.com