{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Currently, this notebook must be run in a NeMo container.\n",
    "An example command to launch the container:\n",
    "```bash\n",
    "docker run --gpus all -it --rm -v <your_nemo_dir>:/NeMo --shm-size=8g -p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit stack=67108864 <your_nemo_container>\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Update megatron version to the newest.\n",
    "!cd /workspace && python -m pip install -e git+https://github.com/NVIDIA/Megatron-LM#egg=megatron-core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%cd /NeMo/tutorials/nlp\n",
    "BRANCH = 'main'\n",
    "import os\n",
    "import wget\n",
    "import sys\n",
    "sys.path.insert(0, \"../..\")  # find the local nemo first before the installed nemo"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "42daf8bf",
   "metadata": {},
   "source": [
    "### Introduction\n",
    "\n",
    "This notebook demonstrates how to apply PEFT in NeMo. For brevity, we have chosen LoRA as the PEFT technique and GPT as the language model, but the same recipe can be used for other PEFT techniques and language models, as described in the [Training](#training) section.\n",
    "\n",
    " The implementation of LoRA is based on the paper, [LoRA: Low-Rank Adaptation of Large Language Models](https://openreview.net/pdf?id=nZeVKeeFYf9) by Hu et al.\n",
    "\n",
    "This example demonstrates how to:\n",
    "\n",
    "    1. Train a LoRA model on a simple Extractive QA task.\n",
    "    2. Inspect the trained LoRA model showing the parameters it contains.\n",
    "    3. Run inference with the base model with the LoRA parameters."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0bfc7709",
   "metadata": {},
   "source": [
    "### Tasks and Datasets\n",
    "We will be using LoRA to teach our GPT model to do Extractive Question Answering.\n",
    "\n",
    "We will be using the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) reading comprehension dataset, consisting of questions posed by crowd workers on a set of Wikipedia articles, where the answer to every question is a segment of text. More information on [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) can be found on their website or in their paper by Rajpurkar et. al \"[Know What You Don’t Know: Unanswerable Questions for SQuAD](https://arxiv.org/pdf/1806.03822.pdf)\".\n",
    "\n",
    "LoRA (and all PEFT tuning) models expect at least two fields in the jsonl files. The `input` field should contain all the tokens necessary for the model to generate the `output`. For example for extractive QA, the `input` should contain the context text as well as the question.\n",
    "\n",
    "```\n",
    "[\n",
    "    {\"input\": \"User: Context: [CONTEXT_1] Question: [QUESTION_1]\\n\\nAssistant:\", \"output\": [ANSWER_1]},\n",
    "    {\"input\": \"User: Context: [CONTEXT_2] Question: [QUESTION_2]\\n\\nAssistant:\", \"output\": [ANSWER_2]},\n",
    "    {\"input\": \"User: Context: [CONTEXT_3] Question: [QUESTION_3]\\n\\nAssistant:\", \"output\": [ANSWER_3]},\n",
    "]\n",
    "```\n",
    "Note that we use keywords in the input like `Context:`, `Question:` to separate the text representing the context and question. We also use the keyword `User:` and end each of the input with `\\n\\nAssistant:` tokens. These are recommended because NeMo's instruction-tuned models are trained with a prefix of `User:` and suffix `\\n\\nAssistant:`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dbd41fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can replace DATA_DIR and NEMO_DIR with your own locations\n",
    "DATA_DIR = \"data\"\n",
    "NEMO_DIR = \".\"\n",
    "os.makedirs(DATA_DIR, exist_ok=True)\n",
    "SQUAD_DIR = os.path.join(DATA_DIR, \"SQuAD\")\n",
    "os.makedirs(SQUAD_DIR, exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "504a7b40",
   "metadata": {},
   "source": [
    "\n",
    "For each dataset we have preprocessing scripts pre-written in NeMo's example directory located in `examples/nlp`. Let's download those now. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e72a1dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the preprocessing scripts from github for the purpose of this tutorial\n",
    "! wget -nc https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "71813919",
   "metadata": {},
   "source": [
    "Now let's down load and process the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa16d8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the SQuAD dataset\n",
    "!wget -nc https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\n",
    "!wget -nc https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\n",
    "!mv train-v1.1.json {SQUAD_DIR}\n",
    "!mv dev-v1.1.json {SQUAD_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e3e25b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess squad data\n",
    "!python prompt_learning_squad_preprocessing.py --sft-format --data-dir {SQUAD_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b562d1de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What the squad dataset looks like after processing\n",
    "! head -200 $SQUAD_DIR/squad_train.jsonl > $SQUAD_DIR/squad_short_train.jsonl\n",
    "! head -20 $SQUAD_DIR/squad_val.jsonl > $SQUAD_DIR/squad_short_val.jsonl\n",
    "! head -4 $SQUAD_DIR/squad_short_val.jsonl\n",
    "! head -4 $SQUAD_DIR/squad_short_train.jsonl"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2e19c8dc",
   "metadata": {},
   "source": [
    "### Model Config Setup\n",
    "Now we will begin setting up the config file needed for PEFT tuning. We use a single config for all supported PEFT methods (LoRA, Adapter, IA3 and P-Tuning, as well as combinations of these).  All PEFT methods use the GPT finetuning class `MegatronGPTSFTModel` as the frozen base network, and use the `add_adapter()` method to add adapter weights for PEFT.\n",
    "\n",
    "Let's create a config object for LoRA training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5749c387",
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "\n",
    "CONFIG_DIR = os.path.join(NEMO_DIR, \"conf\")\n",
    "os.makedirs(CONFIG_DIR, exist_ok=True)\n",
    "\n",
    "# Download the example config file\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_config.yaml', CONFIG_DIR)\n",
    "\n",
    "# Load the example config file so we can start editing it\n",
    "CONFIG_PATH = os.path.join(CONFIG_DIR, \"megatron_gpt_finetuning_config.yaml\")\n",
    "config = OmegaConf.load(CONFIG_PATH)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ce966bcf",
   "metadata": {},
   "source": [
    "The `config` contains several attributes required by the `MegatronGPTSFTModel`. First we will set the training data path and the validation data path in the config.\n",
    "The `config` allows us to set a list of `jsonl` files as training files and sample examples from each file with different probabilities. For simplicity, we are going to use just one training file and thus the sampling probability is set to `1.0`\n",
    "\n",
    "We can also monitor validation loss from multiple validation files during training. Again for simplicity we will use just one validation file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb1590f",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.data.train_ds.file_names = [f\"{SQUAD_DIR}/squad_short_train.jsonl\"]\n",
    "config.model.data.train_ds.concat_sampling_probabilities=[1.0]\n",
    "config.model.data.validation_ds.file_names = [f\"{SQUAD_DIR}/squad_short_val.jsonl\"]\n",
    "config.model.data.validation_ds.names=[\"squad_val\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f6b7831a",
   "metadata": {},
   "source": [
    "### PEFT Config\n",
    "The attribute [config.model.peft](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_config.yaml#L78) contains settings that control the PEFT training method and its related hyperpameters. We currently support `lora`, `adapter`, `ptuning` and `ia3`. We can instruct the training script to use one of these methods by setting the config.model.peft.peft_scheme attribute.\n",
    "\n",
    "The other hyperparams associated with lora tuning are present in the [config.model.peft.lora_tuning](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_config.yaml#L92) attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72c9f966",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.peft.peft_scheme=\"lora\"  # we can also set this to adapter or ptuning or ia3\n",
    "print(OmegaConf.to_yaml(config.model.peft.lora_tuning))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c32e73c3",
   "metadata": {},
   "source": [
    "**Note:** In the original LoRA paper each attention projection (`K`, `Q`, `V` and `O`) can have their own Low-Rank projections. However, NeMo's attention implementation fuses `KQV` into a single projection and thus our LoRA implementation learns a single Low-Rank projection for `KQV` in a combined fashion. We do not support LoRA for the `O` matrix at this point."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e021b24",
   "metadata": {},
   "source": [
    "### Prompt Formatting\n",
    "The `config.model.data.train_ds.prompt_template` attribute allows us to further tweak the format of the input and output if needed. In this example, we have already incorporated our format inside the `jsonl` file during preprocessing, so we can keep the `prompt_template` in the config simple. (See previous section on Data Preparation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6aa5c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.data.train_ds.prompt_template =\"{input} {output}\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a0d5017e",
   "metadata": {},
   "source": [
    "### Setting the Pretrained GPT Model\n",
    "Next we will set the \"base language model\" upon which we will perform LoRA tuning. Obviously, larger base models will have better performance on downstream tasks but for the purposes of this tutorial we will use a small 345M parameter GPT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48cdf868",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check what GPT .nemo models we have available on NGC\n",
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel\n",
    "megatron_gpt_345m_nemo_url = MegatronGPTModel.list_available_models()[0].location\n",
    "megatron_gpt_345m_nemo_url # should point to the 345m megatron gpt model '.nemo' file"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ede350ed",
   "metadata": {},
   "source": [
    "If we wanted to use the GPT model class directly, we could instantiate a trainer then download the model by calling running \n",
    "`gpt_model = MegatronGPTModel.from_pretrained(model_name=\"megatron_gpt_345m\", trainer=trainer).cuda()`. But we just need the `.nemo` file in our working NeMo directory in this tutorial, so we will download it using `wget`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "364439a1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Download the model from NGC\n",
    "gpt_file_name = \"megatron_gpt_345m.nemo\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!wget  -nc --content-disposition {megatron_gpt_345m_nemo_url} -O {NEMO_DIR}/{gpt_file_name}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1d6a8a67",
   "metadata": {},
   "source": [
    "Now that we have a `.nemo` GPT file to work with. We need to add its path in our prompt learning config. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2778a5fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set GPT model path on prompt learning config\n",
    "config.model.restore_from_path = gpt_file_name"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "943a9c83",
   "metadata": {},
   "source": [
    "Next, we will set where we want to save all the intermediate training logs and checkpoints. As well as other training settings such as: number of training steps, batch size and validation check interval, and num_workers for data processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a278cbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.exp_manager.exp_dir=f\"{NEMO_DIR}/peft_lora\"\n",
    "config.exp_manager.explicit_log_dir=\"training_info\"\n",
    "config.trainer.max_steps=100\n",
    "config.model.micro_batch_size=1\n",
    "config.model.global_batch_size=4\n",
    "config.trainer.val_check_interval=50\n",
    "config.model.data.train_ds.num_workers=0  # 0 is recommended which just uses the main thread to process training examples\n",
    "config.model.data.validation_ds.num_workers=0 # 0 is recommended which just uses the main thread to process the validation examples"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a988d16e",
   "metadata": {},
   "source": [
    "Let's have a look at all the values we've set in the model config. You can change any of these values in the same manner we've been using above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a37ada",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Final model config\n",
    "print(OmegaConf.to_yaml(config.model))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4c048852",
   "metadata": {},
   "source": [
    "### Building the PyTorch Lightning Trainer\n",
    "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem.\n",
    "\n",
    "Let's first instantiate a Trainer object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90f85b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy\n",
    "import torch\n",
    "import lightning.pytorch as pl\n",
    "from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronTrainerBuilder\n",
    "\n",
    "# let's modify some trainer configs\n",
    "# check if we have GPU available and uses it\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "config.trainer.accelerator = accelerator\n",
    "config.trainer.devices = 1\n",
    "config.trainer.max_epochs = 4\n",
    "config.trainer.val_check_interval = 1.0\n",
    "\n",
    "# for PyTorch Native AMP set precision=16\n",
    "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
    "\n",
    "# setup cluster environment parameters\"\n",
    "os.environ[\"LOCAL_RANK\"] = '0'\n",
    "os.environ[\"RANK\"] = '0'\n",
    "os.environ[\"WORLD_SIZE\"] = '1'\n",
    "\n",
    "trainer = MegatronTrainerBuilder(config).create_trainer()\n",
    "\n",
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "890f0dc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(config.exp_manager))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4d0124c1",
   "metadata": {},
   "source": [
    "### Setting up a NeMo Experiment\n",
    "\n",
    "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c943ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "# Set name of the experiment \n",
    "config.name = 'lora_example_tuning'\n",
    "config.exp_manager.resume_if_exists = False\n",
    "\n",
    "# Init the experiment manager and view the exp_dir\n",
    "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "exp_dir = str(exp_dir)\n",
    "print(exp_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "298b3dce",
   "metadata": {},
   "source": [
    "### Training\n",
    "We now set up the process for training a LoRA model. We first require a config that contains details about the base language model upon which we will train our LoRA model. So we first extract the `model_cfg` from the checkpoint and update it with any new settings we employ in our current (LoRA) `config`. These are combined in the `merge_cfg_with` function.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edb38445",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel\n",
    "\n",
    "model_cfg = MegatronGPTSFTModel.merge_cfg_with(config.model.restore_from_path, config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dfc55a1c",
   "metadata": {},
   "source": [
    "Next, we instantiate the GPT model class and add the LoRA adapter\n",
    "When we call `add_adapter`, the model prints out the parameter count before and after the operation. We can clearly see the number of trainable parameters increase after adding the adapter.\n",
    "To print the parameter count manually, we can call `model.summarize()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a81d8741",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.parts.peft_config import LoraPEFTConfig\n",
    "\n",
    "model = MegatronGPTSFTModel.restore_from(config.model.restore_from_path, model_cfg, trainer=trainer)\n",
    "model.add_adapter(LoraPEFTConfig(model_cfg))\n",
    "# print(\"Parameter count manually:\\n\", model.summarize())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Simply substitute with the `MegatronT5SFTModel` class to use T5 instead of GPT.\n",
    "\n",
    "To use a different PEFT method, you can use a different config class in place of `LoraPEFTConfig`, such as `CanonicalAdaptersPEFTConfig`, `IA3PEFTConfig`, `PtuningPEFTConfig`. You can also use a combination of the methods by passing in a list:\n",
    "`model.add_adapter([LoraPEFTConfig(model_cfg), PtuningPEFTConfig(model_cfg)])`\n",
    "\n",
    "We're now ready to start training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d99f433",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer.fit(model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b8210d6d",
   "metadata": {},
   "source": [
    "Once training is completed you should see a saved '.nemo' file in this folder `{config.exp_manager.explicit_log_dir}/checkpoints`. This checkpoint will only contain the trained adapter weights, and not the frozen base model weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e19e65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The trained '.nemo' model is saved in the location below:\n",
    "! ls -lh {config.exp_manager.explicit_log_dir}/checkpoints\n",
    "print(config.exp_manager.explicit_log_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6aab09d4",
   "metadata": {},
   "source": [
    "### Inference\n",
    "The model object from `trainer.fit(model)` is also capable of doing inference. For the tutorial, however, we will re-load the saved `.nemo` lora model along with a `.nemo` base language model to simulate a more realistic scenario (where training does not happen right before inference).\n",
    "\n",
    "Run the cell below to reimport libraries and classes in case you did not run the training cells above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# reimport libraries and classes in case one wants to only run cells from the Inference section\n",
    "%cd /NeMo/tutorials/nlp\n",
    "import wget, os, sys\n",
    "sys.path.insert(0, \"../..\")  # find the local nemo first before the installed nemo\n",
    "from omegaconf import OmegaConf\n",
    "from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronTrainerBuilder\n",
    "from nemo.collections.nlp.parts.peft_config import LoraPEFTConfig\n",
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel\n",
    "\n",
    "NEMO_DIR = \".\"\n",
    "DATA_DIR = \"data\"\n",
    "CONFIG_DIR = os.path.join(NEMO_DIR, \"conf\")\n",
    "SQUAD_DIR = os.path.join(DATA_DIR, \"SQuAD\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "First, we will load and modify a config file that will be used for inference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Download the example config file\n",
    "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/language_modeling/tuning/conf/megatron_gpt_generate_config.yaml', CONFIG_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41ab98a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the example config file so we can start editing it\n",
    "CONFIG_EVAL_PATH = os.path.join(CONFIG_DIR, \"megatron_gpt_generate_config.yaml\")\n",
    "config_eval = OmegaConf.load(CONFIG_EVAL_PATH)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "36c58c18",
   "metadata": {},
   "source": [
    "We are going to modify the `config_eval` object that we created above. We will set the base language model as the `345m` model we downloaded earlier.\n",
    "\n",
    "Additionally, we will also set the `model.peft.restore_from_path` with the lora model we just trained. For the tutorial we will just use the validation data for inference as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a4e71a",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_eval.model.restore_from_path=\"megatron_gpt_345m.nemo\"\n",
    "config_eval.model.peft.restore_from_path=\"./training_info/checkpoints/lora_example_tuning.nemo\"\n",
    "config_eval.model.data.test_ds.file_names=[f\"{SQUAD_DIR}/squad_short_val.jsonl\"]\n",
    "config_eval.model.data.test_ds.names=[\"test_set\"]\n",
    "config_eval.model.data.test_ds.global_batch_size=1\n",
    "config_eval.model.data.test_ds.micro_batch_size=1\n",
    "config_eval.model.data.test_ds.tokens_to_generate=30\n",
    "config_eval.inference.greedy=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ace8f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_eval = MegatronTrainerBuilder(config_eval).create_trainer()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e745ac5e",
   "metadata": {},
   "source": [
    "The `config_eval` object is the hydra config at \"inference/test time\". This means it should contain information relevant for inference/test time, although some properties that were set at training time are still relevant. For example, whether training was done with `BOS` enabled or not, and other model specific attributes.\n",
    "\n",
    "So we extract the relevant information from the '.nemo' file of the lora model we just trained using the `merge_inference_cfg` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04a2201",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_model_cfg = MegatronGPTSFTModel.merge_inference_cfg(config_eval.model.peft.restore_from_path, config_eval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The cell below is required if you are running the notebook end-to-end, and if you use a different batch size for training and evaluation. In this case, the microbatch calculator needs to be rest. If you are running training only or inference only, feel free to ignore this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from megatron.core.num_microbatches_calculator import reconfigure_num_microbatches_calculator\n",
    "reconfigure_num_microbatches_calculator(\n",
    "    rank=0,\n",
    "    rampup_batch_size=None,\n",
    "    global_batch_size=config_eval.model.global_batch_size,\n",
    "    micro_batch_size=config_eval.model.micro_batch_size,\n",
    "    data_parallel_size=1,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "132ae378",
   "metadata": {},
   "source": [
    "Then, we load the base language model as well as the lora model we just trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b19cd0ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_eval = MegatronGPTSFTModel.restore_from(config_eval.model.restore_from_path, eval_model_cfg, trainer=trainer_eval)\n",
    "model_eval.load_adapters(config_eval.model.peft.restore_from_path, LoraPEFTConfig(eval_model_cfg))\n",
    "model_eval.freeze()\n",
    "\n",
    "print(\"Parameter count manually:\\n\", model_eval.summarize())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "012439d9",
   "metadata": {},
   "source": [
    "Next, we prepare the dataset and the dataloader objects that the model will perform inference on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12c390f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "_test_ds = model_eval._build_dataset(eval_model_cfg.data.test_ds, is_train=False)\n",
    "from torch.utils.data import DataLoader\n",
    "request_dl = DataLoader(\n",
    "    dataset=_test_ds[0],\n",
    "    batch_size=eval_model_cfg.data.test_ds.global_batch_size,\n",
    "    collate_fn=_test_ds[0].collate_fn,\n",
    ")\n",
    "config_inference = OmegaConf.to_container(config_eval.inference, resolve=True)\n",
    "model_eval.set_inference_config(config_inference)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "76592a1e",
   "metadata": {},
   "source": [
    "And finally, we call `trainer.predict` which triggers the inference process. The `response` object contains the outputs of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ba6a70c",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = trainer_eval.predict(model_eval, request_dl)\n",
    "for batch in response:\n",
    "    for s in batch['sentences']:\n",
    "        print(f\"{s}\\n\\n\")"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}