tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a434f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "BRANCH='r1.17.0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "developmental-gibraltar",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42daf8bf",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this notebook we demonstrate how to use p-tuning and prompt tuning within NeMo-Megatron. Both methods are parameter efficient alternatives to fine-tuning pretrained language models. Our NeMo implementation makes it possible to use one pretrained GPT model on many downstream tasks without needing to tune the model’s full set of parameters. It also allows for adding new tasks to your model without overwriting or disrupting previous tasks for which the model has already been p-tuned/prompt-tuned. Because the original model parameters are frozen and never altered by either method, p-tuning/prompt-tuning also avoid cartographic forgetting issues often encountered when fine-tuning models.\n",
    "\n",
    "- Our prompt tuning implementation is based off Lester et. al’s EMNLP 2021 paper [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/abs/2104.08691)\n",
    "\n",
    "- Our p-tuning implementation is based off Liu et al's paper [GPT Understands, Too](https://arxiv.org/abs/2103.10385).\n",
    "\n",
    "- Command line usage examples and API documentation can be found in [our user docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/nemo_megatron/prompt_learning.html). \n",
    "\n",
    "<img src=\"images/prompt_learning_forward_pass.png\" alt=\"Prompt Learning Forward Pass\"/>\n",
    "\n",
    "Our continuous learning capability for combined p-tuning and prompt tuning with GPT style models is a NeMo specific extension of the author’s original work.\n",
    "\n",
    "# The Plan\n",
    "\n",
    "We are going to show you how to:\n",
    "    \n",
    "    1. P-Tune/Prompt Tune a model on multiple tasks at the same time\n",
    "    2. Add a new task to a model that has already been P-Tuned/Prompt Tuned previously\n",
    "    \n",
    "We will first p-tune a GPT model on sentiment analysis, and intent and slot classification tasks. Then we will show how to add the squad question answering task to the same model we already p-tuned once.\n",
    "\n",
    "\n",
    "# Technical Overview\n",
    "Instead of selecting discrete text prompts in a manual or automated fashion, prompt tuning and p-tuning utilize virtual prompt embeddings that can be optimized via gradient decent. The only difference between prompt tuning and p-tuning within NeMo-Megatron is the architecture used to tune the soft prompt tokens during training.\n",
    "\n",
    "### Terminology\n",
    "We will be using the terms `continuous`, `soft`, and `virtual` token interchangeably to refer to embeddings inserted into the model prompt that have no concrete mapping to strings or characters within the model’s vocabulary. These virtual token embeddings exist in contrast to the `discrete`, `hard`, or `real` tokens that do make up the model’s vocabulary. Virtual tokens are purely 1D vectors with dimensionality equal to that of each real token embedding, matching the `hidden_size` hyperparameter. In training and inference, continuous token embeddings are inserted among discrete token embeddings according to a template you provide in the model’s config. We will demonstrate how to do this below.\n",
    "\n",
    "When referring to p-tuning and prompt tuning together, we will be using the phrase prompt learning for simplicity.\n",
    "\n",
    "### Prompt-Tuning\n",
    "In prompt-tuning a pretrained GPT model, soft prompt embeddings are initialized as a 2D matrix of size `total_virtual_tokens X hidden_size`. Each task the model is prompt-tuned to perform has its own 2D embedding matrix associated with it. Tasks do not share any parameters during training or inference. All GPT model parameters are frozen and only the embedding parameters for each task are updated during training.\n",
    "\n",
    "In prompt tuning you can specify how the embeddings are initialized for each task. You can either\n",
    "\n",
    "1. Initialize embedding parameters according to some random distribution\n",
    "2. Initialize embedding parameters from existing vocabulary embeddings (recommended)\n",
    "\n",
    "If you choose to initialize virtual token embeddings from existing embedding weights, you can provide the string of words you want to use for initialization in the model’s config. This string will be tokenized and tiled or truncated to match the specified number of virtual tokens you would like to use (`total_virtual_tokens`). Vocab embeddings are copied and used to initialize the soft prompt embedding matrix for each task. The vocab embeddings themselves are not updated or changed during prompt tuning.\n",
    "\n",
    "\n",
    "### P-Tuning\n",
    "In p-tuning, an LSTM model is used to predict virtual token embeddings. We refer to this LSTM model as our `prompt_encoder`. LSTM parameters are randomly initialized at the start of p-tuning. All GPT model parameters are frozen, and only the LSTM weights are updated at each training step. LSTM parameters are shared between all tasks that are p-tuned at the same time, but the LSTM model outputs unique virtual token embeddings for each task. The virtual tokens predicted by the LSTM are inserted among the discrete token input in the exact same manner as with prompt-tuning. You still specify the number of virtual tokens you want to use by setting `total_virtual_tokens` and each virtual token embedding is still a 1D vector of size `hidden_size`.\n",
    "\n",
    "\n",
    "\n",
    "# The Best of Both\n",
    "A single pretrained GPT model can use both p-tuning and prompt-tuning. While you must decide to use either p-tuning or prompt-tuning for each task you want your model to perform, you can p-tune your model on a set of tasks A, then prompt tune your same model on a different set of tasks B, then finally run inference on tasks from both A and B at the same time. During prompt-tuning or p-tuning, tasks tuned at the same time must use the same number of virtual tokens. During inference, tasks using differing amounts of virtual tokens can be run at the same time.\n",
    "\n",
    "Please see our [docs for more comparisons between prompt and p-tuning](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/nemo_megatron/prompt_learning.html). \n",
    "\n",
    "With all that covered, let's get started!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c27562",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import wget \n",
    "import pathlib"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0bfc7709",
   "metadata": {},
   "source": [
    "# Tasks and Datasets\n",
    "We will be using p-tuning to teach our GPT model to do **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)\"."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e0b0072a",
   "metadata": {},
   "source": [
    "# Data Preparation\n",
    "\n",
    "The prompt learning dataset loader accepts a list of json/dictionary objects or a list of json file names where each json file contains a collection of json objects. Each json object must include the field `taskname` which is a string identifier for the task the data example corresponds to. They should also include one or more fields corresponding to different sections of the discrete text prompt. The input data might look like:\n",
    "\n",
    "```\n",
    "[\n",
    "    {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT1], \"question\": [QUESTION_TEXT1], \"answer\": [ANSWER_TEXT1]},\n",
    "    {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT2], \"question\": [QUESTION_TEXT2], \"answer\": [ANSWER_TEXT2]},\n",
    "]\n",
    "```\n",
    "\n",
    "These additional fields can be unlimited in number and will be used to help map different parts of the discrete text input to a prompt template that you define. We will show how this mapping works and how to construct your prompt template in the `Prompt Formatting` section. "
   ]
  },
  {
   "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",
    "\n",
    "os.makedirs(DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "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.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py', NEMO_DIR)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "71813919",
   "metadata": {},
   "source": [
    "Now let's down load and process the dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "816791de",
   "metadata": {},
   "source": [
    "### SQuAD Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa16d8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "SQUAD_DIR = os.path.join(DATA_DIR, \"SQuAD\")\n",
    "os.makedirs(SQUAD_DIR, exist_ok=True)\n",
    "\n",
    "# Download the SQuAD dataset\n",
    "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\n",
    "!wget 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 $NEMO_DIR/prompt_learning_squad_preprocessing.py --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 -4 $SQUAD_DIR/squad_train.jsonl"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a385d319",
   "metadata": {},
   "source": [
    "We made a `.jsonl` file for each of the train, validation, and testing splits of the squad data. Every `.jsonl` file contains json objects with the fields `taskname`, `context`, `question`, and `answer`. The preprocessing script is called `prompt_learning_squad_preprocessing.py`. It should be in your `NEMO_DIR` and at `scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py` in the NeMo repo. \n",
    "\n",
    "The SQuAD dataset consists of various topics like `Beyoncé`, `IPod`, and `Symbiosis`. Each topic has several paragraphs associated with it, and each paragraph has several questions and answers related to it. When we separated the train/validation/test splits, we separated them on the topic level. For example, if the training set contains paragraphs and questions about the topic `Beyoncé`, neither the validation nor test sets will contain any questions on this topic. All questions about a certain topic are isolated to one split of the data. \n",
    "\n",
    "Like the Financial PhraseBank Dataset, we randomly selected 80% of the questions for training, 10% for validation, and 10% for test. This resulted in `69125` test examples, `8952` validation examples, and `8744` testing examples. The `answer` field was removed from test examples.\n",
    "\n",
    "Training on the full train split could take a lot of time, so we are going to clip the train split to 2k examples for the sake of this tutorial, and limit the validation dataset to 200 samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1473ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "! head -2000 $SQUAD_DIR/squad_train.jsonl > $SQUAD_DIR/squad_short_train.jsonl\n",
    "! head -200 $SQUAD_DIR/squad_val.jsonl > $SQUAD_DIR/squad_short_val.jsonl\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e19c8dc",
   "metadata": {},
   "source": [
    "# P-Tuning Model Config Setup\n",
    "\n",
    "Now we will begin setting up the config file used for prompt/p-tuning our GPT models! GPT Prompt learning within NeMo uses a class called `MegatronGPTPromptLearningModel` which has its own config file. We will start by loading an example prompt learning config file, then make changes to it to fit our tasks and training plans. "
   ]
  },
  {
   "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/conf/megatron_gpt_prompt_learning_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_prompt_learning_config.yaml\")\n",
    "config = OmegaConf.load(CONFIG_PATH)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ce966bcf",
   "metadata": {},
   "source": [
    "First let's set the datasets we've created in the config. We are going to start by p-tuning a GPT model on a small subset of the **Squad** task. We do this by setting the following config params below: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb1590f",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.data.train_ds = [f\"{SQUAD_DIR}/squad_short_train.jsonl\"]\n",
    "config.model.data.validation_ds = [f\"{SQUAD_DIR}/squad_short_val.jsonl\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e021b24",
   "metadata": {},
   "source": [
    "### Prompt Formatting\n",
    "Now that we have our dataset, lets define what we want the prompt to look like. \n",
    "\n",
    "The squad dataset json files contain fields named \"context\", \"question\" and \"answer\". The prompt formatting template allows us to arrange these fields and decide where to insert virtual prompts. We can add the `<|VIRTUAL_PROMPT_0|>` token anywere between the fields (although we recommend simply adding it in the leftmost position will be sufficient).\n",
    "\n",
    "For example, given a data jsonl file with examples like this: \n",
    "\n",
    "\n",
    "**{\"taskname\": \"squad\", \"context\": \"Super Bowl 50 was an American football ga... numerals 50.\", \"question\": \"What does AFC stand for?\", \"answer\": \"American Football Conference\"}**. \n",
    "\n",
    "\n",
    "We can create a prompt template set to `prompt_template = \"<|VIRTUAL_PROMPT_0|> Context: {context}\\n\\nquestion: {question}\\n\\nanswer: {answer}\"` other options are also possible, for example the `\\n` can be replaced with whitespace or the other of the context and question can be swapped. The answer however, should be at the end.\n",
    "\n",
    "Let's configure the prompt template for the task below:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f935b411",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.task_templates = [\n",
    "    \n",
    "    {\n",
    "        \"taskname\": \"squad\",\n",
    "        \"prompt_template\": \"<|VIRTUAL_PROMPT_0|> Context: {context}\\n\\nQuestion: {question}\\n\\nAnswer:{answer}\",\n",
    "        \"total_virtual_tokens\": 15,\n",
    "        \"virtual_token_splits\": [15],\n",
    "        \"truncate_field\": \"context\",\n",
    "        \"answer_only_loss\": True,\n",
    "        \"answer_field\": \"answer\",\n",
    "    },\n",
    "    \n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc438b5",
   "metadata": {},
   "source": [
    "Note each `task_template` item has 5 fields. \n",
    "\n",
    "- **`prompt_template`** is a string showing the model where to place virtual tokens and how to map dataset json fields to where they belong in the model prompt. \n",
    "\n",
    "\n",
    "- **`taskname`** refers to the same `taskname` in the dataset json objects. \n",
    "\n",
    "\n",
    "- **`total_virtual_tokens`** specifies the total number of virtual tokens that will be inserted into the model prompt.\n",
    "\n",
    "\n",
    "- **`virtual_token_splits`** specifies the number of virtual tokens that belong at each `<|VIRTUAL_PROMPT_#|>` marker. `virtual_token_splits` values should add up to `total_virtual_tokens`. The number of `virtual_token_splits` should match the number of `<|VIRTUAL_PROMPT_#|>` markers. \n",
    "\n",
    "\n",
    "- **`truncate_field`** specifies which field in the data json to truncate if the length of the input exceeds the maximum sequence length of the model. If `truncate_field` is set to `None`, examples that are too long are simply dropped from the dataset.\n",
    "\n",
    "\n",
    "- **`answer_only_loss`** Whether to limit loss calculation to only the answer portion of the prompt during tuning. `True` Strongly recommended for long prompts, but shorter prompts with single word answers seem to benefit from setting this to `False`. \n",
    "\n",
    "\n",
    "- **`answer_field`** The field in the data json corresponding to the answer. The loss will only be calculated on this portion of the prompt if `answer_only_loss` is `True`. The answer field must be at the end of the prompt template.\n",
    "\n",
    "In the `task_templates` we set above, `squad` has a different number of virtual tokens than `sentiment` and `intent_and_slot`. This is because we will be p-tuning on `squad` after we p-tune on the other two tasks and **we do not need to use the same number of virtual tokens between sessions**. We also set the `truncate` field for squad because the context can sometimes be longer than the model's max sequence length, and we want that field to be truncated if the example is too long. Lastly, we set `answer_only_loss` to true for `squad` due to the longer prompt. We've found `answer_only_loss=True` to work significantly better for this task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84579c7a",
   "metadata": {},
   "source": [
    "### Setting New Tasks\n",
    "After you p-tune your model this time, you can always go back and p-tune or prompt-tune your model on more tasks without over writing the virtual prompts who've trained this time. You can also use a different number of `total_virtual_tokens` between each training session as long as tasks ptuned or prompt tuned at the same time have the same number of `total_virtual_tokens`. For this reason, when you p-tune on a new task, you need to tell your model which of your tasks are new and which ones already exist (and thus you don't want to tune them). \n",
    "\n",
    "You do this by setting the `new_tasks` and `existing_tasks` values in the config file. Because we are p-tuning a model with no existing tasks, you should set `existing_tasks=[]` and `new_tasks=[\"sentiment\", \"intent_and_slot\"]` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57a73e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.existing_tasks = []\n",
    "config.model.new_tasks = [\"squad\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b77e88c",
   "metadata": {},
   "source": [
    "After p-tuning and/or prompt tuning is complete, you can run inference on all tasks at the same time, regardless of their `total_virtual_tokens` value."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d5017e",
   "metadata": {},
   "source": [
    "### Setting The Pre-Trained GPT Model\n",
    "We still need to set which GPT model we want to p-tune/prompt tune. Prompt learning methods work best with large GPT language models (5B or above), but the purposes of this tutorial, we are going to download a 345M parameter GPT model from NVIDIA NGC."
   ]
  },
  {
   "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",
    "MegatronGPTModel.list_available_models()"
   ]
  },
  {
   "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\"\n",
    "!wget  -nc --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/megatron_gpt_345m/versions/1/files/megatron_gpt_345m.nemo -O {NEMO_DIR}/{gpt_file_name}"
   ]
  },
  {
   "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.language_model_path = gpt_file_name"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "943a9c83",
   "metadata": {},
   "source": [
    "We can also set where we want the final prompt tuned model to be saved by setting `model.nemo_path`. By default the tuned prompt learning model will be saved in your current working directory to a `.nemo` file with the same name as your experiment (`config.name`). Let's change the save name to be `p_tuned_gpt.nemo`. **Your model path must end in `.nemo`.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a278cbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.exp_manager.checkpoint_callback_params.save_nemo_on_train_end= True\n",
    "config.exp_manager.checkpoint_callback_params.always_save_nemo= True\n",
    "config.exp_manager.checkpoint_callback_params.save_best_model= True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378a73e7",
   "metadata": {},
   "source": [
    "### Setting P-Tuning Specific Params\n",
    "Within the config file, p-tuning and prompt-tuning each have a couple of hyperparameters specific to them. We first need to tell the model that we want to do p-tuning, not prompt-tuning. To do this, we set the **`model.virtual_prompt_style`** hyperparameter like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68763763",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.modules.common import VirtualPromptStyle\n",
    "config.model.virtual_prompt_style = VirtualPromptStyle.P_TUNING"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "947dec63",
   "metadata": {},
   "source": [
    "Then we can set the 2 p-tuning specific parameters. Reminder, pp-tuning uses an LSTM prompt encoder to predict virtual tokens. \n",
    "\n",
    "- **`p_tuning.dropout`** the LSTM prompt encoder dropout probability \n",
    "- **`p_tuning.num_layers`** the number of LSTM layers you want your p-tuning prompt encoder to have\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f893ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.p_tuning.dropout = 0.0\n",
    "config.model.p_tuning.num_layers = 2\n",
    "config.model.global_batch_size = 2\n",
    "config.model.micro_batch_size = 1"
   ]
  },
  {
   "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))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b4bc7f3",
   "metadata": {},
   "source": [
    "### Setting Prompt-Tuning Specific Params\n",
    "\n",
    "Though we are not using prompt tuning in this training session, lets go over the prompt tuning specific parameters we would use if we were. \n",
    "\n",
    "- **`prompt_tuning.new_prompt_init_methods`** Whether you want to initialize virtual token embeddings from the embeddings of existing parts of the model's vocabulary (either 'text' or 'random')\n",
    "- **`prompt_tuning.new_prompt_init_text`** The text you want to use if you have 'text' in the list above, should be None otherwise. \n",
    "\n",
    "Each of the above hyperparameters are a list of strings. \n",
    "\n",
    "`new_prompt_init_methods` would look like `[\"text\", \"random\", \"text\", \"text\"]` if you were prompt tuning on 4 tasks at once, and you wanted the second task in `new_tasks` to use random initialization. \n",
    "\n",
    "`new_prompt_init_text` might look like `[\"some text I want to use\", None, \"some other text\", \"task text goes here\"]` for those four new tasks. \n",
    "\n",
    "The order of both should correspond to the order of the tasks you have listed in `model.new_tasks`. "
   ]
  },
  {
   "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": [
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy\n",
    "from pytorch_lightning.plugins.environments import TorchElasticEnvironment\n",
    "\n",
    "# lets modify some trainer configs\n",
    "# checks 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",
    "# use torch elastic cluster environment so `create_process_externally` is True\n",
    "# the launcher is set to None. It will not try to spawn new processes.\n",
    "# It won't create the misconfiguration error because of the `interactive session`\n",
    "os.environ[\"LOCAL_RANK\"] = '0'\n",
    "os.environ[\"RANK\"] = '0'\n",
    "os.environ[\"WORLD_SIZE\"] = '1'\n",
    "\n",
    "strategy = NLPDDPStrategy(find_unused_parameters=False, no_ddp_communication_hook=True)\n",
    "plugins = [TorchElasticEnvironment()]\n",
    "trainer = pl.Trainer(plugins= plugins, strategy=strategy, **config.trainer)\n",
    "\n",
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "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 = 'p_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)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5860bd90",
   "metadata": {},
   "source": [
    "We can also set learning hyperparameters as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4ec542",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set some of the learning parameters\n",
    "config.model.optim.lr = 1e-4\n",
    "config.model.precision = config.trainer.precision"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298b3dce",
   "metadata": {},
   "source": [
    "# First P-Tuning Session\n",
    "The only thing left to do is load up the model and begin p-tuning!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4bda19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model import MegatronGPTPromptLearningModel\n",
    "\n",
    "model = MegatronGPTPromptLearningModel(cfg=config.model, trainer=trainer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d99f433",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training set to 2 epochs by default in a cell above\n",
    "# Each epoch will take around 1min 15sec, but training time can vary\n",
    "trainer.fit(model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6aab09d4",
   "metadata": {},
   "source": [
    "# Inference After P-Tuning\n",
    "One way to run inference after p-tuning or prompt-tuning your model is to call `model.generate()`. `model.generate()` takes in \n",
    "\n",
    "- `inputs` which can be either a list of dictionary objects or `.jsonl` files containing dictionary objects, \n",
    "- `length_params`\n",
    "- `sampling_params`\n",
    "\n",
    "as arguments. More information about the [text generation API can be found here](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/transformer/text_generation.py).\n",
    "\n",
    "If `length_params` and `sampling_params` are set to `None`, the model generates output with a greedy decoding strategy and generates up to `30` new tokens. Most predictive downstream tasks (not text generation tasks), use greedy sampling. To see other ways to run inference with your prompt learning model and more details on how to define various inference parameters, visit `examples/nlp/language_modeling/megatron_gpt_eval.py`.\n",
    "\n",
    "Below are some randomly selected test examples from the sentiment classification and intent and slot classification test files. Notice that the `label` field is dropped from all test examples. The `MegatronPromptLearningDataset` called within `.generate()` automatically leaves fields in the prompt template empty when they are not provided in the data json. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc95e764",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_examples = [\n",
    "    {\"taskname\": \"squad\", \"context\": \"The build was released for download later in the day in standard 32-bit and 64-bit versions, plus a special 64-bit version which included SDKs and developer tools (Visual Studio Express and Expression Blend) for developing Metro-style apps. The Windows Store was announced during the presentation, but was not available in this build. According to Microsoft, there were about 535,000 downloads of the developer preview within the first 12 hours of its release. Originally set to expire on March 11, 2012, in February 2012 the Developer Preview's expiry date was changed to January 15, 2013.\", \"question\": \"When was the Developer preview initially intended to expire?\"},\n",
    "    {\"taskname\": \"squad\", \"context\": \"The structures of most federal governments incorporate mechanisms to protect the rights of component states. One method, known as 'intrastate federalism', is to directly represent the governments of component states in federal political institutions. Where a federation has a bicameral legislature the upper house is often used to represent the component states while the lower house represents the people of the nation as a whole. A federal upper house may be based on a special scheme of apportionment, as is the case in the senates of the United States and Australia, where each state is represented by an equal number of senators irrespective of the size of its population.\", \"question\": \"What is a bicameral legislature?\"},\n",
    "    {\"taskname\": \"squad\", \"context\": \"Imported mystery religions, which offered initiates salvation in the afterlife, were a matter of personal choice for an individual, practiced in addition to carrying on one's family rites and participating in public religion. The mysteries, however, involved exclusive oaths and secrecy, conditions that conservative Romans viewed with suspicion as characteristic of \\\"magic\\\", conspiratorial (coniuratio), or subversive activity. Sporadic and sometimes brutal attempts were made to suppress religionists who seemed to threaten traditional morality and unity, as with the senate's efforts to restrict the Bacchanals in 186 BC.\", \"question\": \"What was the practice of religion to the Romans?\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a5a358",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "response = model.generate(inputs=test_examples, length_params=None)\n",
    "\n",
    "print('The prediction results of some sample queries with the trained model:')\n",
    "for result in response['sentences']:\n",
    "    print(result)\n",
    "    print(\"-\" * 30)"
   ]
  }
 ],
 "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
}