Fast.ai in production. Real-word text classification with ULMFiT.

I’ve overcome my skepticism about fast.ai for production and trained a text classification system in non-English language, small dataset and lots of classes with ULMFiT.

About the project

My friend and classmate, who is one of the founders of RocketBank (leading online-only bank in Russia), asked me to develop a classifier to help first-line of customer support.

The initial set of constraints was pretty restricted —

no labeled historical data
obfuscated personally identifiable information
“we want it yesterday”
mostly Russian language (hard to find pre-trained model)
no access to cloud solutions due to privacy regulations
ability to retrain the model if new classes arise without my involvement

The scope of work was pretty straight forward — develop a model and serving solution for incoming messages classification into 25 classes.

Initially, after thinking about restrictions — I was pretty sure, that no neural networks should be used for this approach. Why?

We would never label enough data for the neural model in reasonable time.
Building an environment for the reliable serving of neural model is a kind of pain.
I was skeptical about reaching the requested performance (requests per second) with reasonable resources.

Based on that, I made a sad face and created new conda environment_._ I was classic.

Dataset collection and initial research.

RocketBank had set up a task force consisting of a project manager and devops on their side plus a bunch of people handling dataset labeling. It was extremely smart and helpful, as, in my opinion, this constitutes a perfect team for handling a data science project in the industrial world.

We started with analyzing historical data and came to a number of conclusions:

To train a system we take into account only messages received by the bank before any response from customer support.
There are 2 distinct meta-classes of incoming messages — coming from existing customers and from new leads. Adding this information as an input to classifier should provide additional information to the system and boost classification scores.
Bank on their side decided on 25 distinct classes of messages ranging from ‘Credit request’ up to ‘Harassment’.

I requested around 25000 not labeled historical messages and in just a few days a task force was able to classify around 1500 historical messages into 25 classes. Initially, I assumed that this number (1.5k) is too low to even try any neural model (I was wrong).

TD-IDF + TPOT and DEAP

I will fast forward through non-interesting part of the thing. I decided to test various flavors of TD-IDF, embeddings and optimize machine learning model using TPOT and DEAP.

TPOT and DEAP, for those unaware, are two secret weapons in data scientist arsenal that make model search CPU-intensive and hands-free.

TPOT runs all stack of machine learning methods embedded in sklearn plus few extra(XGBoost) and finds the optimal pipeline. I played around various embeddings, fed them into TPOT and after 24 hours said that Random Forest performs best for my model (ha-ha, what a surprise!).

Then I needed to find an optimal set of hyperparameters. I always do this using directed evolution strategies with DEAP library. That actually deserves a separate post.

Anyway, at the end of the day, I received an optimal set of settings and my precision was around 63%. I think this was close to the maximum that I was able to get from classical methods and 1.5k dataset. While 63% for 25 classes sounds good from the machine learning perspective, it's quite bad for real-world usage. So, I decided to take a look into neural nets as a last chance.

Fast.ai comes into play.

So, I needed a fast way to check the performance of a neural-based model on the same task. While implementing a model from scratch using Tensorflow was the most viable option, I decided to run a fast test with fast.ai and their recent discovery of ULMFiT. Problem is — I needed a pretrained language model for Russian text, which isn’t available in fast.ai. After looking at fast.ai forums I discovered an ongoing effort to create a set of a language model for most languages. There was a thread for Russian language and a pre-trained model from a Russian Kaggler Pavel Pleskov, that he used to get a second place at Yandex competition.

From there it was mostly writing 20 lines of code and few hours of GPU training time to get to 70% precision. After a few more days of tuning hyperparameters, I get to 80% precision. Some tips:

Use FocalLoss as a training goal.
Have a pretrained language model, but finetune it on a non-labeled data available.
Convert text to lowercase and make a token for uppercase, make a special token for pieces of obfuscated data.
Put token of meta-class not only in the beginning but also at the end of the message.

Deployment

Ok, great. Should I convert the model into TensorFlow? Nope, I was lazy and decided to test model performance using native Fast.ai + Pytorch + Docker.

After running stress-tests in a single-core Docker container I was surprised to see less than 300 milliseconds response time for an average request and no crashes. What else did I need? Nothing.

Fast.ai showed that it is a perfect solution for fast and precise development of production ML systems.

Retraining

The beauty of fast.ai + transfer learning is a pretty predictable result of retraining in terms of quality and speed. I’ve shared a script inside the docker container coping my final training notebook and providing a new model as an asset. I’ve run a few cycles of retraining and cross-validation and obtained highly repeatable results, so this is a simple way to deliver not only a model but a training script as well.

I can’t share the actual code and system configs, but I am ready to answer any questions.