Monitoring NLP pipelines

Rubrix currently gives users several ways to monitor and observe model predictions.

This brief guide introduces the different methods and expected usages.

Using rb.monitor

For widely-used libraries Rubrix includes an “auto-monitoring” option via the rb.monitor method. Currently supported libraries are Hugging Face Transformers and spaCy, if you’d like to see another library supported feel free to add a discussion or issue on GitHub.

rb.monitor will wrap HF and spaCy pipelines so every time you call them, the output of these calls will be logged into the dataset of your choice, as a background process, in a non-blocking way. Additionally, rb.monitor will add several tags to your dataset such as the library build version, the model name, the language, etc. This should also work for custom (private) pipelines, not only the Hub’s or official spaCy models.

It is worth noting that this feature is useful beyond monitoring, and can be used for data collection (e.g., bootstrapping data annotation with pre-trained pipelines), model development (e.g., error analysis), and model evaluation (e.g., combined with data annotation to obtain evaluation metrics).

Let’s see it in action using the IMDB dataset:

from datasets import load_dataset

dataset = load_dataset("imdb", split="test[0:1000]")

Hugging Face Transformer Pipelines

Rubrix currently supports monitoringtext-classification and zero-shot-classification pipelines, but token-classification and text2text pipelines will be added in coming releases.

from transformers import pipeline
import rubrix as rb

nlp = pipeline("sentiment-analysis", return_all_scores=True, padding=True, truncation=True)
nlp = rb.monitor(nlp, dataset="nlp_monitoring")

dataset.map(lambda example: {"prediction": nlp(example["text"])})

Once the map operation starts, you can start browsing the predictions in the Web-app:

The default Rubrix installation comes with Elastic’s Kibana pre-configured, so you can easily build custom monitoring dashboards and alerts (for your team and other stakeholders):

Record-level metadata is a key element of Rubrix datasets, enabling users to do fine-grained analysis and dataset slicing. Let’s see how we can log metadata while using rb.monitor. Let’s use the label in ag_news to add a news_category field for each record.

dataset

dataset.map(lambda example: {"prediction": nlp(example["text"], metadata={"news_category": example["label"]})})

spaCy

Rubrix currently supports monitoring the NER pipeline component, but textcat will be added soon.

import spacy
import rubrix as rb

nlp = spacy.load("en_core_web_sm")
nlp = rb.monitor(nlp, dataset="nlp_monitoring_spacy")

dataset.map(lambda example: {"prediction": nlp(example["text"])})

Once the map operation starts, you can start browsing the predictions in the Web-app:

Flair

Rubrix currently supports monitoring Flair NER pipelines component.

import rubrix as rb

from flair.data import Sentence
from flair.models import SequenceTagger

# load tagger
tagger = rb.monitor(SequenceTagger.load("flair/ner-english"), dataset="flair-example", sample_rate=1.0)

# make example sentence
sentence = Sentence("George Washington went to Washington")

# predict NER tags. This will log the prediction in Rubrix
tagger.predict(sentence)

The following logs the predictions over the IMDB dataset:

def make_prediction(example):
    tagger.predict(Sentence(example["text"]))
    return {"prediction": True}

dataset.map(make_prediction)

Using rb.log in background mode

You can monitor your own models without adding a response delay by using the background param in rb.log

Let’s see an example using BentoML with a spaCy NER pipeline:

import spacy

nlp = spacy.load("en_core_web_sm")

%%writefile spacy_model.py


from bentoml import BentoService, api, artifacts, env
from bentoml.adapters import JsonInput
from bentoml.frameworks.spacy import SpacyModelArtifact

import rubrix as rb


@env(infer_pip_packages=True)
@artifacts([SpacyModelArtifact("nlp")])
class SpacyNERService(BentoService):

    @api(input=JsonInput(), batch=True)
    def predict(self, parsed_json_list):
        result, rb_records = ([], [])
        for index, parsed_json in enumerate(parsed_json_list):
            doc = self.artifacts.nlp(parsed_json["text"])
            prediction = [{"entity": ent.text, "label": ent.label_} for ent in doc.ents]
            rb_records.append(
                rb.TokenClassificationRecord(
                    text=doc.text,
                    tokens=[t.text for t in doc],
                    prediction=[
                        (ent.label_, ent.start_char, ent.end_char) for ent in doc.ents
                    ],
                )
            )
            result.append(prediction)

        rb.log(
            name="monitor-for-spacy-ner",
            records=rb_records,
            tags={"framework": "bentoml"},
            background=True,
            verbose=False
        ) # By using the background=True, the model latency won't be affected

        return result

from spacy_model import SpacyNERService

svc = SpacyNERService()
svc.pack('nlp', nlp)

saved_path = svc.save()

You can predict some data without serving the model. Just launch following command:

!bentoml run SpacyNERService:latest predict --input "{\"text\":\"I am driving BMW\"}"

If you’re running Rubrix in local, go to http://localhost:6900/datasets/rubrix/monitor-for-spacy-ner and see that the new dataset monitor-for-spacy-ner contains your data

Using the ASGI middleware

For using the ASGI middleware, see this tutorial