nemo/data-flywheel/tool-calling/2_finetuning_and_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "675acc69-4d0c-481c-b093-1f1b849019fa",
   "metadata": {},
   "source": [
    "# Part II: LoRA Fine-tuning Using NeMo Customizer\n",
    "\n",
    "This notebook covers the following:\n",
    "\n",
    "0. [Prerequisites: Configurations, Health Checks, and Namespaces](#step-0)\n",
    "1. [Upload Data to NeMo Datastore](#step-1)\n",
    "2. [LoRA Customization with NeMo Customizer](#step-2)\n",
    "3. [Running Inference on the Customized Model with NVIDIA NIM](#step-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "96adf453-5f61-435c-9ed2-257cdcebe24e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import requests\n",
    "import random\n",
    "from openai import OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e84c255-d91b-4ee8-9db2-61634ab70179",
   "metadata": {},
   "source": [
    "<a id=\"step-0\"></a>\n",
    "## Prerequisites: Configurations, Health Checks, and Namespaces"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b175f7f-4ada-465e-ac84-e40aac6abb38",
   "metadata": {},
   "source": [
    "Before you proceed, make sure that you completed the first notebook on data preparation to obtain the assets required to follow along."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fd6ee1a-0789-4724-8f12-5157363ebf9c",
   "metadata": {},
   "source": [
    "### Configure NeMo Microservices Endpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa34fe25",
   "metadata": {},
   "source": [
    "This section includes importing required libraries, configuring endpoints, and performing health checks to ensure that the NeMo Data Store, NIM, and other services are running correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d29578e1-6757-4e89-8f29-a0f6739015c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data Store endpoint: http://data-store.test\n",
      "Entity Store, Customizer, Evaluator endpoint: http://nemo.test\n",
      "NIM endpoint: http://nim.test\n",
      "Namespace: xlam-tutorial-ns\n",
      "Base Model for Customization: meta/llama-3.2-1b-instruct\n"
     ]
    }
   ],
   "source": [
    "from config import *\n",
    "\n",
    "print(f\"Data Store endpoint: {NDS_URL}\")\n",
    "print(f\"Entity Store, Customizer, Evaluator endpoint: {NEMO_URL}\")\n",
    "print(f\"NIM endpoint: {NIM_URL}\")\n",
    "print(f\"Namespace: {NMS_NAMESPACE}\")\n",
    "print(f\"Base Model for Customization: {BASE_MODEL}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1b039f7-f07d-4189-add1-d47941a4d78e",
   "metadata": {},
   "source": [
    "### Configure Path to Prepared data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa89d41",
   "metadata": {},
   "source": [
    "The following code sets the paths to the prepared dataset files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "26e8c768-498e-47df-b8c4-805b25b63d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path where data preparation notebook saved finetuning and evaluation data\n",
    "DATA_ROOT = os.path.join(os.getcwd(), \"data\")\n",
    "CUSTOMIZATION_DATA_ROOT = os.path.join(DATA_ROOT, \"customization\")\n",
    "VALIDATION_DATA_ROOT = os.path.join(DATA_ROOT, \"validation\")\n",
    "EVALUATION_DATA_ROOT = os.path.join(DATA_ROOT, \"evaluation\")\n",
    "\n",
    "# Sanity checks\n",
    "train_fp = f\"{CUSTOMIZATION_DATA_ROOT}/training.jsonl\"\n",
    "assert os.path.exists(train_fp), f\"The training data at '{train_fp}' does not exist. Please ensure that the data was prepared successfully.\"\n",
    "\n",
    "val_fp = f\"{VALIDATION_DATA_ROOT}/validation.jsonl\"\n",
    "assert os.path.exists(val_fp), f\"The validation data at '{val_fp}' does not exist. Please ensure that the data was prepared successfully.\"\n",
    "\n",
    "test_fp = f\"{EVALUATION_DATA_ROOT}/xlam-test-single.jsonl\"\n",
    "assert os.path.exists(test_fp), f\"The test data at '{test_fp}' does not exist. Please ensure that the data was prepared successfully.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "062c7805-727f-4265-9248-5babc4a32fc5",
   "metadata": {},
   "source": [
    "### Resource Organization Using Namespace\n",
    "\n",
    "You can use a [namespace](https://developer.nvidia.com/docs/nemo-microservices/manage-entities/namespaces/index.html) to isolate and organize the artifacts in this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e57ea9a-05b0-4e5e-9455-fc7c2c515ca1",
   "metadata": {},
   "source": [
    "#### Create Namespace\n",
    "\n",
    "Both Data Store and Entity Store use namespaces. The following code creates namespaces for the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4dff1528-2c7c-4118-8921-f54167ba57ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Response [200]>\n",
      "<Response [409]>\n"
     ]
    }
   ],
   "source": [
    "def create_namespaces(entity_host, ds_host, namespace):\n",
    "    # Create namespace in Entity Store\n",
    "    entity_store_url = f\"{entity_host}/v1/namespaces\"\n",
    "    resp = requests.post(entity_store_url, json={\"id\": namespace})\n",
    "    assert resp.status_code in (200, 201, 409, 422), \\\n",
    "        f\"Unexpected response from Entity Store during namespace creation: {resp.status_code}\"\n",
    "    print(resp)\n",
    "\n",
    "    # Create namespace in Data Store\n",
    "    nds_url = f\"{ds_host}/v1/datastore/namespaces\"\n",
    "    resp = requests.post(nds_url, data={\"namespace\": namespace})\n",
    "    assert resp.status_code in (200, 201, 409, 422), \\\n",
    "        f\"Unexpected response from Data Store during namespace creation: {resp.status_code}\"\n",
    "    print(resp)\n",
    "\n",
    "create_namespaces(entity_host=NEMO_URL, ds_host=NDS_URL, namespace=NMS_NAMESPACE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "057acae8-4d0b-48ee-bd8f-6f98bd941ba6",
   "metadata": {},
   "source": [
    "#### Verify Namespaces\n",
    "\n",
    "The following [Data Store API](https://developer.nvidia.com/docs/nemo-microservices/api/datastore.html) and [Entity Store API](https://developer.nvidia.com/docs/nemo-microservices/api/entity-store.html) list the namespace created in the previous cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "413166f0-a314-4db0-a341-914495b583ef",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status Code: 201\n",
      "Response JSON: {'namespace': 'xlam-tutorial-ns', 'created_at': '2025-04-07T23:37:44Z', 'updated_at': '2025-04-08T07:03:53Z'}\n",
      "Status Code: 200\n",
      "Response JSON: {'id': 'xlam-tutorial-ns', 'created_at': '2025-04-08T07:05:07.362008', 'updated_at': '2025-04-08T07:05:07.362012', 'description': None, 'project': None, 'custom_fields': {}, 'ownership': None}\n"
     ]
    }
   ],
   "source": [
    "# Verify Namespace in Data Store\n",
    "response = requests.get(f\"{NDS_URL}/v1/datastore/namespaces/{NMS_NAMESPACE}\")\n",
    "print(f\"Status Code: {response.status_code}\\nResponse JSON: {response.json()}\")\n",
    "\n",
    "# Verify Namespace in Entity Store\n",
    "response = requests.get(f\"{NEMO_URL}/v1/namespaces/{NMS_NAMESPACE}\")\n",
    "print(f\"Status Code: {response.status_code}\\nResponse JSON: {response.json()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556e6df9-bb6b-4606-80da-e936f68a5d23",
   "metadata": {},
   "source": [
    "**Tips**:\n",
    "* You may generally use `{DATASTORE_HOST}/v1/datastore/namespaces/` and `{ENTITYSTORE_HOST}/v1/namespaces/` GET APIs to list **all** available namespaces.\n",
    "* Send DELETE requests to `{DATASTORE_HOST}/v1/datastore/namespaces/{namespace}` and `{ENTITYSTORE_HOST}/v1/namespaces/{namespace}` APIs to delete a namespace."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2e83d04-e092-49bd-8769-3422526e135b",
   "metadata": {},
   "source": [
    "---\n",
    "<a id=\"step-1\"></a>\n",
    "## Step 1: Upload Data to NeMo Data Store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efbdb147-9994-4e08-a742-38098f4dc911",
   "metadata": {},
   "source": [
    "The NeMo Data Store supports data management using the Hugging Face `HfApi` Client. \n",
    "\n",
    "**Note that this step does not interact with Hugging Face at all, it just uses the client library to interact with NeMo Data Store.** This is in comparison to the previous notebook, where we used the `load_dataset` API to download the xLAM dataset from Hugging Face's repository.\n",
    "\n",
    "More information can be found in [documentation](https://developer.nvidia.com/docs/nemo-microservices/manage-entities/tutorials/manage-dataset-files.html#set-up-hugging-face-client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e29d515c-bd7d-4fdb-9c6d-42370d975e69",
   "metadata": {},
   "source": [
    "### 1.1 Create Repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c1a1178a-b516-4f50-96a3-2dee114f9cda",
   "metadata": {},
   "outputs": [],
   "source": [
    "repo_id = f\"{NMS_NAMESPACE}/{DATASET_NAME}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bd240212-3272-419e-b19b-d31d8aede2b3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-28T14:02:46.351447Z",
     "start_time": "2025-02-28T14:02:46.343156Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RepoUrl('datasets/xlam-tutorial-ns/xlam-ft-dataset', endpoint='http://data-store.test/v1/hf', repo_type='dataset', repo_id='xlam-tutorial-ns/xlam-ft-dataset')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "hf_api = HfApi(endpoint=f\"{NDS_URL}/v1/hf\", token=\"\")\n",
    "\n",
    "# Create repo\n",
    "hf_api.create_repo(\n",
    "    repo_id=repo_id,\n",
    "    repo_type='dataset',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97ac352a-31b9-4144-ad0f-699fcceebfc2",
   "metadata": {},
   "source": [
    "Next, creating a dataset programmatically requires two steps: uploading and registration. More information can be found in [documentation](https://developer.nvidia.com/docs/nemo-microservices/manage-entities/datasets/create-dataset.html#how-to-create-a-dataset)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f927b0-b216-46f9-a8d0-0f9d6836868f",
   "metadata": {},
   "source": [
    "### 1.2 Upload Dataset Files to NeMo Data Store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "50be5cec-5df6-4092-82ff-29d39aceaa1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ff5b5fb2a7d74c1d8017080e76c426c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "training.jsonl:   0%|          | 0.00/6.06M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9215b2796df4cc783364896e2f8cfc0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "validation.jsonl:   0%|          | 0.00/1.30M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26b794a398fe4c7fbf3571b9bacf778a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "xlam-test-single.jsonl:   0%|          | 0.00/1.19M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='', commit_message='Upload testing/xlam-test-single.jsonl with huggingface_hub', commit_description='', oid='f3f1705c7652aaa06823693293c44ce228884a85', pr_url=None, repo_url=RepoUrl('', endpoint='https://huggingface.co', repo_type='model', repo_id=''), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_api.upload_file(path_or_fileobj=train_fp,\n",
    "    path_in_repo=\"training/training.jsonl\",\n",
    "    repo_id=repo_id,\n",
    "    repo_type='dataset',\n",
    ")\n",
    "\n",
    "hf_api.upload_file(path_or_fileobj=val_fp,\n",
    "    path_in_repo=\"validation/validation.jsonl\",\n",
    "    repo_id=repo_id,\n",
    "    repo_type='dataset',\n",
    ")\n",
    "\n",
    "hf_api.upload_file(path_or_fileobj=test_fp,\n",
    "    path_in_repo=\"testing/xlam-test-single.jsonl\",\n",
    "    repo_id=repo_id,\n",
    "    repo_type='dataset',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e5d0ab1-69b2-4ca5-9a25-d514cf72b7cd",
   "metadata": {},
   "source": [
    "Other tips:\n",
    "* Take a look at the `path_in_repo` argument above. If there are more than one files in the subfolders:\n",
    "    * All the .jsonl files in `training/` will be merged and used for training by customizer.\n",
    "    * All the .jsonl files in `validation/` will be merged and used for validation by customizer.\n",
    "* NeMo Data Store generally supports data management using the [HfApi API](https://huggingface.co/docs/huggingface_hub/en/package_reference/hf_api). For example, to delete a repo, you may use - \n",
    "```python\n",
    "   hf_api.delete_repo(\n",
    "     repo_id=repo_id,\n",
    "     repo_type=\"dataset\"\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a371c044-56df-412d-9fb3-7e0e191dd3a8",
   "metadata": {},
   "source": [
    "### 1.3 Register the Dataset with NeMo Entity Store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66ab2d1d-17d5-4c7a-9b1b-8f8b1c68204c",
   "metadata": {},
   "source": [
    "To use a dataset for operations such as evaluations and customizations, register a dataset using the `/v1/datasets` endpoint.\n",
    "Register the dataset to refer to it by its namespace and name afterward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7a45fadd-df51-48e0-b30c-336c5bca071a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'created_at': '2025-04-08T07:05:42.895214',\n",
       " 'updated_at': '2025-04-08T07:05:42.895217',\n",
       " 'name': 'xlam-ft-dataset',\n",
       " 'namespace': 'xlam-tutorial-ns',\n",
       " 'description': 'Tool calling xLAM dataset in OpenAI ChatCompletions format',\n",
       " 'format': None,\n",
       " 'files_url': 'hf://datasets/xlam-tutorial-ns/xlam-ft-dataset',\n",
       " 'hf_endpoint': None,\n",
       " 'split': None,\n",
       " 'limit': None,\n",
       " 'id': 'dataset-3ozLsXkX7TqvQm9fCuCUGT',\n",
       " 'project': 'tool_calling',\n",
       " 'custom_fields': {}}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resp = requests.post(\n",
    "    url=f\"{NEMO_URL}/v1/datasets\",\n",
    "    json={\n",
    "        \"name\": DATASET_NAME,\n",
    "        \"namespace\": NMS_NAMESPACE,\n",
    "        \"description\": \"Tool calling xLAM dataset in OpenAI ChatCompletions format\",\n",
    "        \"files_url\": f\"hf://datasets/{NMS_NAMESPACE}/{DATASET_NAME}\",\n",
    "        \"project\": \"tool_calling\",\n",
    "    },\n",
    ")\n",
    "assert resp.status_code in (200, 201), f\"Status Code {resp.status_code} Failed to create dataset {resp.text}\"\n",
    "resp.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8a2b58c7-30d7-4c04-b3a5-872032dbf4f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files URL: hf://datasets/xlam-tutorial-ns/xlam-ft-dataset\n"
     ]
    }
   ],
   "source": [
    "# Sanity check to validate dataset\n",
    "res = requests.get(url=f\"{NEMO_URL}/v1/datasets/{NMS_NAMESPACE}/{DATASET_NAME}\")\n",
    "assert res.status_code in (200, 201), f\"Status Code {res.status_code} Failed to fetch dataset {res.text}\"\n",
    "dataset_obj = res.json()\n",
    "\n",
    "print(\"Files URL:\", dataset_obj[\"files_url\"])\n",
    "assert dataset_obj[\"files_url\"] == f\"hf://datasets/{repo_id}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e7cebbc-8a5f-492d-820c-b3fbbed6fafb",
   "metadata": {},
   "source": [
    "---\n",
    "<a id=\"step-2\"></a>\n",
    "## 2. LoRA Customization with NeMo Customizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6915c1f5-9687-443f-92cd-cd30b55fc501",
   "metadata": {},
   "source": [
    "### 2.1 Start the Training Job"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e22fa48",
   "metadata": {},
   "source": [
    "\n",
    "Start the training job by sending a POST request to the `/v1/customization/jobs` endpoint.\n",
    "The following code sets the training parameters and sends the request.\n",
    "\n",
    " **The training job will take approximately 45 minutes to complete.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a2e9678c-2785-4e95-b11b-1f41067bc920",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'cust-BTkGbfifLfEAjV2THu3tas',\n",
       " 'created_at': '2025-04-08T07:05:52.236823',\n",
       " 'updated_at': '2025-04-08T07:05:52.236829',\n",
       " 'namespace': 'default',\n",
       " 'dataset': 'xlam-tutorial-ns/xlam-ft-dataset',\n",
       " 'output_model': 'xlam-tutorial-ns/llama-3.2-1b-xlam-run1@cust-BTkGbfifLfEAjV2THu3tas',\n",
       " 'config': {'base_model': 'meta/llama-3.2-1b-instruct',\n",
       "  'precision': 'bf16-mixed',\n",
       "  'num_gpus': 1,\n",
       "  'num_nodes': 1,\n",
       "  'micro_batch_size': 1,\n",
       "  'tensor_parallel_size': 1,\n",
       "  'max_seq_length': 4096,\n",
       "  'prompt_template': '{prompt} {completion}'},\n",
       " 'hyperparameters': {'finetuning_type': 'lora',\n",
       "  'training_type': 'sft',\n",
       "  'batch_size': 16,\n",
       "  'epochs': 2,\n",
       "  'learning_rate': 0.0001,\n",
       "  'lora': {'adapter_dim': 32, 'alpha': 16, 'adapter_dropout': 0.1},\n",
       "  'sequence_packing_enabled': False},\n",
       " 'status': 'created',\n",
       " 'status_details': {'created_at': '2025-04-08T07:05:53.328702',\n",
       "  'updated_at': '2025-04-08T07:05:53.328702',\n",
       "  'steps_completed': 0,\n",
       "  'epochs_completed': 0,\n",
       "  'percentage_done': 0.0,\n",
       "  'status_logs': [{'updated_at': '2025-04-08T07:05:53.328702',\n",
       "    'message': 'created'}]}}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "headers = {\"wandb-api-key\": WANDB_API_KEY} if WANDB_API_KEY else None\n",
    "\n",
    "training_params = {\n",
    "    \"name\": \"llama-3.2-1b-xlam-ft\",\n",
    "    \"output_model\": f\"{NMS_NAMESPACE}/llama-3.2-1b-xlam-run1\",\n",
    "    \"config\": BASE_MODEL,\n",
    "    \"dataset\": {\"name\": DATASET_NAME, \"namespace\" : NMS_NAMESPACE},\n",
    "    \"hyperparameters\": {\n",
    "        \"training_type\": \"sft\",\n",
    "        \"finetuning_type\": \"lora\",\n",
    "        \"epochs\": 2,\n",
    "        \"batch_size\": 16,\n",
    "        \"learning_rate\": 0.0001,\n",
    "        \"lora\": {\n",
    "            \"adapter_dim\": 32,\n",
    "            \"adapter_dropout\": 0.1\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "resp = requests.post(f\"{NEMO_URL}/v1/customization/jobs\", json=training_params, headers=headers)\n",
    "customization = resp.json()\n",
    "customization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69131ef8",
   "metadata": {},
   "source": [
    "The following code sets variables for storing the job ID and customized model name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "57eb5ae6-9b3e-4915-8242-34b65b0c0680",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To track status\n",
    "JOB_ID = customization[\"id\"]\n",
    "\n",
    "# This will be the name of the model that will be used to send inference queries to\n",
    "CUSTOMIZED_MODEL = customization[\"output_model\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9169153d-2ff0-4b34-9a3d-236ecedb7a5d",
   "metadata": {},
   "source": [
    "**Tips**:\n",
    "* If you configured the NeMo Customizer microservice with your own [Weights & Biases (WandB)](https://wandb.ai/) API key, you can find the training graphs and logs in your WandB account, \"nvidia-nemo-customizer\" project. Your run ID is similar to your customization `JOB_ID`.\n",
    "  \n",
    "* To cancel a job that you scheduled incorrectly, run the following code.\n",
    "  \n",
    "  ```python\n",
    "  requests.post(f\"{NEMO_URL}/v1/customization/jobs/{JOB_ID}/cancel\")\n",
    "  ```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ab1b3b7-5e0d-4bbb-a7cb-476e0abb298b",
   "metadata": {},
   "source": [
    "### 2.2 Get Job Status\n",
    "\n",
    "Get the job status by sending a GET request to the `/v1/customization/jobs/{JOB_ID}/status` endpoint.\n",
    "The following code sets the job ID and sends the request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "900f28fb-fb8e-4d57-88d7-6c09699523e2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response JSON: {\n",
      "    \"created_at\": \"2025-04-08T07:05:53.328702\",\n",
      "    \"updated_at\": \"2025-04-08T07:26:36.749083\",\n",
      "    \"status\": \"completed\",\n",
      "    \"steps_completed\": 438,\n",
      "    \"epochs_completed\": 2,\n",
      "    \"percentage_done\": 100.0,\n",
      "    \"best_epoch\": 1,\n",
      "    \"train_loss\": 0.052824027836322784,\n",
      "    \"val_loss\": 0.047389596700668335,\n",
      "    \"metrics\": {\n",
      "        \"keys\": [\n",
      "            \"train_loss\",\n",
      "            \"val_loss\"\n",
      "        ],\n",
      "        \"metrics\": {\n",
      "            \"train_loss\": [\n",
      "                {\n",
      "                    \"value\": 1.7395085096359253,\n",
      "                    \"step\": 9,\n",
      "                    \"timestamp\": \"2025-04-08T07:07:50.318152\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.5097759962081909,\n",
      "                    \"step\": 19,\n",
      "                    \"timestamp\": \"2025-04-08T07:08:21.445975\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.11785753816366196,\n",
      "                    \"step\": 29,\n",
      "                    \"timestamp\": \"2025-04-08T07:08:45.141126\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.044126059859991074,\n",
      "                    \"step\": 39,\n",
      "                    \"timestamp\": \"2025-04-08T07:09:08.064725\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.16011953353881836,\n",
      "                    \"step\": 49,\n",
      "                    \"timestamp\": \"2025-04-08T07:09:27.849022\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.060752492398023605,\n",
      "                    \"step\": 59,\n",
      "                    \"timestamp\": \"2025-04-08T07:09:49.829851\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.061382267624139786,\n",
      "                    \"step\": 69,\n",
      "                    \"timestamp\": \"2025-04-08T07:10:07.325047\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.01653873734176159,\n",
      "                    \"step\": 79,\n",
      "                    \"timestamp\": \"2025-04-08T07:10:26.559717\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.06011583283543587,\n",
      "                    \"step\": 89,\n",
      "                    \"timestamp\": \"2025-04-08T07:10:42.204168\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.06969955563545227,\n",
      "                    \"step\": 99,\n",
      "                    \"timestamp\": \"2025-04-08T07:10:58.948349\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.08384404331445694,\n",
      "                    \"step\": 109,\n",
      "                    \"timestamp\": \"2025-04-08T07:11:26.313792\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.06991764903068542,\n",
      "                    \"step\": 119,\n",
      "                    \"timestamp\": \"2025-04-08T07:11:43.883420\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.03866000473499298,\n",
      "                    \"step\": 129,\n",
      "                    \"timestamp\": \"2025-04-08T07:12:11.505720\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.05358102172613144,\n",
      "                    \"step\": 139,\n",
      "                    \"timestamp\": \"2025-04-08T07:12:41.459384\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.07019069045782089,\n",
      "                    \"step\": 149,\n",
      "                    \"timestamp\": \"2025-04-08T07:13:01.654562\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.07882250845432281,\n",
      "                    \"step\": 159,\n",
      "                    \"timestamp\": \"2025-04-08T07:13:28.198781\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.057352885603904724,\n",
      "                    \"step\": 169,\n",
      "                    \"timestamp\": \"2025-04-08T07:13:55.679940\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.048704877495765686,\n",
      "                    \"step\": 179,\n",
      "                    \"timestamp\": \"2025-04-08T07:14:16.962284\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.11735919862985611,\n",
      "                    \"step\": 189,\n",
      "                    \"timestamp\": \"2025-04-08T07:14:47.678515\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.030179649591445923,\n",
      "                    \"step\": 199,\n",
      "                    \"timestamp\": \"2025-04-08T07:15:09.066542\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.027945276349782944,\n",
      "                    \"step\": 209,\n",
      "                    \"timestamp\": \"2025-04-08T07:15:29.159765\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.04044452682137489,\n",
      "                    \"step\": 219,\n",
      "                    \"timestamp\": \"2025-04-08T07:17:11.140956\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.047244369983673096,\n",
      "                    \"step\": 229,\n",
      "                    \"timestamp\": \"2025-04-08T07:17:31.705232\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.04799420014023781,\n",
      "                    \"step\": 239,\n",
      "                    \"timestamp\": \"2025-04-08T07:18:00.639729\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.017997125163674355,\n",
      "                    \"step\": 249,\n",
      "                    \"timestamp\": \"2025-04-08T07:18:18.967952\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.05586077272891998,\n",
      "                    \"step\": 259,\n",
      "                    \"timestamp\": \"2025-04-08T07:18:37.262997\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.024363242089748383,\n",
      "                    \"step\": 269,\n",
      "                    \"timestamp\": \"2025-04-08T07:18:55.515872\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.007903068326413631,\n",
      "                    \"step\": 279,\n",
      "                    \"timestamp\": \"2025-04-08T07:19:15.973629\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.04577765613794327,\n",
      "                    \"step\": 289,\n",
      "                    \"timestamp\": \"2025-04-08T07:19:38.154098\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.05233767256140709,\n",
      "                    \"step\": 299,\n",
      "                    \"timestamp\": \"2025-04-08T07:19:58.773278\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.02214895375072956,\n",
      "                    \"step\": 309,\n",
      "                    \"timestamp\": \"2025-04-08T07:20:31.187366\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.03232814371585846,\n",
      "                    \"step\": 319,\n",
      "                    \"timestamp\": \"2025-04-08T07:20:53.427308\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.07305795699357986,\n",
      "                    \"step\": 329,\n",
      "                    \"timestamp\": \"2025-04-08T07:21:14.126436\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.010019486770033836,\n",
      "                    \"step\": 339,\n",
      "                    \"timestamp\": \"2025-04-08T07:21:45.127167\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.005295319017022848,\n",
      "                    \"step\": 349,\n",
      "                    \"timestamp\": \"2025-04-08T07:22:04.148566\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.006320458836853504,\n",
      "                    \"step\": 359,\n",
      "                    \"timestamp\": \"2025-04-08T07:22:26.461070\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.005519861355423927,\n",
      "                    \"step\": 369,\n",
      "                    \"timestamp\": \"2025-04-08T07:22:56.395855\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.024514716118574142,\n",
      "                    \"step\": 379,\n",
      "                    \"timestamp\": \"2025-04-08T07:23:14.846578\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.006911200471222401,\n",
      "                    \"step\": 389,\n",
      "                    \"timestamp\": \"2025-04-08T07:23:40.472405\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.02401524968445301,\n",
      "                    \"step\": 399,\n",
      "                    \"timestamp\": \"2025-04-08T07:24:12.313296\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.0362529456615448,\n",
      "                    \"step\": 409,\n",
      "                    \"timestamp\": \"2025-04-08T07:24:43.646847\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.03379900008440018,\n",
      "                    \"step\": 419,\n",
      "                    \"timestamp\": \"2025-04-08T07:25:14.787556\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.01118356455117464,\n",
      "                    \"step\": 429,\n",
      "                    \"timestamp\": \"2025-04-08T07:25:43.727954\"\n",
      "                }\n",
      "            ],\n",
      "            \"val_loss\": [\n",
      "                {\n",
      "                    \"value\": 0.047389596700668335,\n",
      "                    \"step\": 218,\n",
      "                    \"timestamp\": \"2025-04-08T07:16:54.402277\"\n",
      "                },\n",
      "                {\n",
      "                    \"value\": 0.04571458697319031,\n",
      "                    \"step\": 437,\n",
      "                    \"timestamp\": \"2025-04-08T07:26:36.013749\"\n",
      "                }\n",
      "            ]\n",
      "        }\n",
      "    },\n",
      "    \"status_logs\": [\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:05:53\",\n",
      "            \"message\": \"PVCCreated\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:05:53\",\n",
      "            \"message\": \"EntityHandler_0_Created\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:05:53.328702\",\n",
      "            \"message\": \"created\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:06:04\",\n",
      "            \"message\": \"EntityHandler_0_Pending\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:06:04\",\n",
      "            \"message\": \"EntityHandler_0_Completed\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:06:04\",\n",
      "            \"message\": \"TrainingJobCreated\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:06:04\",\n",
      "            \"message\": \"TrainingJobRunning\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:27:02\",\n",
      "            \"message\": \"TrainingJobCompleted\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:27:02\",\n",
      "            \"message\": \"EntityHandler_1_Created\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:27:02\",\n",
      "            \"message\": \"EntityHandler_1_Running\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:27:10\",\n",
      "            \"message\": \"EntityHandler_1_Pending\",\n",
      "            \"detail\": null\n",
      "        },\n",
      "        {\n",
      "            \"updated_at\": \"2025-04-08T07:27:10\",\n",
      "            \"message\": \"EntityHandler_1_Completed\",\n",
      "            \"detail\": null\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "response = requests.get(f\"{NEMO_URL}/v1/customization/jobs/{JOB_ID}/status\")\n",
    "\n",
    "assert response.status_code == 200, (\n",
    "    f\"Status Code {response.status_code}: Failed to get job status. Response: {response.text}\"\n",
    ")\n",
    "print(\"Response JSON:\", json.dumps(response.json(), indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b721be-8ca0-4e8f-99a7-5eb12ea1b47f",
   "metadata": {},
   "source": [
    "**IMPORTANT:** Monitor the job status. Ensure training is completed before proceeding by observing the `percentage_done` key in the response frame."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec56031-ec55-4cec-9f65-b34b87818e01",
   "metadata": {},
   "source": [
    "### 2.3 Validate Availability of Custom Model\n",
    "The following NeMo Entity Store API should display the model when the training job is complete.\n",
    "The list below shows all models filtered by your namespace and sorted by the latest first.\n",
    "For more information about this API, see the [NeMo Entity Store API reference](https://developer.nvidia.com/docs/nemo-microservices/api/entity-store.html).\n",
    "With the following code, you can find all customized models, including the one trained in the previous cells.\n",
    "Look for the `name` fields in the output, which should match your `CUSTOMIZED_MODEL`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ebad3944-70d5-4b23-9a38-a83774bf20c4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response JSON: {\n",
      "    \"object\": \"list\",\n",
      "    \"data\": [\n",
      "        {\n",
      "            \"created_at\": \"2025-04-08T07:05:53.417350\",\n",
      "            \"updated_at\": \"2025-04-08T07:05:53.417354\",\n",
      "            \"name\": \"llama-3.2-1b-xlam-run1@cust-BTkGbfifLfEAjV2THu3tas\",\n",
      "            \"namespace\": \"xlam-tutorial-ns\",\n",
      "            \"description\": \"None\",\n",
      "            \"spec\": {\n",
      "                \"num_parameters\": 1000000000,\n",
      "                \"context_size\": 4096,\n",
      "                \"num_virtual_tokens\": 0,\n",
      "                \"is_chat\": true\n",
      "            },\n",
      "            \"artifact\": {\n",
      "                \"gpu_arch\": \"Ampere\",\n",
      "                \"precision\": \"bf16-mixed\",\n",
      "                \"tensor_parallelism\": 1,\n",
      "                \"backend_engine\": \"nemo\",\n",
      "                \"status\": \"upload_completed\",\n",
      "                \"files_url\": \"hf://xlam-tutorial-ns/llama-3.2-1b-xlam-run1@cust-BTkGbfifLfEAjV2THu3tas\"\n",
      "            },\n",
      "            \"base_model\": \"meta/llama-3.2-1b-instruct\",\n",
      "            \"peft\": {\n",
      "                \"finetuning_type\": \"lora\"\n",
      "            },\n",
      "            \"schema_version\": \"1.0\",\n",
      "            \"custom_fields\": {}\n",
      "        },\n",
      "        {\n",
      "            \"created_at\": \"2025-04-07T23:41:27.787234\",\n",
      "            \"updated_at\": \"2025-04-07T23:41:27.787238\",\n",
      "            \"name\": \"llama-3.2-1b-xlam-run1@cust-ULiqFRom2AnFEUsVpeoLED\",\n",
      "            \"namespace\": \"xlam-tutorial-ns\",\n",
      "            \"description\": \"None\",\n",
      "            \"spec\": {\n",
      "                \"num_parameters\": 1000000000,\n",
      "                \"context_size\": 4096,\n",
      "                \"num_virtual_tokens\": 0,\n",
      "                \"is_chat\": true\n",
      "            },\n",
      "            \"artifact\": {\n",
      "                \"gpu_arch\": \"Ampere\",\n",
      "                \"precision\": \"bf16-mixed\",\n",
      "                \"tensor_parallelism\": 1,\n",
      "                \"backend_engine\": \"nemo\",\n",
      "                \"status\": \"upload_completed\",\n",
      "                \"files_url\": \"hf://xlam-tutorial-ns/llama-3.2-1b-xlam-run1@cust-ULiqFRom2AnFEUsVpeoLED\"\n",
      "            },\n",
      "            \"base_model\": \"meta/llama-3.2-1b-instruct\",\n",
      "            \"peft\": {\n",
      "                \"finetuning_type\": \"lora\"\n",
      "            },\n",
      "            \"schema_version\": \"1.0\",\n",
      "            \"custom_fields\": {}\n",
      "        }\n",
      "    ],\n",
      "    \"pagination\": {\n",
      "        \"page\": 1,\n",
      "        \"page_size\": 1000,\n",
      "        \"current_page_size\": 2,\n",
      "        \"total_pages\": 1,\n",
      "        \"total_results\": 2\n",
      "    },\n",
      "    \"sort\": \"-created_at\",\n",
      "    \"filter\": {\n",
      "        \"namespace\": \"xlam-tutorial-ns\"\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "response = requests.get(f\"{NEMO_URL}/v1/models\", params={\"filter[namespace]\": NMS_NAMESPACE, \"sort\" : \"-created_at\"})\n",
    "\n",
    "assert response.status_code == 200, f\"Status Code {response.status_code}: Request failed. Response: {response.text}\"\n",
    "print(\"Response JSON:\", json.dumps(response.json(), indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1221f1e-04f6-4ceb-a844-e9ea1e94f0d0",
   "metadata": {},
   "source": [
    "**Tips**:\n",
    "\n",
    "* You can also find the model with its name directly:\n",
    "  ```python\n",
    "    # To get specifically the custom model, you may use the following API -\n",
    "    response = requests.get(f\"{NEMO_URL}/v1/models/{CUSTOMIZED_MODEL}\")\n",
    "    \n",
    "    assert response.status_code == 200, f\"Status Code {response.status_code}: Request failed. Response: {response.text}\"\n",
    "    print(\"Response JSON:\", json.dumps(response.json(), indent=4))\n",
    "  ```\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b43d7dc-f442-41cb-a743-7fd5efd4bf6c",
   "metadata": {},
   "source": [
    "NVIDIA NIM directly picks up the LoRA adapters from NeMo Entity Store. You can also query the NIM endpoint to look for it, as shown in the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "43026f8a-3b98-4aa6-b4c6-7441862863fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the custom LoRA model is hosted by NVIDIA NIM\n",
    "resp = requests.get(f\"{NIM_URL}/v1/models\")\n",
    "\n",
    "models = resp.json().get(\"data\", [])\n",
    "model_names = [model[\"id\"] for model in models]\n",
    "\n",
    "assert CUSTOMIZED_MODEL in model_names, \\\n",
    "    f\"Model {CUSTOMIZED_MODEL} not found\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4df0f9c5-98b2-476e-aa2f-3af69c5e45f7",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "<a id=\"step-3\"></a>\n",
    "## Step 3: Sanity Test the Customized Model By Running Sample Inference\n",
    "\n",
    "Once the model is customized, its adapter is automatically saved in NeMo Entity Store and is ready to be picked up by NVIDIA NIM.\n",
    "You can test the model by sending a prompt to its NIM endpoint.\n",
    "\n",
    "First, choose one of the examples from the test set."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf4ada2-2ccb-4f82-bb7c-23134e67a7e0",
   "metadata": {},
   "source": [
    "### 3.1 Get Test Data Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c3f7da82-cbbb-43a4-999b-2adad2cea227",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 713 examples in the test set\n"
     ]
    }
   ],
   "source": [
    "def read_jsonl(file_path):\n",
    "    \"\"\"Reads a JSON Lines file and yields parsed JSON objects\"\"\"\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        for line in file:\n",
    "            line = line.strip()  # Remove leading/trailing whitespace\n",
    "            if not line:\n",
    "                continue  # Skip empty lines\n",
    "            try:\n",
    "                yield json.loads(line)\n",
    "            except json.JSONDecodeError as e:\n",
    "                print(f\"Error decoding JSON: {e}\")\n",
    "                continue\n",
    "\n",
    "\n",
    "test_data = list(read_jsonl(test_fp))\n",
    "\n",
    "print(f\"There are {len(test_data)} examples in the test set\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "92f74f62-8427-438e-b1df-15fe53ce330c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([{'role': 'user',\n",
       "   'content': 'What are the zip codes for New York City in the United States?'}],\n",
       " [{'type': 'function',\n",
       "   'function': {'name': 'zipcode_by_city',\n",
       "    'description': 'Retrieves the zip code(s) of a given city using the GeoSource API.',\n",
       "    'parameters': {'type': 'object',\n",
       "     'properties': {'city': {'description': 'The name of the city for which to retrieve zip code(s). The city must be in the supported countries list (44 countries).',\n",
       "       'type': 'string',\n",
       "       'default': 'Brussels'}}}}},\n",
       "  {'type': 'function',\n",
       "   'function': {'name': 'place_details_google',\n",
       "    'description': 'Fetches contact and opening hours information for a place using its Google Place ID.',\n",
       "    'parameters': {'type': 'object',\n",
       "     'properties': {'is_id': {'description': 'The Google Place ID of the location to retrieve details for.',\n",
       "       'type': 'string',\n",
       "       'default': 'ChIJCxzX9tpwhlQRIFVwULOI2FM'}}}}},\n",
       "  {'type': 'function',\n",
       "   'function': {'name': 'get_states',\n",
       "    'description': 'Fetches a list of states based on various search parameters.',\n",
       "    'parameters': {'type': 'object',\n",
       "     'properties': {'keyword': {'description': \"A keyword to search for states. Defaults to 'asia'.\",\n",
       "       'type': 'string',\n",
       "       'default': 'asia'},\n",
       "      'perpage': {'description': 'The number of results to return per page. Defaults to 50.',\n",
       "       'type': 'integer',\n",
       "       'default': '50'},\n",
       "      'city_id': {'description': 'A city ID to find associated states. Defaults to None.',\n",
       "       'type': 'string'},\n",
       "      'is_id': {'description': 'A specific state ID to search for. Defaults to None.',\n",
       "       'type': 'string'},\n",
       "      'page': {'description': 'The number of results to return per page. Defaults to 50.',\n",
       "       'type': 'integer',\n",
       "       'default': '1'}}}}}])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Randomly choose\n",
    "test_sample = random.choice(test_data)\n",
    "\n",
    "# Visualize the inputs to the LLM - user query and available tools\n",
    "test_sample['messages'], test_sample['tools']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8577c4b-e0f6-4596-9bbc-33c47ca424aa",
   "metadata": {},
   "source": [
    "### 3.2 Send an Inference Call to NIM\n",
    "\n",
    "NIM exposes an OpenAI-compatible completions API endpoint, which you can query using the `OpenAI` client library as shown in the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "af83f794-f8cc-4cc8-8d42-a031dfb7f9aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ChatCompletionMessageToolCall(id='chatcmpl-tool-538bbb6949e642a5be7c6240c69f7caf', function=Function(arguments='{\"city\": \"New York City\"}', name='zipcode_by_city'), type='function')]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inference_client = OpenAI(\n",
    "  base_url = f\"{NIM_URL}/v1\",\n",
    "  api_key = \"None\"\n",
    ")\n",
    "\n",
    "completion = inference_client.chat.completions.create(\n",
    "  model = CUSTOMIZED_MODEL,\n",
    "  messages = test_sample[\"messages\"],\n",
    "  tools = test_sample[\"tools\"],\n",
    "  tool_choice = 'auto',\n",
    "  temperature = 0.1,\n",
    "  top_p = 0.7,\n",
    "  max_tokens = 512,\n",
    "  stream = False\n",
    ")\n",
    "\n",
    "completion.choices[0].message.tool_calls"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "459bb488-3bd9-4092-9c3a-f77f904606e4",
   "metadata": {},
   "source": [
    "Given that the fine-tuning job was successful, you can get an inference result comparable to the ground truth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "dcbdaedf-6fb4-41d6-8297-b7480c907892",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'type': 'function',\n",
       "  'function': {'name': 'zipcode_by_city',\n",
       "   'arguments': {'city': 'New York City'}}}]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The ground truth answer\n",
    "test_sample['tool_calls']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4f0fce4-3169-4ac4-8785-8a8f15959594",
   "metadata": {},
   "source": [
    "### 3.3 Take Note of Your Custom Model Name\n",
    "\n",
    "Take note of your custom model name, as you will use it to run evaluations in the subsequent notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "36398168-1051-4e54-ac3e-7a3e406f393a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name of your custom model is: xlam-tutorial-ns/llama-3.2-1b-xlam-run1@cust-BTkGbfifLfEAjV2THu3tas\n"
     ]
    }
   ],
   "source": [
    "print(f\"Name of your custom model is: {CUSTOMIZED_MODEL}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}