Deploying a NLP model with Docker and FastAPI
As a data scientist, my work usually consists on developing machine learning or deep learning models for a huge variety of tasks. I often start with a Jupyter notebook doing exploration and experimentation with the data. When I am clear about the model that I am going to use, I create some deploy scripts. Normally, these scripts consist of APIs and Docker images to allow fast deployments in the cloud or on-premise systems.
In this article, I am going to provide a simple way to develop and deploy an API for an NLP task. In this case, I will create a model for the Kaggle competition: https://www.kaggle.com/uciml/sms-spam-collection-dataset/. This task consists on classifying a series of SMS messages as spam or not spam (ham). The main problem of this dataset is that the samples are quite unbalanced. We have 4825 samples that are not spam and 747 samples that are spam.
Building the model ⚒
To build the model, I will create a simple Pipeline in scikit-learn where the text is first tokenized using NLTK, then I will create a TF-IDF vectorizer to represent the textual information. Finally, I will train a support vector machine with a bag of words weighted by TF-IDF. For each step, I will explore some hyperparameters using a grid-search method and cross-validation evaluation by 10 folds. As the classes are unbalanced, I will use F1-macro to evaluate the models. The notebook with the exploration is in https://github.com/marescas/spam_ml/blob/main/notebooks/spam.ipynb.
When the model is trained, evaluated and the results are good (98% F1-macro in the test set) the next step is to save the model to allow your production script to use it. For doing that we can use the joblib library.
Building the API ⚙️
Now we have the model trained and saved. It’s time to build our API to serve it on production. For doing that I will use FastAPI. FastAPI is a library built on top of Pydantic, Starlette, and Swagger, and allows developers to build fast and robust APIs in Python. Thanks to Pydantic, check data types are now easier than never. And also, thanks to Swagger it is possible to build automatic interactive documentation for the API.
The API that I developed only has an endpoint to predict the hypothesis class given a new SMS. As we can see in the following snippet the code it is quite simple 😁.
Basically, you need to load the model and then define the endpoint with the inference logic. In this case, I want to define a POST method with a param called “data” as a string. Thanks to strong typification, Pydantic can check if the input data is a string and if not, it can return an exception.
Also, as I say, thanks to the integration between fastAPI and Swagger we can have interactive documentation generated automatically. This is very useful when you need to explain to other people how your API works.
Deploying with Docker 🏎
Now, we have a trained model that works well to detect if an SMS is spam or not. Also, we have an API developed with FastAPI that allows us to serve this model. We need one more thing, a way to encapsulate the code and port the application to different platforms. For that, I will use Docker and Docker-compose. Basically, Docker allows to encapsulate the code in a container and port the application in an easy way.
As we can see in the previous script, our API runs with uvicorn. Uvicorn is an ASGI server implementation, using uvloop and httptools. Thanks to uvicorn we can load 4 workers for the same API.
Also, thanks to Docker-compose we can define in a declarative way how our applications run. For this example, the definition is quite simple and we only map host port 8080 to docker port 80. Using docker-compose up we can launch the API 😁. This API will be running on port 8080…
Conclusions 📚
In this post, we learned how to build and deploy a simple API to predict if an SMS is spam or not. For building this API, we use some tools like Scikit-Learn, FastAPI, and Docker. As I say, this is only a prototype, and it is possible to do it using other technologies like MLflow, but this will be for another post… I hope you found this post helpful and thanks for reading. The code is available at the following link:
