{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This notebook contains Llama Stack implementation of a common end-to-end workflow for customizing and evaluating LLMs using the NVIDIA provider."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prerequisites\n",
"First, ensure the NeMo Microservices platform is up and running, including the model downloading step for `meta/llama-3.2-1b-instruct`. See installation instructions: https://aire.gitlab-master-pages.nvidia.com/microservices/documentation/latest/nemo-microservices/latest-internal/set-up/deploy-as-platform/index.html (TODO: Update to public docs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, set up your development environment on your machine. From the root of the project, set up your virtual environment:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "shellscript"
}
},
"outputs": [],
"source": [
"uv sync --extra dev\n",
"uv pip install -e .\n",
"source .venv/bin/activate"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Build the Llama Stack image using the virtual environment. For local development, set `LLAMA_STACK_DIR` to ensure your local code is use in the image. To use the production version of `llama-stack`, omit `LLAMA_STACK_DIR`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "shellscript"
}
},
"outputs": [],
"source": [
"LLAMA_STACK_DIR=$(pwd) llama stack build --template nvidia --image-type venv"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Configure the environment variables for each service.\n",
"\n",
"If needed, update the URLs for each service to point to your deployment.\n",
"- NDS_URL: NeMo Data Store URL\n",
"- NEMO_URL: NeMo Microservices Platform URL\n",
"- NIM_URL: NIM URL\n",
"\n",
"For more infomation about these variables, please reference the [NVIDIA Distro documentation](docs/source/distributions/remote_hosted_distro/nvidia.md)."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"# NVIDIA URLs\n",
"NDS_URL = \"https://datastore.int.aire.nvidia.com\"\n",
"NEMO_URL = \"https://nmp.int.aire.nvidia.com\"\n",
"NIM_URL = \"https://nim.int.aire.nvidia.com\"\n",
"\n",
"USER_ID = \"llama-stack-user\"\n",
"NAMESPACE = \"default\"\n",
"PROJECT_ID = \"\"\n",
"CUSTOMIZED_MODEL_DIR = \"jg-test-llama-stack@v2\"\n",
"\n",
"# Inference env vars\n",
"os.environ[\"NVIDIA_BASE_URL\"] = NIM_URL\n",
"\n",
"# Customizer env vars\n",
"os.environ[\"NVIDIA_CUSTOMIZER_URL\"] = NEMO_URL\n",
"os.environ[\"NVIDIA_USER_ID\"] = USER_ID\n",
"os.environ[\"NVIDIA_DATASET_NAMESPACE\"] = NAMESPACE\n",
"os.environ[\"NVIDIA_PROJECT_ID\"] = PROJECT_ID\n",
"os.environ[\"NVIDIA_OUTPUT_MODEL_DIR\"] = CUSTOMIZED_MODEL_DIR\n",
"\n",
"# Evaluator env vars\n",
"os.environ[\"NVIDIA_EVALUATOR_URL\"] = NEMO_URL\n",
"\n",
"# Guardrails env vars\n",
"os.environ[\"GUARDRAILS_SERVICE_URL\"] = NEMO_URL\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/jgulabrai/Projects/forks/llama-stack/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
}
],
"source": [
"import asyncio\n",
"import json\n",
"import os\n",
"import pprint\n",
"from time import sleep, time\n",
"from typing import Dict\n",
"\n",
"import aiohttp\n",
"import requests\n",
"from huggingface_hub import HfApi\n",
"\n",
"os.environ[\"HF_ENDPOINT\"] = f\"{NDS_URL}/v1/hf\"\n",
"os.environ[\"HF_TOKEN\"] = \"token\"\n",
"\n",
"hf_api = HfApi(endpoint=os.environ.get(\"HF_ENDPOINT\"), token=os.environ.get(\"HF_TOKEN\"))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Set Up Llama Stack Client\n",
"Begin by importing the necessary components from Llama Stack's client library:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Warning: `bwrap` is not available. Code interpreter tool will not work correctly.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Initializing NVIDIASafetyAdapter(https://nmp.int.aire.nvidia.com)...\n"
]
},
{
"data": {
"text/html": [
"Using config nvidia:\n",
"\n"
],
"text/plain": [
"Using config \u001b[34mnvidia\u001b[0m:\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"apis:\n",
"- agents\n",
"- datasetio\n",
"- eval\n",
"- inference\n",
"- post_training\n",
"- safety\n",
"- scoring\n",
"- telemetry\n",
"- tool_runtime\n",
"- vector_io\n",
"benchmarks: []\n",
"container_image: null\n",
"datasets: []\n",
"image_name: nvidia\n",
"logging: null\n",
"metadata_store:\n",
" db_path: /Users/jgulabrai/.llama/distributions/nvidia/registry.db\n",
" namespace: null\n",
" type: sqlite\n",
"models:\n",
"- metadata: {}\n",
" model_id: meta/llama3-8b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-8b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3-8B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-8b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama3-70b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-70b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3-70B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-70b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.1-8b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-8b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.1-8B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-8b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.1-70b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-70b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.1-70B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-70b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.1-405b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-405b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.1-405B-Instruct-FP8\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.1-405b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.2-1b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-1b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.2-1B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-1b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.2-3b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-3b-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.2-3B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-3b-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.2-11b-vision-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-11b-vision-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.2-11B-Vision-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-11b-vision-instruct\n",
"- metadata: {}\n",
" model_id: meta/llama-3.2-90b-vision-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-90b-vision-instruct\n",
"- metadata: {}\n",
" model_id: meta-llama/Llama-3.2-90B-Vision-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-3.2-90b-vision-instruct\n",
"- metadata:\n",
" context_length: 8192\n",
" embedding_dimension: 2048\n",
" model_id: nvidia/llama-3.2-nv-embedqa-1b-v2\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/llama-3.2-nv-embedqa-1b-v2\n",
"- metadata:\n",
" context_length: 512\n",
" embedding_dimension: 1024\n",
" model_id: nvidia/nv-embedqa-e5-v5\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/nv-embedqa-e5-v5\n",
"- metadata:\n",
" context_length: 512\n",
" embedding_dimension: 4096\n",
" model_id: nvidia/nv-embedqa-mistral-7b-v2\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/nv-embedqa-mistral-7b-v2\n",
"- metadata:\n",
" context_length: 512\n",
" embedding_dimension: 1024\n",
" model_id: snowflake/arctic-embed-l\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: snowflake/arctic-embed-l\n",
"providers:\n",
" agents:\n",
" - config:\n",
" persistence_store:\n",
" db_path: /Users/jgulabrai/.llama/distributions/nvidia/agents_store.db\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: meta-reference\n",
" provider_type: inline::meta-reference\n",
" datasetio:\n",
" - config:\n",
" kvstore:\n",
" db_path: /Users/jgulabrai/.llama/distributions/nvidia/localfs_datasetio.db\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: localfs\n",
" provider_type: inline::localfs\n",
" eval:\n",
" - config:\n",
" evaluator_service_url: https://nmp.int.aire.nvidia.com\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" inference:\n",
" - config:\n",
" api_key: '********'\n",
" url: https://nim.int.aire.nvidia.com\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" post_training:\n",
" - config:\n",
" api_key: '********'\n",
" customizer_url: https://nmp.int.aire.nvidia.com\n",
" dataset_namespace: default\n",
" project_id: test-project\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" safety:\n",
" - config:\n",
" config_id: self-check\n",
" guardrails_service_url: https://nmp.int.aire.nvidia.com\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" scoring:\n",
" - config: {}\n",
" provider_id: basic\n",
" provider_type: inline::basic\n",
" telemetry:\n",
" - config:\n",
" service_name: \"\\u200B\"\n",
" sinks: sqlite\n",
" sqlite_db_path: /Users/jgulabrai/.llama/distributions/nvidia/trace_store.db\n",
" provider_id: meta-reference\n",
" provider_type: inline::meta-reference\n",
" tool_runtime:\n",
" - config: {}\n",
" provider_id: rag-runtime\n",
" provider_type: inline::rag-runtime\n",
" vector_io:\n",
" - config:\n",
" kvstore:\n",
" db_path: /Users/jgulabrai/.llama/distributions/nvidia/faiss_store.db\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: faiss\n",
" provider_type: inline::faiss\n",
"scoring_fns: []\n",
"server:\n",
" auth: null\n",
" port: 8321\n",
" tls_certfile: null\n",
" tls_keyfile: null\n",
"shields: []\n",
"tool_groups:\n",
"- args: null\n",
" mcp_endpoint: null\n",
" provider_id: rag-runtime\n",
" toolgroup_id: builtin::rag\n",
"vector_dbs: []\n",
"version: '2'\n",
"\n",
"\n"
],
"text/plain": [
"apis:\n",
"- agents\n",
"- datasetio\n",
"- eval\n",
"- inference\n",
"- post_training\n",
"- safety\n",
"- scoring\n",
"- telemetry\n",
"- tool_runtime\n",
"- vector_io\n",
"benchmarks: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
"container_image: null\n",
"datasets: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
"image_name: nvidia\n",
"logging: null\n",
"metadata_store:\n",
" db_path: \u001b[35m/Users/jgulabrai/.llama/distributions/nvidia/\u001b[0m\u001b[95mregistry.db\u001b[0m\n",
" namespace: null\n",
" type: sqlite\n",
"models:\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama3-8b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-8b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3\u001b[0m-8B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-8b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama3-70b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-70b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3\u001b[0m-70B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama3-70b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-8b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-8b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-8B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-8b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-70b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-70b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-70B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-70b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-405b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-405b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-405B-Instruct-FP8\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.1\u001b[0m-405b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-1b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-1b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-1B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-1b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-3b-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-3b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-3B-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-3b-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-11b-vision-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-11b-vision-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-11B-Vision-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-11b-vision-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-90b-vision-instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-90b-vision-instruct\n",
"- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-90B-Vision-Instruct\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - llm\n",
" provider_id: nvidia\n",
" provider_model_id: meta/llama-\u001b[1;36m3.2\u001b[0m-90b-vision-instruct\n",
"- metadata:\n",
" context_length: \u001b[1;36m8192\u001b[0m\n",
" embedding_dimension: \u001b[1;36m2048\u001b[0m\n",
" model_id: nvidia/llama-\u001b[1;36m3.2\u001b[0m-nv-embedqa-1b-v2\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/llama-\u001b[1;36m3.2\u001b[0m-nv-embedqa-1b-v2\n",
"- metadata:\n",
" context_length: \u001b[1;36m512\u001b[0m\n",
" embedding_dimension: \u001b[1;36m1024\u001b[0m\n",
" model_id: nvidia/nv-embedqa-e5-v5\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/nv-embedqa-e5-v5\n",
"- metadata:\n",
" context_length: \u001b[1;36m512\u001b[0m\n",
" embedding_dimension: \u001b[1;36m4096\u001b[0m\n",
" model_id: nvidia/nv-embedqa-mistral-7b-v2\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: nvidia/nv-embedqa-mistral-7b-v2\n",
"- metadata:\n",
" context_length: \u001b[1;36m512\u001b[0m\n",
" embedding_dimension: \u001b[1;36m1024\u001b[0m\n",
" model_id: snowflake/arctic-embed-l\n",
" model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
" - embedding\n",
" provider_id: nvidia\n",
" provider_model_id: snowflake/arctic-embed-l\n",
"providers:\n",
" agents:\n",
" - config:\n",
" persistence_store:\n",
" db_path: \u001b[35m/Users/jgulabrai/.llama/distributions/nvidia/\u001b[0m\u001b[95magents_store.db\u001b[0m\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: meta-reference\n",
" provider_type: inline::meta-reference\n",
" datasetio:\n",
" - config:\n",
" kvstore:\n",
" db_path: \u001b[35m/Users/jgulabrai/.llama/distributions/nvidia/\u001b[0m\u001b[95mlocalfs_datasetio.db\u001b[0m\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: localfs\n",
" provider_type: inline::localfs\n",
" eval:\n",
" - config:\n",
" evaluator_service_url: \u001b[4;94mhttps://nmp.int.aire.nvidia.com\u001b[0m\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" inference:\n",
" - config:\n",
" api_key: \u001b[32m'********'\u001b[0m\n",
" url: \u001b[4;94mhttps://nim.int.aire.nvidia.com\u001b[0m\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" post_training:\n",
" - config:\n",
" api_key: \u001b[32m'********'\u001b[0m\n",
" customizer_url: \u001b[4;94mhttps://nmp.int.aire.nvidia.com\u001b[0m\n",
" dataset_namespace: default\n",
" project_id: test-project\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" safety:\n",
" - config:\n",
" config_id: self-check\n",
" guardrails_service_url: \u001b[4;94mhttps://nmp.int.aire.nvidia.com\u001b[0m\n",
" provider_id: nvidia\n",
" provider_type: remote::nvidia\n",
" scoring:\n",
" - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" provider_id: basic\n",
" provider_type: inlin\u001b[1;92me::ba\u001b[0msic\n",
" telemetry:\n",
" - config:\n",
" service_name: \u001b[32m\"\\u200B\"\u001b[0m\n",
" sinks: sqlite\n",
" sqlite_db_path: \u001b[35m/Users/jgulabrai/.llama/distributions/nvidia/\u001b[0m\u001b[95mtrace_store.db\u001b[0m\n",
" provider_id: meta-reference\n",
" provider_type: inline::meta-reference\n",
" tool_runtime:\n",
" - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
" provider_id: rag-runtime\n",
" provider_type: inline::rag-runtime\n",
" vector_io:\n",
" - config:\n",
" kvstore:\n",
" db_path: \u001b[35m/Users/jgulabrai/.llama/distributions/nvidia/\u001b[0m\u001b[95mfaiss_store.db\u001b[0m\n",
" namespace: null\n",
" type: sqlite\n",
" provider_id: faiss\n",
" provider_type: inlin\u001b[1;92me::fa\u001b[0miss\n",
"scoring_fns: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
"server:\n",
" auth: null\n",
" port: \u001b[1;36m8321\u001b[0m\n",
" tls_certfile: null\n",
" tls_keyfile: null\n",
"shields: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
"tool_groups:\n",
"- args: null\n",
" mcp_endpoint: null\n",
" provider_id: rag-runtime\n",
" toolgroup_id: builtin::rag\n",
"vector_dbs: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
"version: \u001b[32m'2'\u001b[0m\n",
"\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from llama_stack.distribution.library_client import LlamaStackAsLibraryClient\n",
"\n",
"client = LlamaStackAsLibraryClient(\"nvidia\")\n",
"client.initialize()"
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {},
"outputs": [],
"source": [
"# Helper functions for waiting on jobs\n",
"from llama_stack.apis.common.job_types import JobStatus\n",
"\n",
"def wait_customization_job(job_id: str, polling_interval: int = 10, timeout: int = 6000):\n",
" start_time = time()\n",
"\n",
" response = client.post_training.job.status(job_uuid=job_id)\n",
" job_status = response.status\n",
"\n",
" print(f\"Waiting for Customization job {job_id} to finish.\")\n",
" print(f\"Job status: {job_status} after {time() - start_time} seconds.\")\n",
"\n",
" while job_status in [JobStatus.scheduled.value, JobStatus.in_progress.value]:\n",
" sleep(polling_interval)\n",
" response = client.post_training.job.status(job_uuid=job_id)\n",
" job_status = response.status\n",
"\n",
" print(f\"Job status: {job_status} after {time() - start_time} seconds.\")\n",
"\n",
" if time() - start_time > timeout:\n",
" raise RuntimeError(f\"Customization Job {job_id} took more than {timeout} seconds.\")\n",
" \n",
" return job_status\n",
"\n",
"def wait_eval_job(benchmark_id: str, job_id: str, polling_interval: int = 10, timeout: int = 6000):\n",
" start_time = time()\n",
" job_status = client.eval.jobs.status(benchmark_id=benchmark_id, job_id=job_id)\n",
"\n",
" print(f\"Waiting for Evaluation job {job_id} to finish.\")\n",
" print(f\"Job status: {job_status} after {time() - start_time} seconds.\")\n",
"\n",
" while job_status.status in [JobStatus.scheduled.value, JobStatus.in_progress.value]:\n",
" sleep(polling_interval)\n",
" job_status = client.eval.jobs.status(benchmark_id=benchmark_id, job_id=job_id)\n",
"\n",
" print(f\"Job status: {job_status} after {time() - start_time} seconds.\")\n",
"\n",
" if time() - start_time > timeout:\n",
" raise RuntimeError(f\"Evaluation Job {job_id} took more than {timeout} seconds.\")\n",
"\n",
" return job_status\n",
"\n",
"def wait_nim_loads_customized_model(model_id: str, namespace: str, polling_interval: int = 10, timeout: int = 300):\n",
" found = False\n",
" start_time = time()\n",
"\n",
" model_path = f\"{namespace}/{model_id}\"\n",
" print(f\"Checking if NIM has loaded customized model {model_path}.\")\n",
"\n",
" while not found:\n",
" sleep(polling_interval)\n",
"\n",
" response = requests.get(f\"{NIM_URL}/v1/models\")\n",
" if model_path in [model[\"id\"] for model in response.json()[\"data\"]]:\n",
" found = True\n",
" print(f\"Model {model_path} available after {time() - start_time} seconds.\")\n",
" break\n",
" else:\n",
" print(f\"Model {model_path} not available after {time() - start_time} seconds.\")\n",
"\n",
" if not found:\n",
" raise RuntimeError(f\"Model {model_path} not available after {timeout} seconds.\")\n",
"\n",
" assert found, f\"Could not find model {model_path} in the list of available models.\"\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Upload Dataset Using the HuggingFace Client"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"sample_squad_test_dataset_name = \"squad-test-dataset\"\n",
"repo_id = f\"{NAMESPACE}/{sample_squad_test_dataset_name}\""
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"# Create the repo\n",
"res = hf_api.create_repo(repo_id, repo_type=\"dataset\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Upload the files from the local folder\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_data/training\",\n",
" path_in_repo=\"training\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_data/validation\",\n",
" path_in_repo=\"validation\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_data/testing\",\n",
" path_in_repo=\"testing\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create the dataset\n",
"# response = client.datasets.register(...)\n",
"response = requests.post(\n",
" url=f\"{NEMO_URL}/v1/datasets\",\n",
" json={\n",
" \"name\": sample_squad_test_dataset_name,\n",
" \"namespace\": NAMESPACE,\n",
" \"description\": \"Dataset created from llama-stack e2e notebook\",\n",
" \"files_url\": f\"hf://datasets/{NAMESPACE}/{sample_squad_test_dataset_name}\",\n",
" },\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to create dataset {response.text}\"\n",
"json.dumps(response.json(), indent=2)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check the files URL\n",
"# response = client.datasets.retrieve(repo_id)\n",
"# dataset = response.model_dump()\n",
"# assert dataset[\"source\"][\"uri\"] == f\"hf://datasets/{repo_id}\"\n",
"response = requests.get(\n",
" url=f\"{NEMO_URL}/v1/datasets/{NAMESPACE}/{sample_squad_test_dataset_name}\",\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to fetch dataset {response.text}\"\n",
"dataset_obj = response.json()\n",
"print(\"Files URL:\", dataset_obj[\"files_url\"])\n",
"assert dataset_obj[\"files_url\"] == f\"hf://datasets/{repo_id}\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Inference"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"import pprint\n",
"\n",
"with open(\"./tmp/sample_squad_data/testing/testing.jsonl\", \"r\") as f:\n",
" examples = [json.loads(line) for line in f]\n",
"\n",
"# Get the user prompt from the last example\n",
"sample_prompt = examples[-1][\"prompt\"]\n",
"pprint.pprint(sample_prompt)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Test inference\n",
"response = client.inference.chat_completion(\n",
" messages=[\n",
" {\"role\": \"user\", \"content\": sample_prompt}\n",
" ],\n",
" model_id=\"meta/llama-3.1-8b-instruct\",\n",
" sampling_params={\n",
" \"max_tokens\": 20,\n",
" \"strategy\": {\n",
" \"type\": \"top_p\",\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9\n",
" }\n",
" }\n",
")\n",
"print(f\"Inference response: {response.completion_message.content}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluation\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"benchmark_id = \"test-eval-config-1\""
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"# Register a benchmark, which creates an Evaluation Config\n",
"simple_eval_config = {\n",
" \"benchmark_id\": benchmark_id,\n",
" \"dataset_id\": \"\",\n",
" \"scoring_functions\": [],\n",
" \"metadata\": {\n",
" \"type\": \"custom\",\n",
" \"params\": {\"parallelism\": 8},\n",
" \"tasks\": {\n",
" \"qa\": {\n",
" \"type\": \"completion\",\n",
" \"params\": {\n",
" \"template\": {\n",
" \"prompt\": \"{{prompt}}\",\n",
" \"max_tokens\": 20,\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9,\n",
" },\n",
" },\n",
" \"dataset\": {\"files_url\": f\"hf://datasets/{repo_id}/testing/testing.jsonl\"},\n",
" \"metrics\": {\n",
" \"bleu\": {\n",
" \"type\": \"bleu\",\n",
" \"params\": {\"references\": [\"{{ideal_response}}\"]},\n",
" },\n",
" \"string-check\": {\n",
" \"type\": \"string-check\",\n",
" \"params\": {\"check\": [\"{{ideal_response | trim}}\", \"equals\", \"{{output_text | trim}}\"]},\n",
" },\n",
" },\n",
" }\n",
" }\n",
" }\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"response = client.benchmarks.register(\n",
" benchmark_id=benchmark_id,\n",
" dataset_id=repo_id,\n",
" scoring_functions=simple_eval_config[\"scoring_functions\"],\n",
" metadata=simple_eval_config[\"metadata\"]\n",
")\n",
"print(f\"Created benchmark {benchmark_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Launch a simple evaluation with the benchmark\n",
"response = client.eval.run_eval(\n",
" benchmark_id=benchmark_id,\n",
" benchmark_config={\n",
" \"eval_candidate\": {\n",
" \"type\": \"model\",\n",
" \"model\": \"meta/llama-3.1-8b-instruct\"\n",
" }\n",
" }\n",
")\n",
"job_id = response.model_dump()[\"job_id\"]\n",
"print(f\"Created evaluation job {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Wait for the job to complete\n",
"job = wait_eval_job(benchmark_id=benchmark_id, job_id=job_id, polling_interval=5, timeout=600)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f\"Job {job_id} status: {job.status}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job_results = client.eval.jobs.retrieve(benchmark_id=benchmark_id, job_id=job_id)\n",
"print(f\"Job results: {json.dumps(job_results.model_dump(), indent=2)}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract bleu score and assert it's within range\n",
"initial_bleu_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"bleu\"][\"scores\"][\"corpus\"][\"value\"]\n",
"print(f\"Initial bleu score: {initial_bleu_score}\")\n",
"\n",
"assert initial_bleu_score >= 2"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract accuracy and assert it's within range\n",
"initial_accuracy_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"string-check\"][\"scores\"][\"string-check\"][\"value\"]\n",
"print(f\"Initial accuracy: {initial_accuracy_score}\")\n",
"\n",
"assert initial_accuracy_score >= 0"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Customization"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Start the customization job\n",
"response = client.post_training.supervised_fine_tune(\n",
" job_uuid=\"\",\n",
" model=\"meta-llama/Llama-3.1-8B-Instruct\",\n",
" training_config={\n",
" \"n_epochs\": 2,\n",
" \"data_config\": {\n",
" \"batch_size\": 16,\n",
" \"dataset_id\": sample_squad_test_dataset_name,\n",
" },\n",
" \"optimizer_config\": {\n",
" \"lr\": 0.0001,\n",
" }\n",
" },\n",
" algorithm_config={\n",
" \"type\": \"LoRA\",\n",
" \"adapter_dim\": 16,\n",
" \"adapter_dropout\": 0.1,\n",
" \"alpha\": 16,\n",
" # NOTE: These fields are required by `AlgorithmConfig` model, but not directly used by NVIDIA\n",
" \"rank\": 8,\n",
" \"lora_attn_modules\": [],\n",
" \"apply_lora_to_mlp\": True,\n",
" \"apply_lora_to_output\": False\n",
" },\n",
" hyperparam_search_config={},\n",
" logger_config={},\n",
" checkpoint_dir=\"\",\n",
")\n",
"\n",
"job_id = response.job_uuid\n",
"print(f\"Created job with ID: {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Wait for the job to complete\n",
"job_status = wait_customization_job(job_id=job_id)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f\"Job {job_id} status: {job_status}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check that inference with the new model works\n",
"from llama_stack.apis.models.models import ModelType\n",
"\n",
"# TODO: Uncomment after https://github.com/meta-llama/llama-stack/pull/1859 is merged\n",
"# client.models.register(\n",
"# model_id=CUSTOMIZED_MODEL_DIR,\n",
"# model_type=ModelType.llm,\n",
"# provider_id=\"nvidia\",\n",
"# )\n",
"\n",
"# TODO: This won't work until the code above works - errors with model_id not found.\n",
"# response = client.inference.completion(\n",
"# content=\"Complete the sentence using one word: Roses are red, violets are \",\n",
"# stream=False,\n",
"# model_id=f\"default/{CUSTOMIZED_MODEL_DIR}\",\n",
"# sampling_params={\n",
"# \"max_tokens\": 50,\n",
"# },\n",
"# )\n",
"\n",
"res = requests.post(\n",
" url=f\"{NIM_URL}/v1/completions\",\n",
" json={\n",
" \"model\": f\"{namespace}/{CUSTOMIZED_MODEL_DIR}\",\n",
" \"prompt\": sample_prompt,\n",
" \"max_tokens\": 20,\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9,\n",
" },\n",
")\n",
"assert res.status_code in (200, 201), f\"Status Code {res.status_code} Failed to get adapted model completion {res.text}\"\n",
"json.dumps(res.json(), indent=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TODO: Evaluate Customized Model\n",
"Implement this section after we can register Customized model in Model Registry."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Upload Chat Dataset\n",
"Repeat fine-tuning and evaluation with a chat style dataset, which has a list of `messages` instead of a `prompt` and `completion`."
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [],
"source": [
"sample_squad_messages_dataset_name = \"test-squad-messages-dataset\"\n",
"repo_id = f\"{NAMESPACE}/{sample_squad_messages_dataset_name}\""
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {},
"outputs": [],
"source": [
"# Create the repo\n",
"# hf_api.create_repo(repo_id, repo_type=\"dataset\")\n",
"res = hf_api.create_repo(repo_id, repo_type=\"dataset\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Upload the files from the local folder\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_messages/training\",\n",
" path_in_repo=\"training\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_messages/validation\",\n",
" path_in_repo=\"validation\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")\n",
"hf_api.upload_folder(\n",
" folder_path=\"./tmp/sample_squad_messages/testing\",\n",
" path_in_repo=\"testing\",\n",
" repo_id=repo_id,\n",
" repo_type=\"dataset\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create the dataset\n",
"# response = client.datasets.register(...)\n",
"response = requests.post(\n",
" url=f\"{NEMO_URL}/v1/datasets\",\n",
" json={\n",
" \"name\": sample_squad_messages_dataset_name,\n",
" \"namespace\": NAMESPACE,\n",
" \"description\": \"Dataset created from llama-stack e2e notebook\",\n",
" \"files_url\": f\"hf://datasets/{NAMESPACE}/{sample_squad_messages_dataset_name}\",\n",
" \"project\": \"default/project-7tLfD8Lt59wFbarFceF3xN\",\n",
" },\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to create dataset {response.text}\"\n",
"json.dumps(response.json(), indent=2)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check the files URL\n",
"# response = client.datasets.retrieve(repo_id)\n",
"# dataset = response.model_dump()\n",
"# assert dataset[\"source\"][\"uri\"] == f\"hf://datasets/{repo_id}\"\n",
"response = requests.get(\n",
" url=f\"{NEMO_URL}/v1/datasets/{NAMESPACE}/{sample_squad_messages_dataset_name}\",\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to fetch dataset {response.text}\"\n",
"dataset_obj = response.json()\n",
"print(\"Files URL:\", dataset_obj[\"files_url\"])\n",
"assert dataset_obj[\"files_url\"] == f\"hf://datasets/{repo_id}\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Inference with chat/completions"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(\"./tmp/sample_squad_messages/testing/testing.jsonl\", \"r\") as f:\n",
" examples = [json.loads(line) for line in f]\n",
"\n",
"# get the user and assistant messages from the last example\n",
"sample_messages = examples[-1][\"messages\"][:-1]\n",
"pprint.pprint(sample_messages)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Test inference\n",
"response = client.inference.chat_completion(\n",
" messages=sample_messages,\n",
" model_id=\"meta/llama-3.1-8b-instruct\",\n",
" sampling_params={\n",
" \"max_tokens\": 20,\n",
" \"strategy\": {\n",
" \"type\": \"top_p\",\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9\n",
" }\n",
" }\n",
")\n",
"assert response.completion_message.content is not None\n",
"print(f\"Inference response: {response.completion_message.content}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluate with chat dataset"
]
},
{
"cell_type": "code",
"execution_count": 36,
"metadata": {},
"outputs": [],
"source": [
"benchmark_id = \"test-eval-config-chat-1\""
]
},
{
"cell_type": "code",
"execution_count": 37,
"metadata": {},
"outputs": [],
"source": [
"# Register a benchmark, which creates an Eval Config\n",
"simple_eval_config = {\n",
" \"benchmark_id\": benchmark_id,\n",
" \"dataset_id\": \"\",\n",
" \"scoring_functions\": [],\n",
" \"metadata\": {\n",
" \"type\": \"custom\",\n",
" \"params\": {\"parallelism\": 8},\n",
" \"tasks\": {\n",
" \"qa\": {\n",
" \"type\": \"completion\",\n",
" \"params\": {\n",
" \"template\": {\n",
" \"messages\": [\n",
" {\"role\": \"{{item.messages[0].role}}\", \"content\": \"{{item.messages[0].content}}\"},\n",
" {\"role\": \"{{item.messages[1].role}}\", \"content\": \"{{item.messages[1].content}}\"},\n",
" ],\n",
" \"max_tokens\": 20,\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9,\n",
" },\n",
" },\n",
" \"dataset\": {\"files_url\": f\"hf://datasets/{repo_id}/testing/testing.jsonl\"},\n",
" \"metrics\": {\n",
" \"bleu\": {\n",
" \"type\": \"bleu\",\n",
" \"params\": {\"references\": [\"{{item.messages[2].content | trim}}\"]},\n",
" },\n",
" \"string-check\": {\n",
" \"type\": \"string-check\",\n",
" \"params\": {\"check\": [\"{{item.messages[2].content}}\", \"equals\", \"{{output_text | trim}}\"]},\n",
" },\n",
" },\n",
" }\n",
" }\n",
" }\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"response = client.benchmarks.register(\n",
" benchmark_id=benchmark_id,\n",
" dataset_id=repo_id,\n",
" scoring_functions=simple_eval_config[\"scoring_functions\"],\n",
" metadata=simple_eval_config[\"metadata\"]\n",
")\n",
"print(f\"Created benchmark {benchmark_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Launch a simple evaluation with the benchmark\n",
"response = client.eval.run_eval(\n",
" benchmark_id=benchmark_id,\n",
" benchmark_config={\n",
" \"eval_candidate\": {\n",
" \"type\": \"model\",\n",
" \"model\": \"meta/llama-3.1-8b-instruct\",\n",
" }\n",
" }\n",
")\n",
"job_id = response.model_dump()[\"job_id\"]\n",
"print(f\"Created evaluation job {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Wait for the job to complete\n",
"job = wait_eval_job(benchmark_id=benchmark_id, job_id=job_id, polling_interval=5, timeout=600)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f\"Job {job_id} status: {job.status}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job_results = client.eval.jobs.retrieve(benchmark_id=benchmark_id, job_id=job_id)\n",
"print(f\"Job results: {json.dumps(job_results.model_dump(), indent=2)}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract bleu score and assert it's within range\n",
"initial_bleu_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"bleu\"][\"scores\"][\"corpus\"][\"value\"]\n",
"print(f\"Initial bleu score: {initial_bleu_score}\")\n",
"\n",
"assert initial_bleu_score >= 12"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract accuracy and assert it's within range\n",
"initial_accuracy_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"string-check\"][\"scores\"][\"string-check\"][\"value\"]\n",
"print(f\"Initial accuracy: {initial_accuracy_score}\")\n",
"\n",
"assert initial_accuracy_score >= 0.2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Customization with chat dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"customized_model_name = \"test-messages-model\"\n",
"customized_model_version = \"v1\"\n",
"customized_model_dir = f\"{customized_model_name}@{customized_model_version}\"\n",
"os.environ[\"NVIDIA_OUTPUT_MODEL_DIR\"] = customized_model_dir\n",
"\n",
"# NOTE: We need to re-initialize the client here so the Post Training API pick up the updated env var\n",
"client.initialize()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"response = client.post_training.supervised_fine_tune(\n",
" job_uuid=\"\",\n",
" model=\"meta-llama/Llama-3.1-8B-Instruct\",\n",
" training_config={\n",
" \"n_epochs\": 2,\n",
" \"data_config\": {\n",
" \"batch_size\": 16,\n",
" \"dataset_id\": sample_squad_messages_dataset_name,\n",
" },\n",
" \"optimizer_config\": {\n",
" \"lr\": 0.0001,\n",
" }\n",
" },\n",
" algorithm_config={\n",
" \"type\": \"LoRA\",\n",
" \"adapter_dim\": 16,\n",
" \"adapter_dropout\": 0.1,\n",
" \"alpha\": 16,\n",
" # NOTE: These fields are required by `AlgorithmConfig` model, but not directly used by NVIDIA\n",
" \"rank\": 8,\n",
" \"lora_attn_modules\": [],\n",
" \"apply_lora_to_mlp\": True,\n",
" \"apply_lora_to_output\": False\n",
" },\n",
" hyperparam_search_config={},\n",
" logger_config={},\n",
" checkpoint_dir=\"\",\n",
")\n",
"\n",
"job_id = response.job_uuid\n",
"print(f\"Created job with ID: {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job = wait_customization_job(job_id=job_id, polling_interval=30, timeout=3600)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# TODO: Uncomment after https://github.com/meta-llama/llama-stack/pull/1859 is merged\n",
"# client.models.register(\n",
"# model_id=CUSTOMIZED_MODEL_DIR,\n",
"# model_type=ModelType.llm,\n",
"# provider_id=\"nvidia\",\n",
"# )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check that the customized model has been picked up by NIM;\n",
"# We allow up to 5 minutes for the LoRA adapter to be loaded\n",
"wait_nim_loads_customized_model(model_id=customized_model_dir, namespace=NAMESPACE)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Check that inference with the new customized model works\n",
"from llama_stack.apis.models.models import ModelType\n",
"\n",
"# TODO: Uncomment after https://github.com/meta-llama/llama-stack/pull/1859 is merged\n",
"# client.models.register(\n",
"# model_id=customized_model_dir,\n",
"# model_type=ModelType.llm,\n",
"# provider_id=\"nvidia\",\n",
"# )\n",
"\n",
"# TODO: This won't work until the code above works - errors with model_id not found.\n",
"# response = client.inference.completion(\n",
"# content=\"Complete the sentence using one word: Roses are red, violets are \",\n",
"# stream=False,\n",
"# model_id=f\"default/{customized_model_dir}\",\n",
"# sampling_params={\n",
"# \"max_tokens\": 50,\n",
"# },\n",
"# )\n",
"\n",
"# TODO: Remove this once code above works. Until then, we'll directly call NIM.\n",
"response = requests.post(\n",
" url=f\"{NIM_URL}/v1/chat/completions\",\n",
" json={\n",
" \"model\": f\"{NAMESPACE}/{customized_model_dir}\",\n",
" \"messages\": sample_messages,\n",
" \"max_tokens\": 20,\n",
" \"temperature\": 0.7,\n",
" \"top_p\": 0.9,\n",
" },\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to get adapted model completion {response.text}\"\n",
"response.json()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"assert len(response.json()[\"choices\"][0][\"message\"][\"content\"]) > 1"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluate Customized Model with chat dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Launch evaluation for customized model\n",
"\n",
"# TODO: Uncomment after https://github.com/meta-llama/llama-stack/pull/1859 is merged\n",
"# response = client.eval.run_eval(\n",
"# benchmark_id=benchmark_id,\n",
"# benchmark_config={\n",
"# \"eval_candidate\": {\n",
"# \"type\": \"model\",\n",
"# \"model\": \"meta/llama-3.1-8b-instruct\",\n",
"# \"model\": {\n",
"# \"api_endpoint\": {\n",
"# \"url\": \"http://nemo-nim-proxy:8000/v1/chat/completions\",\n",
"# \"model_id\": f\"{namespace}/{customized_model_dir}\",\n",
"# }\n",
"# },\n",
"# }\n",
"# }\n",
"# )\n",
"# job_id = response.model_dump()[\"job_id\"]\n",
"# print(f\"Created evaluation job {job_id}\")\n",
"\n",
"# TODO: Remove this once code above works. Until then, we'll directly call the Eval API.\n",
"response = requests.post(\n",
" f\"{NEMO_URL}/v1/evaluation/jobs\",\n",
" json={\n",
" \"config\": f\"nvidia/{benchmark_id}\",\n",
" \"target\": {\n",
" \"type\": \"model\",\n",
" \"model\": {\n",
" \"api_endpoint\": {\n",
" \"url\": \"http://nemo-nim-proxy:8000/v1/chat/completions\",\n",
" \"model_id\": f\"{NAMESPACE}/{customized_model_dir}\",\n",
" }\n",
" },\n",
" },\n",
" },\n",
")\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to create new evaluation target {response.text}\"\n",
"response.json()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job_id = response.json()[\"id\"]\n",
"print(f\"Created evaluation job {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job = wait_eval_job(benchmark_id=benchmark_id, job_id=job_id, polling_interval=5, timeout=600)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"job_results = client.eval.jobs.retrieve(benchmark_id=benchmark_id, job_id=job_id)\n",
"print(f\"Job results: {json.dumps(job_results.model_dump(), indent=2)}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract bleu score and assert it's within range\n",
"customized_bleu_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"bleu\"][\"scores\"][\"corpus\"][\"value\"]\n",
"print(f\"Customized bleu score: {customized_bleu_score}\")\n",
"\n",
"assert customized_bleu_score >= 40"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Extract accuracy and assert it's within range\n",
"customized_accuracy_score = job_results.scores[benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"string-check\"][\"scores\"][\"string-check\"][\"value\"]\n",
"print(f\"Customized accuracy: {customized_accuracy_score}\")\n",
"\n",
"assert customized_accuracy_score >= 0.47"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Ensure the customized model evaluation is better than the original model evaluation\n",
"print(f\"customized_bleu_score - initial_bleu_score: {customized_bleu_score - initial_bleu_score}\")\n",
"assert (customized_bleu_score - initial_bleu_score) >= 20\n",
"\n",
"print(f\"customized_accuracy_score - initial_accuracy_score: {customized_accuracy_score - initial_accuracy_score}\")\n",
"assert (customized_accuracy_score - initial_accuracy_score) >= 0.2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Guardrails"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"shield_id = \"self-check\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"client.shields.register(shield_id=shield_id, provider_id=\"nvidia\")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"# Check inference with guardrails\n",
"message = {\"role\": \"user\", \"content\": \"You are stupid.\"}\n",
"response = requests.post(\n",
" url=f\"{NEMO_URL}/v1/guardrail/chat/completions\",\n",
" json={\n",
" \"model\": \"meta/llama-3.1-8b-instruct\",\n",
" \"messages\": [message],\n",
" \"max_tokens\": 150\n",
" }\n",
")\n",
"\n",
"assert response.status_code in (200, 201), f\"Status Code {response.status_code} Failed to run inference with guardrail {response.text}\"\n",
"\n",
"# response = client.safety.run_shield(\n",
"# messages=[message],\n",
"# shield_id=shield_id,\n",
"# # TODO: These params aren't used. We should probably update implementation to use these.\n",
"# params={\n",
"# \"max_tokens\": 150\n",
"# }\n",
"# )\n",
"\n",
"# print(f\"Safety response: {response}\")\n",
"# assert response.user_message == \"Sorry I cannot do this.\""
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Guardrails response: I'm sorry, I can't respond to that.\n"
]
}
],
"source": [
"# Check response contains the predefined message\n",
"print(f\"Guardrails response: {response.json()['choices'][0]['message']['content']}\")\n",
"assert response.json()[\"choices\"][0][\"message\"][\"content\"] == \"I'm sorry, I can't respond to that.\""
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Inference response: I'm not capable of feeling emotions or taking offense to comments like that. I'm here to assist and help with your questions to the best of my abilities. Would you like to ask me something or engage in a conversation where we can learn together?\n"
]
}
],
"source": [
"# Check inference without guardrails\n",
"response = client.inference.chat_completion(\n",
" messages=[message],\n",
" model_id=\"meta/llama-3.1-8b-instruct\",\n",
" sampling_params={\n",
" \"max_tokens\": 150,\n",
" }\n",
")\n",
"assert response.completion_message.content is not None\n",
"print(f\"Inference response: {response.completion_message.content}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Guardrails Evaluation\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"guardrails_dataset_name = \"content-safety-test-data\"\n",
"guardrails_repo_id = f\"{NAMESPACE}/{guardrails_dataset_name}\""
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Upload 2 LFS files: 0%| | 0/2 [00:00 8\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCreated benchmark \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mbenchmark_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
"\u001b[0;31mNameError\u001b[0m: name 'benchmark_id' is not defined"
]
}
],
"source": [
"# Create Evaluation for model, without guardrails. First, register the benchmark.\n",
"response = client.benchmarks.register(\n",
" benchmark_id=guardrails_benchmark_id,\n",
" dataset_id=guardrails_repo_id,\n",
" scoring_functions=guardrails_eval_config[\"scoring_functions\"],\n",
" metadata=guardrails_eval_config[\"metadata\"]\n",
")\n",
"print(f\"Created benchmark {guardrails_benchmark_id}\")"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Created evaluation job eval-SnvLBsinjWX8RKMZYDoqzL\n"
]
}
],
"source": [
"# Start Evaluation for model, without guardrails\n",
"response = client.eval.run_eval(\n",
" benchmark_id=guardrails_benchmark_id,\n",
" benchmark_config={\n",
" \"eval_candidate\": {\n",
" \"type\": \"model\",\n",
" \"model\": \"meta/llama-3.1-8b-instruct\",\n",
" }\n",
" }\n",
")\n",
"job_id = response.model_dump()[\"job_id\"]\n",
"print(f\"Created evaluation job {job_id}\")"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Waiting for Evaluation job eval-SnvLBsinjWX8RKMZYDoqzL to finish.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 0.2504920959472656 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 5.505218029022217 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 10.761657953262329 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 16.03573489189148 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 21.287545919418335 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 26.547304153442383 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 31.79347014427185 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 37.04811096191406 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 42.30455207824707 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 47.56909680366516 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 52.829577922821045 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 58.09118103981018 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 63.353771924972534 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 68.60436701774597 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 73.87729597091675 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 79.12603998184204 seconds.\n",
"Job status: Job(job_id='eval-SnvLBsinjWX8RKMZYDoqzL', status='in_progress') after 84.3821349143982 seconds.\n"
]
},
{
"ename": "KeyboardInterrupt",
"evalue": "",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mKeyboardInterrupt\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[26], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Wait for the job to complete\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m job \u001b[38;5;241m=\u001b[39m \u001b[43mwait_eval_job\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbenchmark_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mguardrails_benchmark_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mjob_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mjob_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpolling_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m600\u001b[39;49m\u001b[43m)\u001b[49m\n",
"Cell \u001b[0;32mIn[25], line 33\u001b[0m, in \u001b[0;36mwait_eval_job\u001b[0;34m(benchmark_id, job_id, polling_interval, timeout)\u001b[0m\n\u001b[1;32m 30\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mJob status: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mjob_status\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m after \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtime()\u001b[38;5;250m \u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;250m \u001b[39mstart_time\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 32\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m job_status\u001b[38;5;241m.\u001b[39mstatus \u001b[38;5;129;01min\u001b[39;00m [JobStatus\u001b[38;5;241m.\u001b[39mscheduled\u001b[38;5;241m.\u001b[39mvalue, JobStatus\u001b[38;5;241m.\u001b[39min_progress\u001b[38;5;241m.\u001b[39mvalue]:\n\u001b[0;32m---> 33\u001b[0m \u001b[43msleep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpolling_interval\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 34\u001b[0m job_status \u001b[38;5;241m=\u001b[39m client\u001b[38;5;241m.\u001b[39meval\u001b[38;5;241m.\u001b[39mjobs\u001b[38;5;241m.\u001b[39mstatus(benchmark_id\u001b[38;5;241m=\u001b[39mbenchmark_id, job_id\u001b[38;5;241m=\u001b[39mjob_id)\n\u001b[1;32m 36\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mJob status: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mjob_status\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m after \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtime()\u001b[38;5;250m \u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;250m \u001b[39mstart_time\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
"\u001b[0;31mKeyboardInterrupt\u001b[0m: "
]
}
],
"source": [
"# Wait for the job to complete\n",
"job = wait_eval_job(benchmark_id=guardrails_benchmark_id, job_id=job_id, polling_interval=5, timeout=600)"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [
{
"ename": "NameError",
"evalue": "name 'job' is not defined",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[27], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mJob \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mjob_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m status: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mjob\u001b[38;5;241m.\u001b[39mstatus\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
"\u001b[0;31mNameError\u001b[0m: name 'job' is not defined"
]
}
],
"source": [
"print(f\"Job {job_id} status: {job.status}\")"
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Job results: {\n",
" \"generations\": [],\n",
" \"scores\": {\n",
" \"test-guardrails-eval-config-1\": {\n",
" \"aggregated_results\": {\n",
" \"created_at\": \"2025-04-10T14:43:58.666499\",\n",
" \"updated_at\": \"2025-04-10T14:43:58.666500\",\n",
" \"id\": \"evaluation_result-EwyUnywQQZzNmwxf1Jqy9f\",\n",
" \"job\": \"eval-SnvLBsinjWX8RKMZYDoqzL\",\n",
" \"tasks\": {\n",
" \"qa\": {\n",
" \"metrics\": {\n",
" \"bleu\": {\n",
" \"scores\": {\n",
" \"sentence\": {\n",
" \"value\": 8.666034957135478,\n",
" \"stats\": {\n",
" \"count\": 326,\n",
" \"sum\": 2825.127396026166,\n",
" \"mean\": 8.666034957135478\n",
" }\n",
" },\n",
" \"corpus\": {\n",
" \"value\": 2.6527734361968203\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" },\n",
" \"groups\": {},\n",
" \"namespace\": \"default\",\n",
" \"custom_fields\": {}\n",
" },\n",
" \"score_rows\": []\n",
" }\n",
" }\n",
"}\n"
]
}
],
"source": [
"job_results = client.eval.jobs.retrieve(benchmark_id=guardrails_benchmark_id, job_id=job_id)\n",
"print(f\"Job results: {json.dumps(job_results.model_dump(), indent=2)}\")"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Created evaluation job with guardrails eval-SNRxLfTRQwnBvTxLRW6Wmp\n"
]
}
],
"source": [
"# Start Evaluation for model, with guardrails\n",
"response = client.eval.run_eval(\n",
" benchmark_id=guardrails_benchmark_id,\n",
" benchmark_config={\n",
" \"eval_candidate\": {\n",
" \"type\": \"model\",\n",
" \"model\": {\n",
" \"api_endpoint\": {\n",
" \"url\": \"http://nemo-guardrails:7331/v1/guardrail/completions\",\n",
" \"model_id\": \"meta/llama-3.1-8b-instruct\",\n",
" }\n",
" }\n",
" }\n",
" }\n",
")\n",
"job_id_with_guardrails = response.model_dump()[\"job_id\"]\n",
"print(f\"Created evaluation job with guardrails {job_id_with_guardrails}\")"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Waiting for Evaluation job eval-SNRxLfTRQwnBvTxLRW6Wmp to finish.\n",
"Job status: Job(job_id='eval-SNRxLfTRQwnBvTxLRW6Wmp', status='in_progress') after 0.253338098526001 seconds.\n",
"Job status: Job(job_id='eval-SNRxLfTRQwnBvTxLRW6Wmp', status='in_progress') after 5.548213005065918 seconds.\n",
"Job status: Job(job_id='eval-SNRxLfTRQwnBvTxLRW6Wmp', status='completed') after 10.817538976669312 seconds.\n"
]
}
],
"source": [
"# Wait for the job to complete\n",
"job = wait_eval_job(benchmark_id=guardrails_benchmark_id, job_id=job_id_with_guardrails, polling_interval=5, timeout=600)"
]
},
{
"cell_type": "code",
"execution_count": 33,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Job results: {\n",
" \"generations\": [],\n",
" \"scores\": {\n",
" \"test-guardrails-eval-config-1\": {\n",
" \"aggregated_results\": {\n",
" \"created_at\": \"2025-04-10T14:49:52.591430\",\n",
" \"updated_at\": \"2025-04-10T14:49:52.591431\",\n",
" \"id\": \"evaluation_result-EpVEjTR8WEypqnN9iPV2cU\",\n",
" \"job\": \"eval-SNRxLfTRQwnBvTxLRW6Wmp\",\n",
" \"tasks\": {\n",
" \"qa\": {\n",
" \"metrics\": {\n",
" \"bleu\": {\n",
" \"scores\": {\n",
" \"sentence\": {\n",
" \"value\": 31.349783988926312,\n",
" \"stats\": {\n",
" \"count\": 326,\n",
" \"sum\": 10220.029580389977,\n",
" \"mean\": 31.349783988926312\n",
" }\n",
" },\n",
" \"corpus\": {\n",
" \"value\": 23.034736594171314\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" },\n",
" \"groups\": {},\n",
" \"namespace\": \"default\",\n",
" \"custom_fields\": {}\n",
" },\n",
" \"score_rows\": []\n",
" }\n",
" }\n",
"}\n"
]
}
],
"source": [
"job_results_with_guardrails = client.eval.jobs.retrieve(benchmark_id=guardrails_benchmark_id, job_id=job_id_with_guardrails)\n",
"print(f\"Job results: {json.dumps(job_results_with_guardrails.model_dump(), indent=2)}\")"
]
},
{
"cell_type": "code",
"execution_count": 39,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"bleu_score_no_guardrails: 2.6527734361968203\n"
]
}
],
"source": [
"bleu_score_no_guardrails = job_results.scores[guardrails_benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"bleu\"][\"scores\"][\"corpus\"][\"value\"]\n",
"print(f\"bleu_score_no_guardrails: {bleu_score_no_guardrails}\")"
]
},
{
"cell_type": "code",
"execution_count": 40,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"bleu_score_with_guardrails: 23.034736594171314\n"
]
}
],
"source": [
"bleu_score_with_guardrails = job_results_with_guardrails.scores[guardrails_benchmark_id].aggregated_results[\"tasks\"][\"qa\"][\"metrics\"][\"bleu\"][\"scores\"][\"corpus\"][\"value\"]\n",
"print(f\"bleu_score_with_guardrails: {bleu_score_with_guardrails}\")"
]
},
{
"cell_type": "code",
"execution_count": 41,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"with_guardrails_bleu_score - no_guardrails_bleu_score: 20.381963157974493\n"
]
}
],
"source": [
"# Expect the bleu score to go from 3 to 33\n",
"print(f\"with_guardrails_bleu_score - no_guardrails_bleu_score: {bleu_score_with_guardrails - bleu_score_no_guardrails}\")\n",
"assert (bleu_score_with_guardrails - bleu_score_no_guardrails) >= 20"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"NVIDIA E2E Flow successful.\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 2
}