Start Here. / GUIDES / Build a Machine Learning Plugin

Build a Machine Learning Plugin

Overview

This guide outlines the full process of building a state-of-the-art sentiment classifier for the Memri app. We will tackle the following steps:

  • Creating a plugin from a template
  • Loading a pretrained Transformer pipeline from 🤗 Hugging Face
  • Testing, configuring, and deploying your project, all made very easy using the plugin template

Creating a plugin

The Pymemri template module offers a way to create a project from a template, with everything around setup, testing, and CI preconfigured. We will use a plugin template to build a classifier plugin in this guide, which enables us to write an app for Memri in just a few lines of code!

First things first, lets install pymemri:

pip install pymemri

On the Memri Gitlab, create a blank public repository for your plugin and clone the repo. From within the new repo, run the pymemri plugin template CLI.

plugin_from_template --template="classifier_plugin"  --description="A transformer based sentiment analyis plugin"

This creates the following folder structure:

├── setup.cfg                           
├── setup.py                            
├── Dockerfile                          <- Dockerfile for your plugin, which builds the plugin docker image
├── metadata.json                       <- Metadata for your plugin, the Pymemri frontend uses this during installation
├── .gitignore                          
├── .gitlab-ci.yml                      <- CI/CD for your plugin, which 1) installs your plugin and pod 2) runs tests 3) deploys your plugin
├── sentiment_plugin                    <- Source code for your plugin
│   ├── model.py                        <- Model definition of your classifier
│   ├── plugin.py                       <- Plugin class
│   ├── schema.py                       <- The schema definition for your plugin
│   └── utils.py                        <- Utility functions for plugins: converting items, converting photos, etc.
├── tests                               <- Tests for your plugin 
│   └── test_plugin.py
├── tools                               
│    └── preload.py                     <- You can define logic here that downloads models and assets required to run your plugin

The resulting plugin in sentiment_plugin/plugin.py is the entrypoint of your project. The Plugin.run method is called when the pod runs your plugin. In most cases, you do not have to edit this file as everything is set up correctly by the template. The model used by this plugin is defined in model.py, which we will edit in the next step to define our sentiment classifier.

Building a 🤗 Hugging Face sentiment classifier

We are building a sentiment analysis plugin that could be used by users from different countries, which are owning data in different languages. In this guide we are not training the model from scatch. Hugging Face Models has many suitable pretrained models for this task. A quick search yields the following RoBERTa model (twitter-xlm-roberta-base-sentiment). For more information about this model, read this blogpost (or the original paper). For this model, we need a few extra requirements, we add these to setup.cfg:

install_requires = 
    pymemri==0.0.8
    pytest
    torch==1.10.0
    transformers
    sentencepiece
    protobuf

With our requirements defined, let’s install the plugin locally

pip install -e .

The model card contains all the code we need to build a functioning sentiment analysis plugin. We are slightly modifying the standard example to deal with messages that are longer than the default max length. We insert this code into the model template in sentiment_plugin/model.py:

from typing import List, Any

"""Add this line"""
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

class Model:
    def __init__(self, name: str = "cardiffnlp/twitter-xlm-roberta-base-sentiment", version: str = None):
        self.name = name
        self.version = version

        """Add these lines"""
        model = AutoModelForSequenceClassification.from_pretrained(self.name)
        tokenizer = AutoTokenizer.from_pretrained(self.name, model_max_length=512)
        self.pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer, return_all_scores=True, truncation=True)

    def predict(self, x: List[str]) -> List[dict]:

        """Add this line"""
        return self.pipeline(x)

This plugin template assumes a specific output format for Model.predict, which is documented in sentiment_plugin/model.py. If you want to use a different model, make sure that the output format is the same.

Test your plugin locally

To test your plugin, its needs to 1) get data from the pod 2) make predictions on that data 3) write data back to the pod. Step 3 is handled by the template, and we just implemented step 2. We need to manually define step 1.

In tests/test_plugin.py, a simple pytest setup is defined. To implement step 1 we write the create_dummy_data method in tests/test_plugin.py, which adds data for our tests to the Pod and returns a query that retrieves said data. In tests/test_plugin.py, we then run the plugin on this data and verify the output.

def create_dummy_data(client: PodClient) -> dict:
    client.add_to_schema(Message)
    
    # Add multilingual data
    sample_data = [
        "this is great",
        "this is awful",
        "c'est incroyable",
        "c'est horrible"
    ]

    client.bulk_action(
        create_items = [Message(content=sample, service="sentiment_test") for sample in sample_data]
    )
    
    # Query to retrieve dummy data
    search_query = {"type": "Message", "service": "sentiment_test"}
    return search_query

With a Pod running, you can now run your tests using:

pytest

Publish your plugin

To be able to use your plugin, you have to publish a docker container with your plugin to the gitlab container registry of your repo. As an example you can look at the registry for the plugin we are creating. You can navigate to the registry of your gitlab repo by going to Packages & Registries -> Container Registry. Using the plugin template, publishing is just a matter of pushing your code to your repo in the dev or prod branch. Let’s try that:

git add -A
git commit -m "publish v0.1"
git push

Now, because of the build_image stage that was defined by the template in our .gitlab-ci.yml, a gitlab ci pipeline is started which builds our plugin docker image as described in our Dockerfile, and automatically uploads it to the registry of your repo. It will take a few minute before the pipeline is completed and the image shows up into your container registry. You can see your ci pipeline and its progress in your repo under CI/CD -> Pipelines, for an example, see this.

Install the plugin in your frontend

To install the plugin in the frontend, we need to link the docker image from our container registry to the frontend. Before we can do this, we need a file called config.json, which we can create this file using the pymemri CLI:

create_plugin_config

The resulting config.json defines the arguments of your plugin that are defined in your Plugin class in sentiment_plugin/plugin.py. For instance, maybe you want to change your model in some cases, or change the language settings. It looks roughly like this:

[
  {
    "name": "model_name",
    "display": "Model Name",
    "data_type": "Text",
    "type": "textbox",
    "default": "cardiffnlp/twitter-xlm-roberta-base-sentiment",
    "optional": true
  },
  {
    "name": "model_version",
    "display": "Model Version",
    "data_type": "Text",
    "type": "textbox",
    "default": "0.1",
    "optional": true
  },
  ...
]

Equiped with our config, we can now register the plugin in the frontend. We can do this by simply entering the url (E.g. https://gitlab.memri.io/eelcovdw/sentiment_plugin) to your repo in the frontend.

Optional step: preload model

To prevent downloading the large RoBERTa model every time the plugin is started, we can preload the model when building the docker image. The script tools/preload.py is called when the CI builds the image, and any assets required to run the plugin can be downloaded here. All Hugging Face models are automatically cached on disk, so we only need to load our model and tokenizer:

from transformers import AutoModelForSequenceClassification, AutoTokenizer

def preload():
    """Download the default model and tokenizer to cache on disk"""
    model_name = "cardiffnlp/twitter-xlm-roberta-base-sentiment"
    print("Preloading RoBERTa model...")
    AutoModelForSequenceClassification.from_pretrained(model_name)
    AutoTokenizer.from_pretrained(model_name)

if __name__=="__main__":
    preload()