{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "675dbb2d-57bb-414c-bf7e-63dc6aa072a5",
   "metadata": {},
   "source": [
    "# Evaluating a NeMo checkpoint with lm-eval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d4e997c-cf60-45f4-bbd1-c71a1c221687",
   "metadata": {},
   "source": [
    "This notebook showcases how to evaluate a model with NeMo 2.0. It will guide you through the process of in-framework deployment, and evaluation of completions and chat endpoints.\n",
    "\n",
    "In this tutorial we will evaluate an LLM on the [MMLU benchmark](https://arxiv.org/abs/2009.03300).\n",
    "The benchmark measures a language model's general knowledge across 57 diverse subjects, ranging from humanities and social sciences to STEM and professional fields, using multiple-choice questions.\n",
    "We will use two variants of the benchmarks: a more general one, that can be used to evaluate both base and instruction-tuned models, and a chat variant, that requires instruction-following capabilities from the model.\n",
    "\n",
    "> NOTE: It is recommended to run this notebook inside a [NeMo Framework container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo) which has all the required dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64387183-fff3-4b40-ae7f-2dd83a719e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import signal\n",
    "import subprocess\n",
    "\n",
    "from nemo.collections.llm import api\n",
    "from nemo.collections.llm.evaluation.api import EvaluationConfig, EvaluationTarget\n",
    "from nemo.collections.llm.evaluation.base import wait_for_fastapi_server\n",
    "from nemo.utils import logging\n",
    "\n",
    "logging.setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd220939-d4e6-45a8-930e-1ad1170ed1eb",
   "metadata": {},
   "source": [
    "## 1. Deploying the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d4a74ae-460d-4a0b-b400-5238ad7febcc",
   "metadata": {},
   "source": [
    "First, you need to prepare a NeMo 2 checkpoint of the model you would like to evaluate. For the purpose of this tutorial, we will use Llama 3.2 1B Instruct checkpoint, which you can download from the [NGC Catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/llama-3_2-1b-instruct). Make sure to mount the directory containing the checkpoint when starting the container. In this tutorial, we assume that the checkpoint is available under `\"/checkpoints/llama-3_2-1b-instruct_v2.0\"` path.\n",
    "\n",
    "> NOTE: Some steps in this tutorial are **only available for instruction-tuned (chat) models**. If you are working with a base model instead, you can still evaluate it using the `completions` endpoint and the standard `mmlu` task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c11c9a-a85e-4d85-a4a3-342bb0dafd1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# modify this variable to point to your checkpoint\n",
    "CHECKPOINT_PATH = \"/checkpoints/llama-3_2-1b-instruct_v2.0\"\n",
    "\n",
    "# if you are not using NeMo FW container, modify this path to point to scripts directory\n",
    "SCRIPTS_PATH = \"/opt/NeMo/scripts\"\n",
    "\n",
    "# modify this path if you would like to save results in a different directory\n",
    "WORKSPACE = \"/workspace\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf9a27f-da0f-4799-a17a-dcc9f5d3de7a",
   "metadata": {},
   "source": [
    "After downloading the model, we can deploy it for evaluation.\n",
    "The command below will start a server for the provided checkpoint in a separate process using the `deploy_in_fw_oai_server_eval.py` script.\n",
    "The script will deploy the model using the [Triton Inference Server](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/index.html) and set up an OpenAI-like endpoints for querying it.\n",
    "\n",
    "If you would like to use multiple devices for the deployment, you can specify deployment parameters to distribute your model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24d2425-ab20-427c-a8ca-8961d0a6e1b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "deploy_script = f\"{SCRIPTS_PATH}/deploy/nlp/deploy_in_fw_oai_server_eval.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08563a68-b872-46b3-a965-dae2c2833be3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python {deploy_script} --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e245c0c-0a12-40a9-ac62-e6e447e3c833",
   "metadata": {},
   "outputs": [],
   "source": [
    "deploy_process = subprocess.Popen(\n",
    "    ['python', deploy_script, '--nemo_checkpoint', CHECKPOINT_PATH],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79a0e537-39f7-41a0-bce0-65b7e1c2b8c6",
   "metadata": {},
   "source": [
    "The server exposes three endpoints:\n",
    "* `/v1/triton_health`\n",
    "* `/v1/completions/`\n",
    "* `/v1/chat/completions/`\n",
    "\n",
    "The `/v1/triton_health` allows you to check if the underlying Triton server is ready.\n",
    "The `/v1/completions/` endpoint allows you to send prompt to the model as-is, without applying the chat template. The model responds with a text completion.\n",
    "Finally, the `/v1/chat/completions/` endpoint allows for multi-turn conversational interactions with the model. This endpoint accepts a structured list of messages with different roles (system, user, assistant) to maintain context and generates chat-like responses. Under the hood, a chat template is applied to turn the conversation into a single input string.\n",
    "\n",
    "**Please note that the chat endpoint will not work correctly for base models, as they do not define a chat template.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb3c1c12-f440-445f-ba1a-bbb0fa33e6d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_url = \"http://0.0.0.0:8886\"\n",
    "model_name = \"triton_model\"\n",
    "\n",
    "completions_url = f\"{base_url}/v1/completions/\"\n",
    "chat_url = f\"{base_url}/v1/chat/completions/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd41d1c-aeb4-4bf7-bcd3-40ef4837a64d",
   "metadata": {},
   "source": [
    "Deployment can take a couple of minutes, especially for larger models. We will check the server status and wait until it is ready:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79e2ae8-758d-4bbc-b4a9-b65ea1140e27",
   "metadata": {},
   "outputs": [],
   "source": [
    "wait_for_fastapi_server(base_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6bc803-da74-4db6-bdda-1ffd6543e22c",
   "metadata": {},
   "source": [
    "After the model was deployed we can query it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd43d967-1a9c-4716-84a0-e162262cea9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "completions_payload = {\n",
    "    \"prompt\": \"My name is\",\n",
    "    \"model\": model_name,\n",
    "    \"max_tokens\": 16,\n",
    "}\n",
    "\n",
    "response = requests.post(completions_url, json=completions_payload)\n",
    "print(response.content.decode())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dfaa539-929b-4ad8-a03f-74e963a858dc",
   "metadata": {},
   "source": [
    "If you are working with a instruction-tuned model, you can also use the chat endpoint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab7251e4-6a26-439f-9b7b-8671692647a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_payload = {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"What is your name?\"}\n",
    "    ],\n",
    "    \"model\": model_name,\n",
    "    \"max_tokens\": 64,\n",
    "}\n",
    "\n",
    "response = requests.post(chat_url, json=chat_payload)\n",
    "print(response.content.decode())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2de593-23f2-4852-bd59-28932212ad64",
   "metadata": {},
   "source": [
    "## 2. Evaluating the completions endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "165bc6ad-aa4c-44a4-98ed-e24c11682d97",
   "metadata": {},
   "source": [
    "Now, we are ready to start the evaluation. First, we will evaluate the completions endpoint on the `mmlu` task.\n",
    "We will load a pre-defined configuration from [NVIDIA Evals Factory](https://pypi.org/project/nvidia-lm-eval/) lm-evaluation-harness.\n",
    "This configuration has a `--num_fewshot 5` flag specified, which means that each question to the model is prepended with five examples of question-answer pairs.\n",
    "This way, the model is guided on the correct way to format the output.\n",
    "\n",
    "For the purpose of this tutorial, we will only use one sample from each subset (by setting the `limit_samples` flag to 1).\n",
    "To run the full evaluation, remove this parameter from the command below.\n",
    "Alternatively, you can set the parameter to, for example, 0.1 to run the evaluation on 10% of the dataset.\n",
    "\n",
    "For more details on arguments in the EvaluationTarget and EvaluationConfig classes for evaluation, refer to [`nemo/collections/llm/evaluation/api.py`](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/llm/evaluation/api.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843b60c8-b852-412f-a358-5d5b04c6f130",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_config = EvaluationTarget(api_endpoint={\"url\": completions_url, \"type\": \"completions\"})\n",
    "eval_config = EvaluationConfig(\n",
    "    type=\"mmlu\",\n",
    "    params={\"limit_samples\": 1},\n",
    "    output_dir=f\"{WORKSPACE}/mmlu\",\n",
    ")\n",
    "\n",
    "completions_results = api.evaluate(target_cfg=target_config, eval_cfg=eval_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fae75bd0-1b3e-4c1a-b7bd-8fc2e4b5f39a",
   "metadata": {},
   "source": [
    "## 3. Evaluating the chat endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4af0504-3462-4ccc-90b0-a3d50dcca0a8",
   "metadata": {},
   "source": [
    "Now, we will use a \"chat\" variant of the same benchmark, `mmlu_instruct`, for evaluating the chat endpoint.\n",
    "In this evaluation scenario, we do not send examples of questions and answers (0-shot setting) but instead provide an instruction to the model on how the output should be formatted.\n",
    "\n",
    "This variant of the benchmark is more challenging as it requires the model to not only provide the correct answer but also to format it according to the instruction.\n",
    "\n",
    "Again, we will only use one sample from each subset.\n",
    "You can modify this behavior by changing or removing the `limit_samples` parameter from the command below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e593329-a26c-4802-8cd6-33c6f1d78a74",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_config = EvaluationTarget(api_endpoint={\"url\": chat_url, \"type\": \"chat\"})\n",
    "eval_config = EvaluationConfig(\n",
    "    type=\"mmlu_instruct\",\n",
    "    params={\"limit_samples\": 1},\n",
    "    output_dir=f\"{WORKSPACE}/mmlu_instruct\",\n",
    ")\n",
    "\n",
    "chat_results = api.evaluate(target_cfg=target_config, eval_cfg=eval_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09c9bf0f-3a96-4e38-be21-159edf4c6bfa",
   "metadata": {},
   "source": [
    "# 4. Inspecting the results and shuting the server down"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b56aa80-e0b7-4cba-9b93-df5d47898fd6",
   "metadata": {},
   "source": [
    "After the evaluation is finished, we can take a look at the results.\n",
    "We can compare the aggregated metrics or examine the scores for particular subtasks.\n",
    "\n",
    "It is often the case that results for the \"instruct\" variant are lower, as it requires strong instruction-following abilities from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2721e87-2573-4d58-91e1-dd563a80326c",
   "metadata": {},
   "outputs": [],
   "source": [
    "completions_results[\"groups\"][\"mmlu_str\"][\"metrics\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca8c371-e470-4f0b-9a5a-e756cc58ae7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_results[\"groups\"][\"mmlu_str\"][\"metrics\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24c065c-4578-4233-89b3-6ddb26521bf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "completions_results[\"tasks\"][\"mmlu_str_professional_medicine\"][\"metrics\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eecf6a8-4171-4642-81b3-6381c28d1fc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_results[\"tasks\"][\"mmlu_str_professional_medicine\"][\"metrics\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ceb940a-a07b-4af6-8f61-8ff173dc92c3",
   "metadata": {},
   "source": [
    "We can also examine the artifacts produced by both jobs.\n",
    "Inside the output directories, we can find `run_config.yml` files, which store details about the evaluation setup; `lm_cache_rank0.db`, file which contains a cache that can be used to resume an interrupted evaluation; and `triton_model` directory, which holds saved metrics as well as detailed logs for each input sample and its corresponding response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45279e09-1559-4529-96ca-23e87eaa2d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {WORKSPACE}/mmlu*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6cdf62d-5b6d-40c5-bcd8-2b42654f4855",
   "metadata": {},
   "source": [
    "Finally we can close the model's server.\n",
    "It can be done by sending `SIGINT` signal to the deployment process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4597fde8-f357-41fc-bf2b-4507d3bdca30",
   "metadata": {},
   "outputs": [],
   "source": [
    "deploy_process.send_signal(signal.SIGINT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3911d24e-5300-4b26-bc95-1a92ae88d6e9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}