Robust Intelligence 🤝 🤗 Integration Walkthrough

RIME natively supports stress testing and monitoring of Hugging Face models through a native integration with their API’s. Below is an example run stress testing a sentiment classification transformer model on the Amazon Polarity dataset. For more information on how to connect with a Hugging Face Model or Hugging Face Dataset, check out the linked documentation.

First, specify the URL and api token to connect with your cluster.

Latest Colab version of this notebook available here

# set these before beginning!
CLUSTER_URL = '' # e.g., rime.<name>.rime.dev
API_TOKEN = ''

Create a Managed Image

Since the Hugging Face integration requires a couple additional dependencies, we will first create a Docker image to be managed by RIME.

%pip install rime-sdk &> /dev/null
from rime_sdk import Client

image_name = "huggingface"

# connect to your cluster
rime_client = Client(CLUSTER_URL, api_key=API_TOKEN)

# Specify pip requirements for the run. Some models require additional dependencies.
requirements = [
     rime_client.pip_requirement("transformers"),
     rime_client.pip_requirement("datasets"),
     # Uncomment if you need sentencepiece
     # rime_client.pip_requirement("sentencepiece"),
   ]

if not rime_client.has_managed_image(image_name, check_status=True):
    # e.g., if the image build job failed
    if rime_client.has_managed_image(image_name):
        rime_client.delete_managed_image(image_name)
    # Start a new image building job
    builder_job = rime_client.create_managed_image(image_name, requirements)

    # Wait until the job has finished and print out status information.
    # Once this prints out the `READY` status, your image is available for use in stress tests.
    builder_job.get_status(verbose=True, wait_until_finish=True)

Run Stress Test

Now it’s time to kick off a stress testing run. We will synchronously monitor the job using the get_status() command.

config = {
    "run_name": "Sentiment Classification",
    "data_info": {
        "type": "huggingface",
        "dataset_uri": "amazon_polarity",
        "text_key": "content"
    },
    "prediction_info": {"n_samples": 5000}, # limit to 5k predictions to speed up run.
    "model_info": {
        "type": "huggingface_classification",
        "model_uri": "BaxterAI/SentimentClassifier"
    },
    "model_task": "Text Classification"
}


job = rime_client.start_stress_test(test_run_config=config, data_type='nlp', rime_managed_image=image_name)
job.get_status(wait_until_finish=True, verbose=True)

Review Results

Now that the test has completed, check out the test run results using the link below.

test_run = job.get_test_run()
test_run