Use pre-trained Huggingface models in TensorFlow Serving | by Luis Duarte | Towards Data Science

HuggingFace simplifies NLP to the point that with a few lines of code you have a complete pipeline capable to perform tasks from sentiment analysis to text generation. Being a Hub for pre-trained models and with its open-source framework Transformers, a lot of the hard work that we used to do is simplified. This allows us to write applications capable of solving complex NLP tasks, but with the downside that we don’t know what is happening behind the curtains. Despite this amazing simplification that HuggingFace and Transformers perform we might want some abstraction from all that code and just simply use one of many available pre-trained models. In this post, we will learn how we can use one of many pre-trained models with TensorFlow Serving, a popular service to put machine learning models into production.
I will use this Distilbert pre-trained model for Sentiment Analysis which will predict if a given text is positive or negative. Unfortunately, Transformers doesn’t have a direct export functionality to TensorFlow Serve so we will have to do a bit of working around to achieve our goals. First, we need to install Tensorflow, Transformers and NumPy libraries.
pip install transformerspip install tensorflowpip install numpy
In this first section of code, we will load both the model and the tokenizer from Transformers and then save it on disk with the correct format to use in TensorFlow Serve.
from transformers import TFAutoModelForSequenceClassificationimport tensorflow as tfMAX_SEQ_LEN = 100model = TFAutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")callable = tf.function(model.call)concrete_function = callable.get_concrete_function([tf.TensorSpec([None, MAX_SEQ_LEN], tf.int32, name="input_ids"), tf.TensorSpec([None, MAX_SEQ_LEN], tf.int32, name="attention_mask")])model.save('saved_model/distilbert/1', signatures=concrete_function)
First of all, we load the model with TFAutoModelForSequenceClassification. One crucial aspect in loading transformers models is choosing the right Class. Since we are using a pre-trained model for Sentiment Analysis we will use the loader for TensorFlow (that’s why we import the TFAutoModel class) for Sequence Classification. If you are unsure what Class to load just check the model card or “Use in transformers” info on Huggingface model page for which class to use.
In the code snippet, we declared a signature function that is necessary to use with TensorFlow Serve. This function is necessary for declaring the input shape of our data to the servable model in TF Serve, which is defined by the variable MAX_SEQ_LEN. In this case, I have defined that the model will accept two inputs, two lists with a size of 200, for the input_ids (the ids that we obtain from the tokenizer) and the attention_mask (used if the input sequence length is smaller than the max input sequence length). It is also important to remember the name we declare for each input because it’s the parameter we will later define in the HTTP request to send to our model. After executing these few lines we should have a new directory with the following files on our working directory:
saved_models

distilbert

1

assets
saved_model.pb
variables

variables.data-00000-of-00001
variables.index








distilbert

1

assets
saved_model.pb
variables

variables.data-00000-of-00001
variables.index






1

assets
saved_model.pb
variables

variables.data-00000-of-00001
variables.index




assets
saved_model.pb
variables

variables.data-00000-of-00001
variables.index


variables.data-00000-of-00001
variables.index
If you haven’t installed TensorFlow Serve (I recommend using it with Docker) do it now. After we have installed we can pull the image from Docker and start running our service with the Distilbert model loaded with the following command:
docker run -p 8501:8501 --mount type=bind, source=/PATH_TO_YOUR_DIRECTORY/saved_models/distilbert, target=/models/distilbert -e MODEL_NAME=distilbert -t tensorflow/serving
This will cause TensorFlow Serve to load Distilbert (The model files must be under a directory named with a number or else TF Serve won’t load the model) in port 8501 which we can request by HTTP. In the next section, I will show how we can make predictions for our loaded model.
import requestsfrom transformers import AutoTokenizer, AutoConfigtokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")text = "I like you. I love you"encoded_input = tokenizer(text, pad_to_max_length=MAX_SEQ_LEN, max_length=MAX_SEQ_LEN)#TF Serve endpointurl = "http://localhost:8501/v1/models/distilbert:predict"payload={"instances": [{"input_ids": encoded_input['input_ids'], "attention_mask": encoded_input['attention_mask']}]}print(payload)>> { "input_ids": [101, 1045, 2066, 2017, 1012, 1045, 2293, 2017,  102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "attention_mask": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] }headers = {  'Content-Type': 'application/json'}response = requests.request("POST", url, headers=headers, data=json.dumps(payload))print(json.loads(response.text)['predictions'])>>[[-4.2159214, 4.58923769]]
We start by calling the model tokenizer from transformers to convert our input text into a list of ids and its corresponding attention mask, both with padding, to obtain the necessary format to send to our loaded model. After tokenizing we perform a POST request to TensorFlow Serving with the correct payload where we obtain the presented values. But wait, isn’t the model suppose to classify the sentiment in our text? Well, we are missing one last step, applying the softmax function to our result vector.
import numpy as npdef softmax(x):    e_x = np.exp(x - np.max(x))    return e_x / e_x.sum(axis=0)print(softmax(json.loads(response.text)['predictions'][0]))>>[0.0001 0.9999]
The resulting vector is the probability of the input text to be negative or positive. In this case, a probability of 0.0001 to be negative and 0.9999 to belong to the positive class. If we compare the result with the output from loading the pipeline with transformers we can see both are the same.
Concluding this post we have used one of many huggingface pre-trained models and loaded it up on TensorFlow Serve so we can make HTTP requests and easily scale these models into production. We also learned how we can make requests to our loaded model with the model tokenizer and how we need to apply the softmax function to our results to obtain the desired values. In my next post, I will show how we can abstract ourselves from the transformers library completely by writing our tokenizer and use it to make requests to our loaded models. Like always feel free to comment with suggestions or improvements and hope that this post helped you achieve your goals.