Implementing a Custom Intent Classification Model with Rasa

This post was originally written by Anna Vechkaeva

Rasa is a framework for developing conversational AI agents, which offers a number of off-the-shelf models which can be trained on custom data. The models range from simple, e.g., keyword-based for intent classification and rule-based models for dialogue management, to more complex transformer-based ones like DIET (Dual Intent and Entity Transformer) and TED (Transformer Embedding Dialogue). The Rasa framework is quite flexible and allows us to easily implement conversational pipelines with various configurations using these pre-built models.

Apart from mixing and matching existing models, Rasa also offers the ability to use your own custom models, making the framework even more flexible. In this blog post, you will learn how to build your own intent classification model compatible with Rasa with the example of a simple SVM classifier implemented using scikit-learn. If you want to follow this tutorial you can find the code in this repository https://github.com/MantisAI/rasa_custom_intent_classification.

Note that there is a Rasa tutorial on how to implement a custom component available here. This post is more specific to implementing the Intent Classification component, and will go into more detail about the data structures used by Rasa. Note that the code shared in this post is specific to Rasa 2.8.0.

First, we need to define the model class, which needs to inherit from the IntentClassifier Rasa class, and can be imported as follows:

The intent classification class should have the following attributes:

• name: the name of the component
• provides: what output it provides e.g. intent, entities, etc.
• requires: what input it requires e.g. text, tokens, etc.
• defaults: default parameters for the model language_list — the list of languages for which it can be used

In our example class named CustomIntentSVM the parameters are defined as follows:

Model initialization. The custom intent classification class is initialized in the __init__() method, which loads the model parameters from the config.yml.

The intent classification class should contain the following methods, which will be used by Rasa to train and use the model:

• train()
• process()
• persist()
• load()

train() method. This method will be called by Rasa to train the model using data specified in nlu.yml and is called when rasa train command is executed. As an input, it gets training data (train_data) and configuration parameters (cfg) specified in the config.yml. The train method does not return anything and updates the model, which in our example is stored in an attribute model of our custom class.

The first thing that the train method does is to define the model using the parameters from cfg.

These parameters are specified in the config.yml as follows.

Then it needs to transform train_data to the input that the sklearn expects, which is a list of texts and a list of labels (intents). Those can be retrieved by iterating the messages through the training_examples attribute of train_data. Each message has the text and label stored as an attribute in a dictionary, with text and intent keys

Finally it can call model.fit with documents, labels as the arguments.

process() method. This method is used to get the prediction out of the model and record it. It is called when the rasa shell command is executed, and Rasa is passed a message. As an input, it gets an argument message, which has only the text key in its data attribute, but no intent. This method does not need to return anything, it just updates the message passed as the input, adding the predicted intent and its confidence score to the dictionary stored in the message.data. The prediction added to the dictionary should be in the specific Rasa format: {“name”: value, “confidence”: confidence}.

We add a helper function _convert_to_rasa to convert the predictions and confidences produced by our sklearn model into Rasa format and _predict to produce the predictions and the confidences.

persist() method. This method saves the trained model as a .pkl, so that it can be loaded for prediction when interacting with Rasa. As the input, it gets the following attributes: file_name, and model_dir, which contain the name of the file and the directory where it needs to be saved. This method needs to return the model’s file name in the form of the dictionary: {“file”: model_file}, so that it later can be retrieved.

load() method. This method loads the model saved by the persist method. In order to load the model, we need to get its directory and the file name. This information is passed in the meta and and model_dir arguments.

Once you define these files and put your data in the nlu.yml, you can train your custom model using rasa train nlu, test it by using rasa test, and interact with the nlu part of it using rasa shell. Any other model can be easily implemented by editing _define_model, _transform_data, and _predict methods. In the accompanying repository we also share an example of a transformer based classifier.

Below you can find the full code for the implemented model as well as the config.yml file which is used to run it with Rasa.