RI NLP Walkthrough using REST API

▶️ Try this in Colab! Run the RI NLP REST Walkthrough in Google Colab.

In this Notebook Walkthrough, we will walkthrough one of RIME’s core products using our REST API - AI Stress Testing.

AI Stress Testing is used in the model development stage. Using AI Stress Testing you can test the developed model. RIME goes beyond simply optimizing for basic model performance like accuracy and automatically discovers the model’s weaknesses.

import json
import requests
import time
import uuid
from datetime import datetime
from typing import List, Optional, Tuple

Establish the RIME Client.

To get started, provide the API credentials and the base domain/address of the RIME service. You can generate and copy an API token from the API Access Tokens Page under Workspace settings. For the domain/address of the RIME service, contact your admin.

API_TOKEN = "" # Paste API_KEY
headers = {"rime-api-key": API_TOKEN}

CLUSTER_URL = "" # PASTE DOMAIN OF RIME SERVICE (e.g., https://rime.example.rbst.io)
API_URL = f"https://{CLUSTER_URL}"

AGENT_ID = '' # PASTE AGENT_ID IF USING AN AGENT THAT IS NOT THE DEFAULT

Create a New Project.

You can create Projects in RIME to organize your Test Runs. Each Project represents a workspace for a given machine learning task. It can contain multiple candidate models, but should only contain one promoted production model.

# Create a Project for the Natural Language Inference model task.
req = {
    "name": "NLP Example",
    "description": "Example stress test on NLP data.",
    "model_task": "MODEL_TASK_NATURAL_LANGUAGE_INFERENCE"
}
api_endpoint = f"{API_URL}/v1/projects"
res = requests.post(
    api_endpoint,
    json=req,
    headers=headers,
)
assert res.status_code == 200
# Get the Project ID from the response. This will be needed to start a Stress Test.
project_id = res.json()["project"]["id"]

Optional: Create Job Tracker.

Once a Stress Test request or Managed Image request has been sent, we have to wait until the Job completes before checking for results. The following helper functions track the status of a job by blocking until the Job has status “Succeeded”.

# Helper function for displaying Job Status.
def _add_padding(string: str) -> str:
    max_str_len = len("JOB_STATUS_REQUESTED")
    padding_len = max(max_str_len - len(string), 0)
    return string + " " * padding_len

# Helper function for periodically querying for Job status.
def track_job(job_id):
    time.sleep(1)
    api_endpoint = f"{API_URL}/v1/jobs/{job_id}"
    res = requests.get(api_endpoint, headers=headers)
    job_status = res.json()["job"]["status"]

    i = 0
    while job_status != "JOB_STATUS_SUCCEEDED":
        res = requests.get(api_endpoint, headers=headers)
        res_json = res.json()
        job_status = res.json()["job"]["status"]
        i += 1
        print(f"Poll count: {i}.\tJob status: {_add_padding(job_status)}", end="\r")
        if job_status == "JOB_STATUS_FAILED":
            logs_url = res_json["job"]["archivedJobLogs"]["url"]["url"]
            raise ValueError(f"Test run failed. Job logs available at {logs_url}.")
        time.sleep(10)

    return res_json

Create Managed Image.

In order to run a Stress Test against a huggingface model, we have to create a Managed Image with the model’s pip requirements pre-installed. The following snippet checks if a huggingface Managed Image exists, and if not, creates one.

managed_image_name = "huggingface_image"

api_endpoint = f"{API_URL}/images/{managed_image_name}"
res = requests.get(api_endpoint, headers=headers)
managed_image_exists = (res.status_code == 200)

if not managed_image_exists:
    req = {
        "name": managed_image_name,
        "pip_requirements": [
            {"name": "transformers"},
            {"name": "datasets"},
        ],
    }
    api_endpoint = f"{API_URL}/images"
    res = requests.post(api_endpoint, headers=headers)
    job_id = res.json()["job"]["jobId"]

    track_job(job_id)

# These are the URIs of the model that we want to Stress Test, and an associated dataset.
DATASET_URI = "rungalileo/snli"
MODEL_URI = "cross-encoder/nli-MiniLM2-L6-H768"

Register Dataset

Now that we have the huggingface model URIs, we can register them and the dataset to RIME. Once they’re registed, we can refer to these resources using their RIME-generated ID’s.

dt = str(datetime.now())

# Register the huggingface dataset.
def register_dataset(dataset_uri, split_name):
    api_endpoint = f"{API_URL}/v1/registry/{project_id['uuid']}/dataset"
    req = {
        # Note: models and datasets need to have unique names.
        "name": f"{dataset_uri}_{split_name}_{dt}",
        "data_info": {
            "connection_info": {
                "hugging_face": {
                    "dataset_uri": dataset_uri,
                    "split_name": split_name,
                    # TODO: remove line below
                    "loading_params_json": json.dumps({}),
                },
            },
            "data_params": {
                "label_col": "label",
                "text_features": ["premise", "hypothesis"],
                "class_names": ["Contradiction", "Neutral", "Entailment"],
            },
        },
    }
    if AGENT_ID:
        req["agent_id"] = {"uuid": AGENT_ID}
    res = requests.post(api_endpoint, json=req, headers=headers)
    return res.json()["datasetId"]


ref_data_id = register_dataset(DATASET_URI, "train")
eval_data_id = register_dataset(DATASET_URI, "test")

Register Model

api_endpoint = f"{API_URL}/v1/registry/{project_id['uuid']}/model"
req = {
    # Note: models and datasets need to have unique names.
    "name": f"{MODEL_URI}_{dt}",
    "model_info": {
        "hugging_face": {
            "model_uri": MODEL_URI,
        },
    },
}
if AGENT_ID:
    req["agent_id"] = {"uuid": AGENT_ID}
res = requests.post(api_endpoint, json=req, headers=headers)
model_id = res.json()["modelId"]

Create Test Run Config

Below is a sample configuration of how to setup and run a RIME Stress Test. The configuration includes important information, such as the IDs of the model, reference dataset, evaluation dataset and the name of a Managed Image to use (optional).

config = {
    "run_name": "NLP Stress Test",
    "model_id": model_id,
    "data_info": {
        "ref_dataset_id": ref_data_id,
        "eval_dataset_id": eval_data_id,
    },
    "run_time_info": {
        "custom_image": {
            "managed_image_name": managed_image_name,
        },
    },
    "test_suite_config": {
        "global_exclude_columns": ["id"]
    },
}

Run Stress Test

With a completed Test Run Configuration, we can now start a Stress Test and track its progress until the tests complete.

req = {"test_run_config": config}
if AGENT_ID:
    req["agent_id"] = {"uuid": AGENT_ID}
api_endpoint = f"{API_URL}/v1/stress-tests/{project_id['uuid']}"
res = requests.post(api_endpoint, json=req, headers=headers)
job_id = res.json()["job"]["jobId"]

# Track Stress Test Job Progress.
res_json = track_job(job_id)

# Once a Stress Test Job has completed, get the Test Run ID to access results.
test_run_id = res_json["job"]["jobData"]["stress"]["testRunId"]

Get Test Run Results.

Once a Test Run completes, we can access the results of the Stress Test using the Test Run ID.

# Print information about a completed Test Run.
api_endpoint = f"{API_URL}/v1/test-runs/{test_run_id}"
res = requests.get(api_endpoint, headers=headers)
print(res.json()["testRun"])

# List individual Test Cases from the Test Run.
api_endpoint = f"{API_URL}/v1/test-cases"
payload = {"listTestCasesQuery.testRunId": test_run_id}
res = requests.get(api_endpoint, headers=headers, params=payload)
print(res.json()["testCases"])