llama31_law.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19a20d1a",
   "metadata": {},
   "source": [
    "# Creating a Llama-3.1 LoRA adapter with the NeMo Framework and Deploy via NVIDIA NIM\n",
    "It's Llama 3.1 Day and we're excited to share our newest notebook in collaboration with the NVIDIA for finetuning using the NeMo framework and deploying it using an NVIDIA NIM. In this notebook, we'll be finetuning our own LoRA with a cleaned up version of the [Law StackExchange](https://huggingface.co/datasets/ymoslem/Law-StackExchange) dataset using NeMo Framework. Law StackExchange is a dataset of legal question/answers. Each record consists of a question, its title, as well as human-provided answers. Given a Law StackExchange forum question our goal is to auto-generate an appropriate title for it.\n",
    "\n",
    "####  NVIDIA NeMo Framework and NVIDIA NIM\n",
    "NVIDIA NeMo Framework is a scalable and cloud-native generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (e.g. Automatic Speech Recognition and Text-to-Speech). It enables users to efficiently create, customize, and deploy new generative AI models by leveraging existing code and pre-trained model checkpoints. After we finetune a LoRa using NeMo, we then deploy it using an NVIDIA NIM. An NVIDIA NIM is an accelerated inference solution for Generative AI models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "487f6804-97b2-4459-914d-e379c57407ad",
   "metadata": {},
   "source": [
    "#### Prerequistes\n",
    "\n",
    "Before you start this notebook, ensure that you have an NGC key available that is able to access the Llama3.1 NIM on NGC. To generate one, please visit build.nvidia.com and click Get API Key!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2527b3c1-d743-45b1-bfd9-d25fdf99f6b2",
   "metadata": {},
   "source": [
    "First we install the NGC CLI and docker and pull the `.nemo` checkpoint that we will use for finetuning. This can take about 5-7 minutes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8b80f5a-cd85-420d-a49c-7ad8148caf49",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "test -f setup-ngc.sh || (wget https://raw.githubusercontent.com/brevdev/notebooks/main/assets/setup-ngc.sh; chmod +x setup-ngc.sh)\n",
    "./setup-ngc.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea08e421-a42b-498c-9ddf-54f622be8e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "!COLUMNS=400 ./ngc-cli/ngc registry model download-version \"nvidia/nemo/llama-3_1-8b-instruct-nemo:1.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3057e525-7957-45c0-bedc-c347d4811081",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# this should the .nemo checkpoint that is saved\n",
    "!ls ./llama-3_1-8b-instruct-nemo_v1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a46c6f7-c3ce-46f9-a573-29b2d4797f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from rouge_score import rouge_scorer, scoring"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3323204-1463-4df3-8c75-5e95b6d66ba1",
   "metadata": {},
   "source": [
    "# Phase 1: Finetuning the LoRa adapter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb6a910-a05e-4ae1-aac4-56e5092be2b4",
   "metadata": {
    "tags": []
   },
   "source": [
    "##  Step-by-step PEFT finetuning instructions\n",
    "\n",
    "1. Prepare the dataset\n",
    "2. Run the PEFT finetuning script\n",
    "3. Inference with NeMo Framework\n",
    "4. Check the model accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ea5bd31",
   "metadata": {},
   "source": [
    "### Step 1: Prepare the dataset\n",
    "\n",
    "The dataset we used is a subset of the [Law-StackExchange dataset](https://huggingface.co/datasets/ymoslem/Law-StackExchange). We've already filtered and processed this dataset and it can be used to train the model for various different tasks - question title generation (summarization), law domain question answering, and question tag generation (multi-label classification). To run your own data cleaning and prepreocessing, please refer to the [data generation notebook](https://github.com/NVIDIA/NeMo-Curator/tree/main/tutorials/peft-curation-with-sdg). That tutorial also allows you to generate synthetic data and increase the size of the dataset.\n",
    "\n",
    "This dataset is licensed under the [CC BY-SA 4.0](https://creativecommons.org/licenses/by-sa/4.0/) license. You can use it for any purpose, including commercial use, without attribution. However, if you use the dataset in a publication, please cite the original authors and the [Law-StackExchange dataset](https://huggingface.co/datasets/ymoslem/Law-StackExchange) repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60de339b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://huggingface.co/datasets/bigmlguy2234/hf-law-qa-dataset/resolve/main/law-qa-curated.zip "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ab82b7-a4ff-45fc-9e0c-34157427a58f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip -j law-qa-curated.zip -d curated-data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "210e8ba7-f1a5-4264-bfcb-43f071eb6274",
   "metadata": {},
   "source": [
    "You should see the `law-qa-{train/val/test}.jsonl` splits in the curated folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4c81f14-9a19-4479-ae45-fded66027b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = os.path.join(\"./curated-data\")\n",
    "\n",
    "TRAIN_DS = os.path.join(DATA_DIR, \"law-qa-train.jsonl\")\n",
    "VAL_DS = os.path.join(DATA_DIR, \"law-qa-val.jsonl\")\n",
    "TEST_DS = os.path.join(DATA_DIR, \"law-qa-test.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42717840-33ca-44fe-b9e6-a43a22318f30",
   "metadata": {},
   "source": [
    "You will see several fields in the `.jsonl`, including `title`, `question`, `answer`, and other associated metadata.\n",
    "\n",
    "For this tutorial, our input will be the `answer` field, and output will be it's `title`. \n",
    "\n",
    "The following cell does two things -\n",
    "* Adds a template - a prompt instruction (which is optional), and format `{PROMPT} \\nQUESTION: {data[\"question\"]} \\nTITLE: `.\n",
    "* Saves the data splits into the same location, also appending a `_preprocessed` marker to them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ba498b-4a63-4a95-83fe-c2048b0d46db",
   "metadata": {},
   "outputs": [],
   "source": [
    " # Add a prompt instruction.\n",
    "PROMPT='''Generate a concise, engaging title for the following legal question on an internet forum. The title should be legally relevant, capture key aspects of the issue, and entice readers to learn more.'''\n",
    "\n",
    "# Creates a preprocessed version of the data files\n",
    "for input_file in [TRAIN_DS, VAL_DS, TEST_DS]:\n",
    "    output_file = input_file.rsplit('.', 1)[0] + '_preprocessed.jsonl'\n",
    "    with open(input_file, 'r') as infile, open(output_file, 'w') as outfile:\n",
    "        for line in infile:\n",
    "            # Parse each line as JSON\n",
    "            data = json.loads(line)\n",
    "\n",
    "            # Create a new dictionary with only the desired fields, renamed and formatted\n",
    "            new_data = {\n",
    "                \"input\": f'''{PROMPT} \\nQUESTION: {data[\"question\"]} \\nTITLE: ''',\n",
    "                \"output\": data['title']\n",
    "            }\n",
    "\n",
    "            # Write the new data as a JSON line to the output file\n",
    "            json.dump(new_data, outfile)\n",
    "            outfile.write('\\n')  # Add a newline after each JSON object\n",
    "\n",
    "    print(f\"Processed {input_file} and created {output_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f74e87bf-d41c-47d3-8e8f-c5e5d448b572",
   "metadata": {},
   "source": [
    "After running the above scripts, you will see  `law-qa-{train/test/val}_preprocessed.jsonl` files appear in the data directory.\n",
    "\n",
    "This is what an example will be formatted like -\n",
    "\n",
    "```json\n",
    "{\"input\": \"Generate a concise, engaging title for the following legal question on an internet forum. The title should be legally relevant, capture key aspects of the issue, and entice readers to learn more. \\nQUESTION: In order to be sued in a particular jurisdiction, say New York, a company must have a minimal business presence in the jurisdiction. What constitutes such a presence? Suppose the company engaged a New York-based Plaintiff, and its representatives signed the contract with the Plaintiff in New York City. Does this satisfy the minimum presence rule? Suppose, instead, the plaintiff and contract signing were in New Jersey, but the company hired a law firm with offices in New York City. Does this qualify? \\nTITLE: \", \n",
    " \"output\": \"What constitutes \\\"doing business in a jurisdiction?\\\"\"}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb1d887",
   "metadata": {},
   "source": [
    "### Step 2: Run PEFT finetuning script for LoRA\n",
    "\n",
    "NeMo framework includes a high level python script for fine-tuning  [megatron_gpt_finetuning.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py) that can abstract away some of the lower level API calls. Once you have your model downloaded and the dataset ready, LoRA fine-tuning with NeMo is essentially just running this script!\n",
    "\n",
    "For this demonstration, this training run is capped by `max_steps`, and validation is carried out every `val_check_interval` steps. If the validation loss does not improve after a few checks, training is halted to avoid overfitting.\n",
    "\n",
    "> `NOTE:` In the block of code below, pass the paths to your train, test and validation data files as well as path to the .nemo model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd4f9b0-49ef-4cf5-9482-5ee7a6bb3db0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# Set paths to the model, train, validation and test sets.\n",
    "MODEL=\"./llama-3_1-8b-instruct-nemo_v1.0/llama3_1_8b_instruct.nemo\"\n",
    "\n",
    "TRAIN_DS=\"[./curated-data/law-qa-train_preprocessed.jsonl]\"\n",
    "VALID_DS=\"[./curated-data/law-qa-val_preprocessed.jsonl]\"\n",
    "TEST_DS=\"[./curated-data/law-qa-test_preprocessed.jsonl]\"\n",
    "TEST_NAMES=\"[law]\"\n",
    "\n",
    "SCHEME=\"lora\"\n",
    "TP_SIZE=1\n",
    "PP_SIZE=1\n",
    "\n",
    "rm -rf results\n",
    "OUTPUT_DIR=\"./results/Meta-llama3.1-8B-Instruct-titlegen\"\n",
    "\n",
    "torchrun --nproc_per_node=1 \\\n",
    "/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py \\\n",
    "    exp_manager.exp_dir=${OUTPUT_DIR} \\\n",
    "    exp_manager.explicit_log_dir=${OUTPUT_DIR} \\\n",
    "    trainer.devices=1 \\\n",
    "    trainer.num_nodes=1 \\\n",
    "    trainer.precision=bf16-mixed \\\n",
    "    trainer.val_check_interval=0.2 \\\n",
    "    trainer.max_steps=50 \\\n",
    "    model.megatron_amp_O2=True \\\n",
    "    ++model.mcore_gpt=True \\\n",
    "    model.tensor_model_parallel_size=${TP_SIZE} \\\n",
    "    model.pipeline_model_parallel_size=${PP_SIZE} \\\n",
    "    model.micro_batch_size=1 \\\n",
    "    model.global_batch_size=32 \\\n",
    "    model.restore_from_path=${MODEL} \\\n",
    "    model.data.train_ds.file_names=${TRAIN_DS} \\\n",
    "    model.data.train_ds.concat_sampling_probabilities=[1.0] \\\n",
    "    model.data.validation_ds.file_names=${VALID_DS} \\\n",
    "    model.peft.peft_scheme=${SCHEME}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf4331fd-da30-4e29-8477-3085118e4a7b",
   "metadata": {},
   "source": [
    "This will create a LoRA adapter - a file named `megatron_gpt_peft_lora_tuning.nemo` in `./results/Meta-Llama-3-8B-Instruct/checkpoints/`. We'll use this later.\n",
    "\n",
    "To further configure the run above -\n",
    "\n",
    "* **A different PEFT technique**: The `peft.peft_scheme` parameter determines the technique being used. In this case, we did LoRA, but NeMo Framework supports other techniques as well - such as P-tuning, Adapters, and IA3. For more information, refer to the [PEFT support matrix](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/nlp/nemo_megatron/peft/landing_page.html). For example, for P-tuning, simply set \n",
    "\n",
    "```bash\n",
    "model.peft.peft_scheme=\"ptuning\" # instead of \"lora\"\n",
    "```\n",
    "\n",
    "* **Tuning Llama-3.1 70B**: You will need 8xA100 or 8xH100 GPUs. Provide the path to it's .nemo checkpoint (similar to the download and conversion steps earlier), and change the model parallelization settings for Llama-3.1 70B PEFT to distribute across the GPUs. It is also recommended to run the fine-tuning script from a terminal directly instead of Jupyter when using more than 1 GPU.\n",
    "\n",
    "```bash\n",
    "# Change the following settings, and run from a terminal directly\n",
    "trainer.devices=8\n",
    "model.tensor_model_parallel_size=8\n",
    "model.pipeline_model_parallel_size=1\n",
    "```\n",
    "\n",
    "You can override many such configurations while running the script. A full set of possible configurations is located in [NeMo Framework Github](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_config.yaml)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53979a4d",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Step 3: Inference with NeMo Framework\n",
    "\n",
    "Running text generation within the framework is also possible with running a Python script. Note that is more for testing and validation, not a full-fledged  deployment solution like NVIDIA NIM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d1e3f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    " # Check that the LORA model file exists\n",
    "!ls -l ./results/Meta-llama3.1-8B-Instruct-titlegen/checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3430a0b0-05a0-4179-8750-151d492bb9ae",
   "metadata": {},
   "source": [
    "In the code snippet below, the following configurations are worth noting - \n",
    "\n",
    "1. `model.restore_from_path` to the path for the Meta-Llama-3.1-8B-Instruct.nemo file.\n",
    "2. `model.peft.restore_from_path` to the path for the PEFT checkpoint that was created in the fine-tuning run in the last step.\n",
    "3. `model.test_ds.file_names` to the path of the preprocessed test file.\n",
    "\n",
    "If you have made any changes in model or experiment paths, please ensure they are configured correctly below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3508c556-5844-47d0-9f52-b25271ba77da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a smaller test subset for a quick eval demonstration.\n",
    "!head -n 128 ./curated-data/law-qa-test_preprocessed.jsonl > ./curated-data/law-qa-test_preprocessed-n128.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "568eb35d",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "MODEL=\"./llama-3_1-8b-instruct-nemo_v1.0/llama3_1_8b_instruct.nemo\"\n",
    "\n",
    "TEST_DS=\"[./curated-data/law-qa-test_preprocessed-n128.jsonl]\" # Smaller test split\n",
    "# TEST_DS=\"[./curated-data/law-qa-test_preprocessed.jsonl]\" # Full test set\n",
    "TEST_NAMES=\"[law]\"\n",
    "\n",
    "TP_SIZE=1\n",
    "PP_SIZE=1\n",
    "\n",
    "# This is where your LoRA checkpoint was saved\n",
    "PATH_TO_TRAINED_MODEL=\"./results/Meta-llama3.1-8B-Instruct-titlegen/checkpoints/megatron_gpt_peft_lora_tuning.nemo\"\n",
    "\n",
    "# The generation run will save the generated outputs over the test dataset in a file prefixed like so\n",
    "OUTPUT_PREFIX=\"law_titlegen_lora\"\n",
    "\n",
    "python /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_generate.py \\\n",
    "    model.restore_from_path=${MODEL} \\\n",
    "    model.peft.restore_from_path=${PATH_TO_TRAINED_MODEL} \\\n",
    "    trainer.devices=1 \\\n",
    "    trainer.num_nodes=1 \\\n",
    "    model.data.test_ds.file_names=${TEST_DS} \\\n",
    "    model.data.test_ds.names=${TEST_NAMES} \\\n",
    "    model.data.test_ds.global_batch_size=32 \\\n",
    "    model.data.test_ds.micro_batch_size=1 \\\n",
    "    model.data.test_ds.tokens_to_generate=50 \\\n",
    "    model.tensor_model_parallel_size=${TP_SIZE} \\\n",
    "    model.pipeline_model_parallel_size=${PP_SIZE} \\\n",
    "    inference.greedy=True  \\\n",
    "    model.data.test_ds.output_file_path_prefix=${OUTPUT_PREFIX} \\\n",
    "    model.data.test_ds.write_predictions_to_file=True \\\n",
    "    model.data.test_ds.add_bos=False \\\n",
    "    model.data.test_ds.add_eos=True \\\n",
    "    model.data.test_ds.add_sep=False \\\n",
    "    model.data.test_ds.label_key=\"output\" \\\n",
    "    model.data.test_ds.prompt_template=\"\\{input\\}\\ \\{output\\}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fe048f9",
   "metadata": {},
   "source": [
    "### Step 4: Check the model accuracy\n",
    "\n",
    "Now that the results are in, let's read the results and calculate the accuracy on the question title generation task.\n",
    "Let's take a look at one of the predictions in the generated output file. The pred key indicates what was generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5c0fdc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Take a look at predictions\n",
    "!head -n1  law_titlegen_lora_test_law_inputs_preds_labels.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1c91df7",
   "metadata": {},
   "source": [
    "For evaluating this task, we will use ROUGE.  It measures overlap of ngrams, and a higher score is better. While it's not perfect and it misses capturing the semantics of the prediction, it is a popular metric in academia and industry for evaluating such systems. The following method uses the rouge_score library to implement scoring. It will report `ROUGE_{1/2/L/Lsum}` metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "900f81c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_rouge(input_file: str) -> dict:\n",
    "    ROUGE_KEYS = [\"rouge1\", \"rouge2\", \"rougeL\", \"rougeLsum\"]\n",
    "    scorer = rouge_scorer.RougeScorer(ROUGE_KEYS, use_stemmer=True)\n",
    "    aggregator = scoring.BootstrapAggregator()\n",
    "    lines = [json.loads(line) for line in open(input_file)]\n",
    "    num_response_words = []\n",
    "    num_ref_words = []\n",
    "    for idx, line in enumerate(lines):\n",
    "        prompt = line['input']\n",
    "        response = line['pred']\n",
    "        answer = line['label']\n",
    "        scores = scorer.score(response, answer)\n",
    "        aggregator.add_scores(scores)\n",
    "        num_response_words.append(len(response.split()))\n",
    "        num_ref_words.append(len(answer.split()))\n",
    "\n",
    "    result = aggregator.aggregate()\n",
    "    rouge_scores = {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()}\n",
    "    print(rouge_scores)\n",
    "    print(f\"Average and stddev of response length: {np.mean(num_response_words):.2f}, {np.std(num_response_words):.2f}\")\n",
    "    print(f\"Average and stddev of ref length: {np.mean(num_ref_words):.2f}, {np.std(num_ref_words):.2f}\")\n",
    "\n",
    "    return rouge_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74e1bbce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compute_rouge(\"./law_titlegen_lora_test_law_inputs_preds_labels.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f08250a-8da8-4d91-b401-4b69b6d52dca",
   "metadata": {
    "tags": []
   },
   "source": [
    "For the Llama-3.1-8B-Instruct model, you should see accuracy comparable to the below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1d6c91d-1e0b-4f41-9e9e-02fdba6ac86f",
   "metadata": {},
   "source": [
    "`{'rouge1': 39.2082, 'rouge2': 18.8573, 'rougeL': 35.4098, 'rougeLsum': 35.3906}`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79712f62",
   "metadata": {},
   "source": [
    "# LoRA inference with NVIDIA NIM\n",
    "\n",
    "Now that we've trained our LoRA, lets go ahead and deploy them with NVIDIA NIM. NIM's let you deploy multiple LoRA adapters and supports the .nemo and Hugging Face model formats. We will deploy the Law LoRA adapter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9abda43f-0fe4-43c6-b9a5-1d7948b9a733",
   "metadata": {},
   "source": [
    "## Before you begin\n",
    "\n",
    "Lets download the NIM from NGC and get it up and running with the LoRa's that we've trained."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce2234e-05cc-4f1b-90e7-652c24663051",
   "metadata": {},
   "source": [
    "Note this cell might take a few minutes as it pulls the NIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03bd3247-fab5-46da-b8e9-493b6b34277b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "wget https://raw.githubusercontent.com/brevdev/notebooks/main/assets/setup-nim.sh -O setup-nim\n",
    "chmod +x setup-nim\n",
    "export NGC_API_KEY=\n",
    "./setup-nim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ba1e1da",
   "metadata": {},
   "source": [
    "This notebook includes instructions to send an inference call to NVIDIA NIM using the Python `requests` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3c106c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66c0e646",
   "metadata": {},
   "source": [
    "## Check available LoRA models\n",
    "\n",
    "Once the NIM server is up and running, check the available models as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a562b350",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'http://0.0.0.0:8000/v1/models'\n",
    "\n",
    "response = requests.get(url)\n",
    "data = response.json()\n",
    "\n",
    "print(json.dumps(data, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6a41604",
   "metadata": {},
   "source": [
    "This will return all the models available for inference by NIM. In this case, it will return the base model, as well as the LoRA adapters that were provided during NIM deployment - `llama3.1-8b-law-titlegen`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f855f078",
   "metadata": {},
   "source": [
    "---\n",
    "## Inference\n",
    "\n",
    "Inference can be performed by sending POST requests to the `/completions` endpoint.\n",
    "\n",
    "A few things to note:\n",
    "* The `model` parameter in the payload specifies the model that the request will be directed to. This can be the base model `meta/llama3.1-8b-instruct`, or any of the LoRA models, such as `llama3.1-8b-law-titlegen`.\n",
    "* `max_tokens` parameter specifies the maximum number of tokens to generate. At any point, the cumulative number of input prompt tokens and specified number of output tokens to generate should not exceed the model's maximum context limit. For llama3-8b-instruct, the context length supported is 8192 tokens.\n",
    "\n",
    "Following code snippets show how it's possible to send requests belonging to different LoRAs (or tasks). NIM dynamically loads the LoRA adapters and serves the requests. It also internally handles the batching of requests belonging to different LoRAs to allow better performance and more efficient of compute."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a5c4a8",
   "metadata": {},
   "source": [
    "### Title Generation\n",
    "\n",
    "Try sending an example from the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e2cebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'http://0.0.0.0:8000/v1/completions'\n",
    "headers = {\n",
    "    'accept': 'application/json',\n",
    "    'Content-Type': 'application/json'\n",
    "}\n",
    "\n",
    "# Example from the test set\n",
    "prompt=\"Generate a concise, engaging title for the following legal question on an internet forum. The title should be legally relevant, capture key aspects of the issue, and entice readers to learn more. \\nQUESTION: In order to be sued in a particular jurisdiction, say New York, a company must have a minimal business presence in the jurisdiction. What constitutes such a presence? Suppose the company engaged a New York-based Plaintiff, and its representatives signed the contract with the Plaintiff in New York City. Does this satisfy the minimum presence rule? Suppose, instead, the plaintiff and contract signing were in New Jersey, but the company hired a law firm with offices in New York City. Does this qualify? \\nTITLE: \"\n",
    "data = {\n",
    "    \"model\": \"llama3.1-8b-law-titlegen\",\n",
    "    \"prompt\": prompt,\n",
    "    \"max_tokens\": 50\n",
    "}\n",
    "\n",
    "response = requests.post(url, headers=headers, json=data)\n",
    "response_data = response.json()\n",
    "\n",
    "print(json.dumps(response_data, indent=4))"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}