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.

# 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