summarizing_dialogue.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summarizing dialogue with generative AI\n",
    "\n",
    "This tutorial shows how you can load a pre-trained Large Language Model (LLM) from Hugging Face; improve its completion results with prompt engineering (zero-, one-, and few-shot inference); refine the model with instruction finetuning; and, finally, use Reinforcement Learning with Human Feedback (RLHF) to make its responses less toxic. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- [ 1 - Install and load dependencies](#1)\n",
    "- [ 2 - Summarizing dialogue without in-context learning](#2)\n",
    "- [ 3 - Summarizing dialogue with in-context learning](#3)\n",
    "  - [ 3.1 - Include instruction prompt but no examples (\"zero-shot inference\")](#3.1)\n",
    "  - [ 3.2 - Include instruction prompt and a single example (\"one-shot inference\")](#3.2)\n",
    "  - [ 3.3 - Include instruction prompt and several examples (\"few-shot inference\")](#3.3)\n",
    "  - [ 3.4 - Tweaking the generative Configuration Parameters](#3.4)\n",
    "- [ 4 - Tweaking the generation configuration parameters](#4)\n",
    "- [ 5 - Fine-tuning and automatically evaluating the model](#6)\n",
    "  - [ 5.1 - Preprocessing Dialogsum dataset to add instructions to prompt](#5.1)\n",
    "  - [ 5.2 - Fine-tune the model updating all parameters](#5.2)\n",
    "  - [ 5.3 - Visually inspect model results](#5.3)\n",
    "  - [ 5.4 - Automatically quantify model results with ROUGE metrics](#5.4)\n",
    "- [ 6 - Save compute resources with Parameter Efficient Fine-Tuning (PEFT)](#6)\n",
    "  - [ 6.1 - Setting up LoRA](#6.1)\n",
    "  - [ 6.2 - Training our model with LoRA](#6.2)\n",
    "  - [ 6.3 - Visually inspect LoRA model results](#6.3)\n",
    "  - [ 6.4 - Assessing LoRA model with ROUGE metrics](#6.4)\n",
    "- [ 7 - Using Reinforcement Learning with Human Feedback (RLHF) for further refinement](#7)\n",
    "  - [ 7.1 - Load the reward model](#7.1)\n",
    "  - [ 7.2 - Calculating baseline toxicity stats](#7.2)\n",
    "  - [ 7.3 - Fine-tuning our model to reduce toxicity](#7.3)\n",
    "  - [ 7.4 - Quantifying the toxicity of our detoxified model](#7.4)\n",
    "  - [ 7.5 - Visually comparing pre-detoxified vs. detoxified models](#7.5)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='1'></a>\n",
    "## 1 - Install and load dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: pip in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (23.2.1)\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting git+https://github.com/lvwerra/trl.git@25fa1bd\n",
      "  Cloning https://github.com/lvwerra/trl.git (to revision 25fa1bd) to c:\\users\\david abugaber\\appdata\\local\\temp\\pip-req-build-yguy0jm9\n",
      "  Resolved https://github.com/lvwerra/trl.git to commit 25fa1bd\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n",
      "Requirement already satisfied: torch>=1.4.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from trl==0.4.2.dev0) (1.13.1)\n",
      "Requirement already satisfied: transformers>=4.18.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from trl==0.4.2.dev0) (4.27.2)\n",
      "Requirement already satisfied: numpy>=1.18.2 in c:\\programdata\\anaconda3\\lib\\site-packages (from trl==0.4.2.dev0) (1.21.5)\n",
      "Requirement already satisfied: accelerate in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from trl==0.4.2.dev0) (0.21.0)\n",
      "Requirement already satisfied: datasets in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from trl==0.4.2.dev0) (2.11.0)\n",
      "Requirement already satisfied: typing-extensions in c:\\programdata\\anaconda3\\lib\\site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (4.3.0)\n",
      "Requirement already satisfied: filelock in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (3.6.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (0.13.2)\n",
      "Requirement already satisfied: packaging>=20.0 in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (6.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (2022.7.9)\n",
      "Requirement already satisfied: requests in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (2.28.1)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (0.13.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in c:\\programdata\\anaconda3\\lib\\site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (4.64.1)\n",
      "Requirement already satisfied: psutil in c:\\programdata\\anaconda3\\lib\\site-packages (from accelerate->trl==0.4.2.dev0) (5.9.0)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (11.0.0)\n",
      "Requirement already satisfied: dill<0.3.7,>=0.3.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (0.3.6)\n",
      "Requirement already satisfied: pandas in c:\\programdata\\anaconda3\\lib\\site-packages (from datasets->trl==0.4.2.dev0) (1.4.4)\n",
      "Requirement already satisfied: xxhash in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (3.2.0)\n",
      "Requirement already satisfied: multiprocess in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (0.70.14)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in c:\\programdata\\anaconda3\\lib\\site-packages (from datasets->trl==0.4.2.dev0) (2022.7.1)\n",
      "Requirement already satisfied: aiohttp in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (3.8.4)\n",
      "Requirement already satisfied: responses<0.19 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from datasets->trl==0.4.2.dev0) (0.18.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in c:\\programdata\\anaconda3\\lib\\site-packages (from packaging>=20.0->transformers>=4.18.0->trl==0.4.2.dev0) (3.0.9)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in c:\\programdata\\anaconda3\\lib\\site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\programdata\\anaconda3\\lib\\site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\programdata\\anaconda3\\lib\\site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (1.26.11)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\programdata\\anaconda3\\lib\\site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (2022.9.14)\n",
      "Requirement already satisfied: colorama in c:\\programdata\\anaconda3\\lib\\site-packages (from tqdm>=4.27->transformers>=4.18.0->trl==0.4.2.dev0) (0.4.5)\n",
      "Requirement already satisfied: attrs>=17.3.0 in c:\\programdata\\anaconda3\\lib\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (21.4.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.8.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in c:\\users\\david abugaber\\appdata\\roaming\\python\\python39\\site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.3.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in c:\\programdata\\anaconda3\\lib\\site-packages (from pandas->datasets->trl==0.4.2.dev0) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in c:\\programdata\\anaconda3\\lib\\site-packages (from pandas->datasets->trl==0.4.2.dev0) (2022.1)\n",
      "Requirement already satisfied: six>=1.5 in c:\\programdata\\anaconda3\\lib\\site-packages (from python-dateutil>=2.8.1->pandas->datasets->trl==0.4.2.dev0) (1.16.0)\n",
      "Building wheels for collected packages: trl\n",
      "  Building wheel for trl (setup.py): started\n",
      "  Building wheel for trl (setup.py): finished with status 'done'\n",
      "  Created wheel for trl: filename=trl-0.4.2.dev0-py3-none-any.whl size=67534 sha256=aa85fb906bae541cb69824152ac71b29baa89b085a8b010920c49b2ed807f684\n",
      "  Stored in directory: C:\\Users\\David Abugaber\\AppData\\Local\\Temp\\pip-ephem-wheel-cache-xq5ct0j1\\wheels\\7e\\d9\\20\\22e5aa1b64dfc536ba8515021cd5d17743333eea116903f19d\n",
      "Successfully built trl\n",
      "Installing collected packages: trl\n",
      "Successfully installed trl-0.4.2.dev0\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  Running command git clone --filter=blob:none --quiet https://github.com/lvwerra/trl.git 'C:\\Users\\David Abugaber\\AppData\\Local\\Temp\\pip-req-build-yguy0jm9'\n",
      "  WARNING: Did not find branch or tag '25fa1bd', assuming revision or ref.\n",
      "  Running command git checkout -q 25fa1bd\n"
     ]
    }
   ],
   "source": [
    "%pip install --upgrade pip\n",
    "%pip install --disable-pip-version-check \\\n",
    "    torch==1.13.1 \\\n",
    "    torchdata==0.5.1 --quiet\n",
    "\n",
    "%pip install \\\n",
    "    transformers==4.27.2 \\\n",
    "    datasets==2.11.0 \\\n",
    "    evaluate==0.4.0 \\\n",
    "    rouge_score==0.1.2 \\\n",
    "    loralib==0.1.1 \\\n",
    "    peft==0.3.0 --quiet\n",
    "\n",
    "# Install reinforcement Learning library from github.\n",
    "%pip install git+https://github.com/lvwerra/trl.git@25fa1bd "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Load various modules: datasets, Large Language Model, tokenizer, configurator, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import pipeline,AutoModelForSequenceClassification,AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig, TrainingArguments, Trainer\n",
    "from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead\n",
    "from peft import PeftModel, PeftConfig, LoraConfig, TaskType\n",
    "from trl import create_reference_model\n",
    "from trl.core import LengthSampler\n",
    "import torch\n",
    "import time\n",
    "import evaluate\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='2'></a>\n",
    "## 2 - Summarizing dialogue without in-context learning\n",
    "\n",
    "Load in dialogue data from DialogSum dataset (https://huggingface.co/datasets/knkarthick/dialogsum). This comprises 10,000+ short dialogues with manually annotated summaries. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset csv (C:/Users/David Abugaber/.cache/huggingface/datasets/knkarthick___csv/knkarthick--dialogsum-c8fac5d84cd35861/0.0.0/6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f84730841cd4767a20f053f0abc0fc7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_name = \"knkarthick/dialogsum\"\n",
    "\n",
    "dataset = load_dataset(data_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Data is already split into \"train\" and \"test\" subsets. Each subset contains the dialogue itself plus a human-written summary. Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#Person1#: Hi, Mr. Smith. I'm Doctor Hawkins. Why are you here today?\n",
      "#Person2#: I found it would be a good idea to get a check-up.\n",
      "#Person1#: Yes, well, you haven't had one for 5 years. You should have one every year.\n",
      "#Person2#: I know. I figure as long as there is nothing wrong, why go see the doctor?\n",
      "#Person1#: Well, the best way to avoid serious illnesses is to find out about them early. So try to come at least once a year for your own good.\n",
      "#Person2#: Ok.\n",
      "#Person1#: Let me see here. Your eyes and ears look fine. Take a deep breath, please. Do you smoke, Mr. Smith?\n",
      "#Person2#: Yes.\n",
      "#Person1#: Smoking is the leading cause of lung cancer and heart disease, you know. You really should quit.\n",
      "#Person2#: I've tried hundreds of times, but I just can't seem to kick the habit.\n",
      "#Person1#: Well, we have classes and some medications that might help. I'll give you more information before you leave.\n",
      "#Person2#: Ok, thanks doctor.\n",
      "Mr. Smith's getting a check-up, and Doctor Hawkins advises him to have one every year. Hawkins'll give some information about their classes and medications to help Mr. Smith quit smoking.\n"
     ]
    }
   ],
   "source": [
    "print(dataset['train'][0]['dialogue'])\n",
    "\n",
    "print(dataset['train'][0]['summary'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use the FLAN-T5 model from Hugging Face (more models available at https://huggingface.co/docs/transformers/index). \n",
    "\n",
    "Should be loaded as \"Seq2Seq\" because this is an encoder/decoder model. For autoregressive (decoder-only) models, would instead use AutoModelForCausalLM . See: https://stackoverflow.com/questions/75549632/difference-between-automodelforseq2seqlm-and-automodelforcausallm\n",
    "\n",
    "Model is already trained hence use the `.from_pretrained()` method. \n",
    "\n",
    "Model is ~1Gb, so might take a minute to load."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "id": "iAYlS40Z3l-v",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = 'google/flan-t5-base'\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sPqQA3TT3l_I",
    "tags": []
   },
   "source": [
    "Begin by tokenizing dataset.\n",
    "\n",
    "Use tokenizer that corresponds to FLAN-T5. use_fast=True speeds up the process (see https://huggingface.co/docs/transformers/v4.28.1/en/model_doc/auto#transformers.AutoTokenizer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "id": "sPqQA3TT3l_I",
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Example of the tokenizer encoding/decoding with our example dialogue from above. \n",
    "\n",
    "Use return_tensors='pt' for PyTorch, otherwise return_tensors='tf' for TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens represented as numbers:\n",
      "tensor([ 1713,   345, 13515,   536,  4663,    10,  2018,     6,  1363,     5,\n",
      "         3931,     5,    27,    31,    51,  7582, 12833,    77,     7,     5,\n",
      "         1615,    33,    25,   270,   469,    58,  1713,   345, 13515,   357,\n",
      "         4663,    10,    27,   435,    34,   133,    36,     3,     9,   207,\n",
      "          800,    12,   129,     3,     9,   691,    18,   413,     5,  1713,\n",
      "          345, 13515,   536,  4663,    10,  2163,     6,   168,     6,    25,\n",
      "           43,    29,    31,    17,   141,    80,    21,   305,   203,     5,\n",
      "          148,   225,    43,    80,   334,   215,     5,  1713,   345, 13515,\n",
      "          357,  4663,    10,    27,   214,     5,    27,  2320,    38,   307,\n",
      "           38,   132,    19,  1327,  1786,     6,   572,   281,   217,     8,\n",
      "         2472,    58,  1713,   345, 13515,   536,  4663,    10,  1548,     6,\n",
      "            8,   200,   194,    12,  1792,  2261, 21154,    19,    12,   253,\n",
      "           91,    81,   135,   778,     5,   264,   653,    12,   369,    44,\n",
      "          709,   728,     3,     9,   215,    21,    39,   293,   207,     5,\n",
      "         1713,   345, 13515,   357,  4663,    10,  8872,     5,  1713,   345,\n",
      "        13515,   536,  4663,    10,  1563,   140,   217,   270,     5,   696,\n",
      "         2053,    11, 11581,   320,  1399,     5,  2321,     3,     9,  1659,\n",
      "         6522,     6,   754,     5,   531,    25,  7269,     6,  1363,     5,\n",
      "         3931,    58,  1713,   345, 13515,   357,  4663,    10,  2163,     5,\n",
      "         1713,   345, 13515,   536,  4663,    10, 14627,    53,    19,     8,\n",
      "         1374,  1137,    13,  5084,  1874,    11,   842,  1994,     6,    25,\n",
      "          214,     5,   148,   310,   225, 10399,     5,  1713,   345, 13515,\n",
      "          357,  4663,    10,    27,    31,   162,  1971,  3986,    13,   648,\n",
      "            6,    68,    27,   131,    54,    31,    17,  1727,    12,  4583,\n",
      "            8,  7386,     5,  1713,   345, 13515,   536,  4663,    10,  1548,\n",
      "            6,    62,    43,  2287,    11,   128, 11208,    24,   429,   199,\n",
      "            5,    27,    31,   195,   428,    25,    72,   251,   274,    25,\n",
      "         1175,     5,  1713,   345, 13515,   357,  4663,    10,  8872,     6,\n",
      "         2049,  2472,     5,     1])\n",
      "\n",
      "Tokens decoded back into words:\n",
      "#Person1#: Hi, Mr. Smith. I'm Doctor Hawkins. Why are you here today? #Person2#: I found it would be a good idea to get a check-up. #Person1#: Yes, well, you haven't had one for 5 years. You should have one every year. #Person2#: I know. I figure as long as there is nothing wrong, why go see the doctor? #Person1#: Well, the best way to avoid serious illnesses is to find out about them early. So try to come at least once a year for your own good. #Person2#: Ok. #Person1#: Let me see here. Your eyes and ears look fine. Take a deep breath, please. Do you smoke, Mr. Smith? #Person2#: Yes. #Person1#: Smoking is the leading cause of lung cancer and heart disease, you know. You really should quit. #Person2#: I've tried hundreds of times, but I just can't seem to kick the habit. #Person1#: Well, we have classes and some medications that might help. I'll give you more information before you leave. #Person2#: Ok, thanks doctor.\n"
     ]
    }
   ],
   "source": [
    "test_input = dataset['train'][0]['dialogue']\n",
    "\n",
    "test_encoded = tokenizer(test_input, return_tensors='pt')\n",
    "\n",
    "test_decoded = tokenizer.decode(test_encoded[\"input_ids\"][0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Tokens represented as numbers:\\n{test_encoded[\"input_ids\"][0]}\\n')\n",
    "print(f'Tokens decoded back into words:\\n{test_decoded}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might notice that the length of the text string in words is longer than the length of the list of numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of our list of tokens is 284\n",
      "Length of our input text string in words is 168\n"
     ]
    }
   ],
   "source": [
    "print(\"Length of our list of tokens is \" + str(len(test_encoded[\"input_ids\"][0])))\n",
    "print(\"Length of our input text string in words is \" + str(len(test_decoded.split())))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reason for this is that sometimes a word is split up into several tokens. For instance, \"I'll\" becomes \"I\"+\"'\"+\"ll\", \"Smoking\" is split into \"Smok\"+\"ing\", and (perhaps more unconventionally) \"Person\" is split into \"P\"+\"erson\", due to the capitalization. Here is how you can view the tokens as strings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Hi',\n",
       " ',',\n",
       " '▁Mr',\n",
       " '.',\n",
       " '▁Smith',\n",
       " '.',\n",
       " '▁I',\n",
       " \"'\",\n",
       " 'm',\n",
       " '▁Doctor',\n",
       " '▁Hawk',\n",
       " 'in',\n",
       " 's',\n",
       " '.',\n",
       " '▁Why',\n",
       " '▁are',\n",
       " '▁you',\n",
       " '▁here',\n",
       " '▁today',\n",
       " '?',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁I',\n",
       " '▁found',\n",
       " '▁it',\n",
       " '▁would',\n",
       " '▁be',\n",
       " '▁',\n",
       " 'a',\n",
       " '▁good',\n",
       " '▁idea',\n",
       " '▁to',\n",
       " '▁get',\n",
       " '▁',\n",
       " 'a',\n",
       " '▁check',\n",
       " '-',\n",
       " 'up',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Yes',\n",
       " ',',\n",
       " '▁well',\n",
       " ',',\n",
       " '▁you',\n",
       " '▁have',\n",
       " 'n',\n",
       " \"'\",\n",
       " 't',\n",
       " '▁had',\n",
       " '▁one',\n",
       " '▁for',\n",
       " '▁5',\n",
       " '▁years',\n",
       " '.',\n",
       " '▁You',\n",
       " '▁should',\n",
       " '▁have',\n",
       " '▁one',\n",
       " '▁every',\n",
       " '▁year',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁I',\n",
       " '▁know',\n",
       " '.',\n",
       " '▁I',\n",
       " '▁figure',\n",
       " '▁as',\n",
       " '▁long',\n",
       " '▁as',\n",
       " '▁there',\n",
       " '▁is',\n",
       " '▁nothing',\n",
       " '▁wrong',\n",
       " ',',\n",
       " '▁why',\n",
       " '▁go',\n",
       " '▁see',\n",
       " '▁the',\n",
       " '▁doctor',\n",
       " '?',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Well',\n",
       " ',',\n",
       " '▁the',\n",
       " '▁best',\n",
       " '▁way',\n",
       " '▁to',\n",
       " '▁avoid',\n",
       " '▁serious',\n",
       " '▁illnesses',\n",
       " '▁is',\n",
       " '▁to',\n",
       " '▁find',\n",
       " '▁out',\n",
       " '▁about',\n",
       " '▁them',\n",
       " '▁early',\n",
       " '.',\n",
       " '▁So',\n",
       " '▁try',\n",
       " '▁to',\n",
       " '▁come',\n",
       " '▁at',\n",
       " '▁least',\n",
       " '▁once',\n",
       " '▁',\n",
       " 'a',\n",
       " '▁year',\n",
       " '▁for',\n",
       " '▁your',\n",
       " '▁own',\n",
       " '▁good',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁Ok',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Let',\n",
       " '▁me',\n",
       " '▁see',\n",
       " '▁here',\n",
       " '.',\n",
       " '▁Your',\n",
       " '▁eyes',\n",
       " '▁and',\n",
       " '▁ears',\n",
       " '▁look',\n",
       " '▁fine',\n",
       " '.',\n",
       " '▁Take',\n",
       " '▁',\n",
       " 'a',\n",
       " '▁deep',\n",
       " '▁breath',\n",
       " ',',\n",
       " '▁please',\n",
       " '.',\n",
       " '▁Do',\n",
       " '▁you',\n",
       " '▁smoke',\n",
       " ',',\n",
       " '▁Mr',\n",
       " '.',\n",
       " '▁Smith',\n",
       " '?',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁Yes',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Smok',\n",
       " 'ing',\n",
       " '▁is',\n",
       " '▁the',\n",
       " '▁leading',\n",
       " '▁cause',\n",
       " '▁of',\n",
       " '▁lung',\n",
       " '▁cancer',\n",
       " '▁and',\n",
       " '▁heart',\n",
       " '▁disease',\n",
       " ',',\n",
       " '▁you',\n",
       " '▁know',\n",
       " '.',\n",
       " '▁You',\n",
       " '▁really',\n",
       " '▁should',\n",
       " '▁quit',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁I',\n",
       " \"'\",\n",
       " 've',\n",
       " '▁tried',\n",
       " '▁hundreds',\n",
       " '▁of',\n",
       " '▁times',\n",
       " ',',\n",
       " '▁but',\n",
       " '▁I',\n",
       " '▁just',\n",
       " '▁can',\n",
       " \"'\",\n",
       " 't',\n",
       " '▁seem',\n",
       " '▁to',\n",
       " '▁kick',\n",
       " '▁the',\n",
       " '▁habit',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '1',\n",
       " '#',\n",
       " ':',\n",
       " '▁Well',\n",
       " ',',\n",
       " '▁we',\n",
       " '▁have',\n",
       " '▁classes',\n",
       " '▁and',\n",
       " '▁some',\n",
       " '▁medications',\n",
       " '▁that',\n",
       " '▁might',\n",
       " '▁help',\n",
       " '.',\n",
       " '▁I',\n",
       " \"'\",\n",
       " 'll',\n",
       " '▁give',\n",
       " '▁you',\n",
       " '▁more',\n",
       " '▁information',\n",
       " '▁before',\n",
       " '▁you',\n",
       " '▁leave',\n",
       " '.',\n",
       " '▁#',\n",
       " 'P',\n",
       " 'erson',\n",
       " '2',\n",
       " '#',\n",
       " ':',\n",
       " '▁Ok',\n",
       " ',',\n",
       " '▁thanks',\n",
       " '▁doctor',\n",
       " '.']"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.tokenize(test_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example of using FLAN-T5 to summarize dialogue, using our sample text from above.\n",
    "\n",
    "Zero-shot inference: we're feeding in the raw prompt without including any examples/instructions within it.\n",
    "\n",
    "Here I'm using the first example dialogue (corresponds to index 0 because Python starts counting at zero). Am storing the index number as a separate object so that it can be easily changed if I want to look at other example inputs.\n",
    "\n",
    "Note that max_new_tokens=80 limits the output to 80 tokens (remember that a single word may comprise several tokens). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test prompt:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Human-generated summary:\n",
      "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore.\n",
      "\n",
      "Completion - no in-context learning:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n"
     ]
    }
   ],
   "source": [
    "index = 0\n",
    "\n",
    "dialogue = dataset['test'][index]['dialogue']\n",
    "human_summary = dataset['test'][index]['summary']\n",
    "\n",
    "input_tokens = tokenizer(dialogue, return_tensors='pt')\n",
    "\n",
    "output = tokenizer.decode(model.generate(input_tokens[\"input_ids\"], max_new_tokens=80,)[0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Test prompt:\\n{dialogue}\\n')\n",
    "print(f'Human-generated summary:\\n{human_summary}\\n')\n",
    "print(f'Completion - no in-context learning:\\n{output}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feeding in the raw dialogue produces crappy results. In-context learning with zero-/one-/few-shot inference might help."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='3'></a>\n",
    "## 3 - Summarizing dialogue with in-context learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='3.1'></a>\n",
    "### 3.1 - Include instruction prompt but no examples (\"zero-shot inference\")\n",
    "\n",
    "This time, add an explicit instruction at the end of the prompt to see if our dialogue summaries improve. Still no examples to learn from, though, hence \"zero-shot.\"\n",
    "\n",
    "Note that more example prompts can be found in the FLAN-T5 Github repository (https://github.com/google-research/FLAN/tree/main/flan/v2) in this script (https://github.com/google-research/FLAN/blob/main/flan/v2/templates.py)\n",
    "\n",
    "For instance, some of the instruction prompts they use are: \"Briefly summarize that dialogue\", \"Here is a dialogue ... Write a short summary!\", \"What was that dialogue about, in two sentences or less?, \"Here is a dialogue... What were they talking about?\", etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test prompt:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Human-generated summary:\n",
      "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore.\n",
      "\n",
      "Completion - zero-shot learning:\n",
      "The memo will go out to all employees by this afternoon.\n"
     ]
    }
   ],
   "source": [
    "index = 0\n",
    "\n",
    "dialogue = dataset['test'][index]['dialogue']\n",
    "human_summary = dataset['test'][index]['summary']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "\n",
    "input_tokens = tokenizer(prompt, return_tensors='pt')\n",
    "\n",
    "output = tokenizer.decode(model.generate(input_tokens[\"input_ids\"], max_new_tokens=80,)[0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Test prompt:\\n{dialogue}\\n')\n",
    "print(f'Human-generated summary:\\n{human_summary}\\n')\n",
    "print(f'Completion - zero-shot learning:\\n{output}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A little better, but omits a lot of detail as to the actual content of the memo (at least for the dialogue in index zero -- might perform better for some of the other exampels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='3.2'></a>\n",
    "### 3.2 - Include instruction prompt and a single example (\"one-shot inference\")\n",
    "\n",
    "Now I'll try including a single (\"one-shot\") example of a good dialogue / summary combination, before prompting the model to produce its own on a different dialogue.\n",
    "\n",
    "For this, I'll have two separate indices: one for the example summary, and one for the dialogue to be summarized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (752 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test prompt:\n",
      "\n",
      "Produce a summary of the following dialogue.\n",
      "\n",
      "#Person1#: You're finally here! What took so long?\n",
      "#Person2#: I got stuck in traffic again. There was a terrible traffic jam near the Carrefour intersection.\n",
      "#Person1#: It's always rather congested down there during rush hour. Maybe you should try to find a different route to get home.\n",
      "#Person2#: I don't think it can be avoided, to be honest.\n",
      "#Person1#: perhaps it would be better if you started taking public transport system to work.\n",
      "#Person2#: I think it's something that I'll have to consider. The public transport system is pretty good.\n",
      "#Person1#: It would be better for the environment, too.\n",
      "#Person2#: I know. I feel bad about how much my car is adding to the pollution problem in this city.\n",
      "#Person1#: Taking the subway would be a lot less stressful than driving as well.\n",
      "#Person2#: The only problem is that I'm going to really miss having the freedom that you have with a car.\n",
      "#Person1#: Well, when it's nicer outside, you can start biking to work. That will give you just as much freedom as your car usually provides.\n",
      "#Person2#: That's true. I could certainly use the exercise!\n",
      "#Person1#: So, are you going to quit driving to work then?\n",
      "#Person2#: Yes, it's not good for me or for the environment.\n",
      "\n",
      "Summary:\n",
      "#Person2# arrives late because of traffic jam. #Person1# persuades #Person2# to use public transportations to keep healthy and to protect the environment.\n",
      "\n",
      "Produce a summary of the following dialogue.\n",
      "\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Summary:\n",
      "\n",
      "\n",
      "Human-generated summary:\n",
      "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore.\n",
      "\n",
      "Completion - one-shot learning:\n",
      "The memo will go out to all employees by this afternoon.\n"
     ]
    }
   ],
   "source": [
    "test_index = 0\n",
    "test_dialogue = dataset['test'][test_index]['dialogue']\n",
    "test_human_summary = dataset['test'][test_index]['summary']\n",
    "\n",
    "example_index = 3\n",
    "example_dialogue = dataset['test'][example_index]['dialogue']\n",
    "example_human_summary = dataset['test'][example_index]['summary']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{example_dialogue}\n",
    "\n",
    "Summary:\n",
    "{example_human_summary}\n",
    "\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{test_dialogue}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "\n",
    "input_tokens = tokenizer(prompt, return_tensors='pt')\n",
    "\n",
    "output = tokenizer.decode(model.generate(input_tokens[\"input_ids\"], max_new_tokens=80,)[0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Test prompt:\\n{prompt}\\n')\n",
    "print(f'Human-generated summary:\\n{test_human_summary}\\n')\n",
    "print(f'Completion - one-shot learning:\\n{output}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Maybe a slight improvement (?). Instead of \"the memo WILL go out\" now we get \"the memo IS TO BE DISTRIBUTED.\" Note that bizarrely, using indices 1 or 2 for the example produced some bizarre behavior: those example dialogues were also about a conversation between Ms. Dawson and the boss, so FLAN-T5 just reproduced the human-generated summaries for THOSE dialogues. So, having an example that was too similar worked against us here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a name='3.3'></a>\n",
    "### 3.3 - Include instruction prompt and several examples (\"few-shot inference\")\n",
    "\n",
    "Instead of just a single example in the prompt, I'll try adding several. So, \"few-shot inference\" instead of \"one-shot inference.\"\n",
    "\n",
    "Rather than tediously writing out separate slots in the prompt for example 1, example 2, example 3, etc., I can build a function to automatically add the new examples. The code below prints out an example of such an aggregated prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: You're finally here! What took so long?\n",
      "#Person2#: I got stuck in traffic again. There was a terrible traffic jam near the Carrefour intersection.\n",
      "#Person1#: It's always rather congested down there during rush hour. Maybe you should try to find a different route to get home.\n",
      "#Person2#: I don't think it can be avoided, to be honest.\n",
      "#Person1#: perhaps it would be better if you started taking public transport system to work.\n",
      "#Person2#: I think it's something that I'll have to consider. The public transport system is pretty good.\n",
      "#Person1#: It would be better for the environment, too.\n",
      "#Person2#: I know. I feel bad about how much my car is adding to the pollution problem in this city.\n",
      "#Person1#: Taking the subway would be a lot less stressful than driving as well.\n",
      "#Person2#: The only problem is that I'm going to really miss having the freedom that you have with a car.\n",
      "#Person1#: Well, when it's nicer outside, you can start biking to work. That will give you just as much freedom as your car usually provides.\n",
      "#Person2#: That's true. I could certainly use the exercise!\n",
      "#Person1#: So, are you going to quit driving to work then?\n",
      "#Person2#: Yes, it's not good for me or for the environment.\n",
      "\n",
      "Summary:\n",
      "#Person2# complains to #Person1# about the traffic jam, #Person1# suggests quitting driving and taking public transportation instead.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# tells Kate that Masha and Hero get divorced. Kate is surprised because she thought they are perfect couple.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Summary:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def prompt_builder(example_indices, test_index):\n",
    "    prompt = ''\n",
    "    for i in example_indices:\n",
    "        dialogue = dataset['test'][i]['dialogue']\n",
    "        summary = dataset['test'][i]['summary']\n",
    "        \n",
    "        prompt += f\"\"\"\n",
    "Produce a summary of the following dialogue:\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary:\n",
    "{summary}\n",
    "\n",
    "\"\"\"\n",
    "    \n",
    "    dialogue = dataset['test'][test_index]['dialogue']\n",
    "    \n",
    "    prompt += f\"\"\"\n",
    "Produce a summary of the following dialogue:\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "        \n",
    "    return prompt\n",
    "\n",
    "example_indices = [5,6]\n",
    "test_index = 0\n",
    "\n",
    "print(prompt_builder(example_indices, test_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Following the same procedure as before, I can feed this few-shot inference prompt to FLAN-T5. I'll use 6 prompts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test prompt:\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: You're finally here! What took so long?\n",
      "#Person2#: I got stuck in traffic again. There was a terrible traffic jam near the Carrefour intersection.\n",
      "#Person1#: It's always rather congested down there during rush hour. Maybe you should try to find a different route to get home.\n",
      "#Person2#: I don't think it can be avoided, to be honest.\n",
      "#Person1#: perhaps it would be better if you started taking public transport system to work.\n",
      "#Person2#: I think it's something that I'll have to consider. The public transport system is pretty good.\n",
      "#Person1#: It would be better for the environment, too.\n",
      "#Person2#: I know. I feel bad about how much my car is adding to the pollution problem in this city.\n",
      "#Person1#: Taking the subway would be a lot less stressful than driving as well.\n",
      "#Person2#: The only problem is that I'm going to really miss having the freedom that you have with a car.\n",
      "#Person1#: Well, when it's nicer outside, you can start biking to work. That will give you just as much freedom as your car usually provides.\n",
      "#Person2#: That's true. I could certainly use the exercise!\n",
      "#Person1#: So, are you going to quit driving to work then?\n",
      "#Person2#: Yes, it's not good for me or for the environment.\n",
      "\n",
      "Summary:\n",
      "#Person2# complains to #Person1# about the traffic jam, #Person1# suggests quitting driving and taking public transportation instead.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# tells Kate that Masha and Hero get divorced. Kate is surprised because she thought they are perfect couple.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# tells Kate that Masha and Hero are getting a peaceful divorce. Kate feels surprised and asks about their kids.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# and Kate talk about the divorce between Masha and Hero. Kate feels surprised because she thought they are well matched\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Happy Birthday, this is for you, Brian.\n",
      "#Person2#: I'm so happy you remember, please come in and enjoy the party. Everyone's here, I'm sure you have a good time.\n",
      "#Person1#: Brian, may I have a pleasure to have a dance with you?\n",
      "#Person2#: Ok.\n",
      "#Person1#: This is really wonderful party.\n",
      "#Person2#: Yes, you are always popular with everyone. and you look very pretty today.\n",
      "#Person1#: Thanks, that's very kind of you to say. I hope my necklace goes with my dress, and they both make me look good I feel.\n",
      "#Person2#: You look great, you are absolutely glowing.\n",
      "#Person1#: Thanks, this is a fine party. We should have a drink together to celebrate your birthday\n",
      "\n",
      "Summary:\n",
      "#Person1# and Brian are at the birthday party of Brian. Brian thinks #Person1# looks great and is popular.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Happy Birthday, this is for you, Brian.\n",
      "#Person2#: I'm so happy you remember, please come in and enjoy the party. Everyone's here, I'm sure you have a good time.\n",
      "#Person1#: Brian, may I have a pleasure to have a dance with you?\n",
      "#Person2#: Ok.\n",
      "#Person1#: This is really wonderful party.\n",
      "#Person2#: Yes, you are always popular with everyone. and you look very pretty today.\n",
      "#Person1#: Thanks, that's very kind of you to say. I hope my necklace goes with my dress, and they both make me look good I feel.\n",
      "#Person2#: You look great, you are absolutely glowing.\n",
      "#Person1#: Thanks, this is a fine party. We should have a drink together to celebrate your birthday\n",
      "\n",
      "Summary:\n",
      "#Person1# attends Brian's birthday party. Brian thinks #Person1# looks great and charming.\n",
      "\n",
      "\n",
      "Produce a summary of the following dialogue:\n",
      "\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Summary:\n",
      "\n",
      "\n",
      "Human-generated summary:\n",
      "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore.\n",
      "\n",
      "Completion - one-shot learning:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n"
     ]
    }
   ],
   "source": [
    "example_indices = [5,6,7,8,9,10]\n",
    "test_index = 0\n",
    "\n",
    "test_human_summary = dataset['test'][test_index]['summary']\n",
    "\n",
    "prompt = prompt_builder(example_indices, test_index)\n",
    "\n",
    "input_tokens = tokenizer(prompt, return_tensors='pt')\n",
    "\n",
    "output = tokenizer.decode(model.generate(input_tokens[\"input_ids\"], max_new_tokens=80,)[0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Test prompt:\\n{prompt}\\n')\n",
    "print(f'Human-generated summary:\\n{test_human_summary}\\n')\n",
    "print(f'Completion - few-shot learning:\\n{output}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Oddly enough, adding more examples didn't seem to improve our summary much. In any case, adding at least ONE example was better than having no in-context learning (which just spat out \"#Person1#: Ms. Dawson, I need you to take a dictation for me.\").\n",
    "\n",
    "Two things to note:\n",
    "\n",
    "-best practices suggest that adding more than 6 examples (\"6-shot learning\") does not improve results\n",
    "\n",
    "-models have limited input context length. In the case of FLAN-T5, it's 512 tokens. So, after that ceiling, any subsequent enty gets ignored."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a name='4'></a>\n",
    "## 4 - Tweaking the generation configuration parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "There are other parameters that we can play with, other than the max_new_tokens limit to the output length. For more, see: https://huggingface.co/docs/transformers/v4.29.1/en/main_classes/text_generation#transformers.GenerationConfig. \n",
    "\n",
    "For instance, setting \"do_sample = True\" makes the model *not* just choose the highest-probability output token (what's called \"greedy decoding\") but instead sample from the top candidates. \n",
    "\n",
    "For instance, setting \"top_k = 20\" limits the output to choosing from the 20 highest-probability tokens. Similarly, \"top_p = .50\" limits the output to the set of highest-probability tokens whose probabilities cumulatively sum to .50.\n",
    "\n",
    "Finally, setting \"temperature\" above 1 makes the model use a broader (less peaked) probability function, with a higher chance of choosing low-probability candidates. \n",
    "\n",
    "Here is an example of such parameters at play."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test prompt:\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: You're finally here! What took so long?\n",
      "#Person2#: I got stuck in traffic again. There was a terrible traffic jam near the Carrefour intersection.\n",
      "#Person1#: It's always rather congested down there during rush hour. Maybe you should try to find a different route to get home.\n",
      "#Person2#: I don't think it can be avoided, to be honest.\n",
      "#Person1#: perhaps it would be better if you started taking public transport system to work.\n",
      "#Person2#: I think it's something that I'll have to consider. The public transport system is pretty good.\n",
      "#Person1#: It would be better for the environment, too.\n",
      "#Person2#: I know. I feel bad about how much my car is adding to the pollution problem in this city.\n",
      "#Person1#: Taking the subway would be a lot less stressful than driving as well.\n",
      "#Person2#: The only problem is that I'm going to really miss having the freedom that you have with a car.\n",
      "#Person1#: Well, when it's nicer outside, you can start biking to work. That will give you just as much freedom as your car usually provides.\n",
      "#Person2#: That's true. I could certainly use the exercise!\n",
      "#Person1#: So, are you going to quit driving to work then?\n",
      "#Person2#: Yes, it's not good for me or for the environment.\n",
      "\n",
      "Summary:\n",
      "#Person2# complains to #Person1# about the traffic jam, #Person1# suggests quitting driving and taking public transportation instead.\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# tells Kate that Masha and Hero get divorced. Kate is surprised because she thought they are perfect couple.\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# tells Kate that Masha and Hero are getting a peaceful divorce. Kate feels surprised and asks about their kids.\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Kate, you never believe what's happened.\n",
      "#Person2#: What do you mean?\n",
      "#Person1#: Masha and Hero are getting divorced.\n",
      "#Person2#: You are kidding. What happened?\n",
      "#Person1#: Well, I don't really know, but I heard that they are having a separation for 2 months, and filed for divorce.\n",
      "#Person2#: That's really surprising. I always thought they are well matched. What about the kids? Who get custody?\n",
      "#Person1#: Masha, it seems quiet and makable, no quarrelling about who get the house and stock and then contesting the divorce with other details worked out.\n",
      "#Person2#: That's the change from all the back stepping we usually hear about. Well, I still can't believe it, Masha and Hero, the perfect couple. When would they divorce be final?\n",
      "#Person1#: Early in the New Year I guess.\n",
      "\n",
      "Summary:\n",
      "#Person1# and Kate talk about the divorce between Masha and Hero. Kate feels surprised because she thought they are well matched\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Happy Birthday, this is for you, Brian.\n",
      "#Person2#: I'm so happy you remember, please come in and enjoy the party. Everyone's here, I'm sure you have a good time.\n",
      "#Person1#: Brian, may I have a pleasure to have a dance with you?\n",
      "#Person2#: Ok.\n",
      "#Person1#: This is really wonderful party.\n",
      "#Person2#: Yes, you are always popular with everyone. and you look very pretty today.\n",
      "#Person1#: Thanks, that's very kind of you to say. I hope my necklace goes with my dress, and they both make me look good I feel.\n",
      "#Person2#: You look great, you are absolutely glowing.\n",
      "#Person1#: Thanks, this is a fine party. We should have a drink together to celebrate your birthday\n",
      "\n",
      "Summary:\n",
      "#Person1# and Brian are at the birthday party of Brian. Brian thinks #Person1# looks great and is popular.\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Happy Birthday, this is for you, Brian.\n",
      "#Person2#: I'm so happy you remember, please come in and enjoy the party. Everyone's here, I'm sure you have a good time.\n",
      "#Person1#: Brian, may I have a pleasure to have a dance with you?\n",
      "#Person2#: Ok.\n",
      "#Person1#: This is really wonderful party.\n",
      "#Person2#: Yes, you are always popular with everyone. and you look very pretty today.\n",
      "#Person1#: Thanks, that's very kind of you to say. I hope my necklace goes with my dress, and they both make me look good I feel.\n",
      "#Person2#: You look great, you are absolutely glowing.\n",
      "#Person1#: Thanks, this is a fine party. We should have a drink together to celebrate your birthday\n",
      "\n",
      "Summary:\n",
      "#Person1# attends Brian's birthday party. Brian thinks #Person1# looks great and charming.\n",
      "\n",
      "\n",
      "Here is a dialogue:\n",
      "\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "\n",
      "Summary:\n",
      "\n",
      "\n",
      "Human-generated summary:\n",
      "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore.\n",
      "\n",
      "Completion - one-shot learning:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me. #Person2#: Yes, sir. Go ahead. #Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready? #Person2#: Yes, sir. Go\n"
     ]
    }
   ],
   "source": [
    "generation_config = GenerationConfig(max_new_tokens=50, do_sample=True, top_p = .8, temperature=1.5)\n",
    "\n",
    "output = tokenizer.decode(model.generate(input_tokens[\"input_ids\"], max_new_tokens=80,)[0],skip_special_tokens=True)\n",
    "\n",
    "print(f'Test prompt:\\n{prompt}\\n')\n",
    "print(f'Human-generated summary:\\n{test_human_summary}\\n')\n",
    "print(f'Completion - few-shot learning:\\n{output}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "#resetting parameters\n",
    "generation_config = GenerationConfig(max_new_tokens=50, do_sample=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='5'></a>\n",
    "## 5 - Fine-tuning and automatically evaluating the model\n",
    "\n",
    "Instead of relying on modified prompts to get the model to spit out the desired output (which doesn't actually change the model but puts the burden on the human providing the input), we can modify the parameter weights to change the model itself. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a first set, we can get the size of the model and determine how many of the parameters are \"trainable\" with a function like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable model parameters: 247577856\n",
      "all model parameters: 247577856\n",
      "percentage of trainable model parameters: 100.00%\n"
     ]
    }
   ],
   "source": [
    "def print_number_of_trainable_model_parameters(model):\n",
    "    trainable_model_params = 0\n",
    "    all_model_params = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_model_params += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_model_params += param.numel()\n",
    "    return f\"trainable model parameters: {trainable_model_params}\\nall model parameters: {all_model_params}\\npercentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%\"\n",
    "\n",
    "print(print_number_of_trainable_model_parameters(model))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='5.1'></a>\n",
    "### 5.1 - Preprocessing Dialogsum dataset to add instructions to prompt\n",
    "\n",
    "Not unlike what we did above with the in-context learning, we need to reformat the dialog/summary pairs in the Dialogsum dataset by adding instructions: a message before each dialogue to orient the model, an explicit instruction to create a summary, and a slot for the humman-generated annotation for training:\n",
    "\n",
    "Then, this entire dataset should be tokenized and the 'input_id' should be extracted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/12460 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1500 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/500 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def tokenize_function(example):\n",
    "    start_prompt = 'Produce a summary of the following dialogue.\\n\\n'\n",
    "    end_prompt = '\\n\\nSummary: '\n",
    "    prompt = [start_prompt + dialogue + end_prompt for dialogue in example[\"dialogue\"]]\n",
    "    example['input_ids'] = tokenizer(prompt, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "    example['labels'] = tokenizer(example[\"summary\"], padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "    \n",
    "    return example\n",
    "\n",
    "# set \"batched=True\" because dataset is already divided into training, validation, and test splits\n",
    "# training data: original exemplars used to train the model\n",
    "# validation data: used during model training to fine-tune parameters across iterations\n",
    "# test data: used after all training is complete to give final model performance\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "tokenized_datasets = tokenized_datasets.remove_columns(['id', 'topic', 'dialogue', 'summary',])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, we can just take a subset of the data by taking only every 100th entry. \n",
    "In the code below, \"index % 100 == 0\" is a Boolean expression that looks for entries whose indices are divisible by 100 -- i.e., the remainder after dividing by 100 is zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at C:\\Users\\David Abugaber\\.cache\\huggingface\\datasets\\knkarthick___csv\\knkarthick--dialogsum-c8fac5d84cd35861\\0.0.0\\6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1\\cache-4f905de34fd2d5b2.arrow\n",
      "Loading cached processed dataset at C:\\Users\\David Abugaber\\.cache\\huggingface\\datasets\\knkarthick___csv\\knkarthick--dialogsum-c8fac5d84cd35861\\0.0.0\\6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1\\cache-b568002b0437a33e.arrow\n",
      "Loading cached processed dataset at C:\\Users\\David Abugaber\\.cache\\huggingface\\datasets\\knkarthick___csv\\knkarthick--dialogsum-c8fac5d84cd35861\\0.0.0\\6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1\\cache-8b96e13cd31466b7.arrow\n"
     ]
    }
   ],
   "source": [
    "subsample_tokenized_datasets = tokenized_datasets.filter(lambda example, index: index % 100 == 0, with_indices=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the shape of the dataset that we'll be working on: a train dataset with 125 entries, a test dataset with 15 entries, and a validation dataset with 5 entries. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['input_ids', 'labels'],\n",
      "        num_rows: 125\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['input_ids', 'labels'],\n",
      "        num_rows: 15\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['input_ids', 'labels'],\n",
      "        num_rows: 5\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(subsample_tokenized_datasets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each entry consists of a string of 'input_ids' and a string of 'labels.' As before, the raw values are just numerical representations corresponding to each token..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[947, 19, 3, 9, 7478, 24, 523, 12, 36, 21603, 26, 5, 1713, 345, 13515, 536, 4663, 10, 2018, 6, 1363, 5, 3931, 5, 27, 31, 51, 7582, 12833, 77, 7, 5, 1615, 33, 25, 270, 469, 58, 1713, 345, 13515, 357, 4663, 10, 27, 435, 34, 133, 36, 3, 9, 207, 800, 12, 129, 3, 9, 691, 18, 413, 5, 1713, 345, 13515, 536, 4663, 10, 2163, 6, 168, 6, 25, 43, 29, 31, 17, 141, 80, 21, 305, 203, 5, 148, 225, 43, 80, 334, 215, 5, 1713, 345, 13515, 357, 4663, 10, 27, 214, 5, 27, 2320, 38, 307, 38, 132, 19, 1327, 1786, 6, 572, 281, 217, 8, 2472, 58, 1713, 345, 13515, 536, 4663, 10, 1548, 6, 8, 200, 194, 12, 1792, 2261, 21154, 19, 12, 253, 91, 81, 135, 778, 5, 264, 653, 12, 369, 44, 709, 728, 3, 9, 215, 21, 39, 293, 207, 5, 1713, 345, 13515, 357, 4663, 10, 8872, 5, 1713, 345, 13515, 536, 4663, 10, 1563, 140, 217, 270, 5, 696, 2053, 11, 11581, 320, 1399, 5, 2321, 3, 9, 1659, 6522, 6, 754, 5, 531, 25, 7269, 6, 1363, 5, 3931, 58, 1713, 345, 13515, 357, 4663, 10, 2163, 5, 1713, 345, 13515, 536, 4663, 10, 14627, 53, 19, 8, 1374, 1137, 13, 5084, 1874, 11, 842, 1994, 6, 25, 214, 5, 148, 310, 225, 10399, 5, 1713, 345, 13515, 357, 4663, 10, 27, 31, 162, 1971, 3986, 13, 648, 6, 68, 27, 131, 54, 31, 17, 1727, 12, 4583, 8, 7386, 5, 1713, 345, 13515, 536, 4663, 10, 1548, 6, 62, 43, 2287, 11, 128, 11208, 24, 429, 199, 5, 27, 31, 195, 428, 25, 72, 251, 274, 25, 1175, 5, 1713, 345, 13515, 357, 4663, 10, 8872, 6, 2049, 2472, 5, 8733, 3, 9, 9251, 270, 10, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[1363, 5, 3931, 31, 7, 652, 3, 9, 691, 18, 413, 6, 11, 7582, 12833, 77, 7, 7786, 7, 376, 12, 43, 80, 334, 215, 5, 12833, 77, 7, 31, 195, 428, 128, 251, 81, 70, 2287, 11, 11208, 12, 199, 1363, 5, 3931, 10399, 10257, 5, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "print(subsample_tokenized_datasets['train']['input_ids'][0])\n",
    "print(subsample_tokenized_datasets['train']['labels'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "We can use tokenizer.decode() to turn these back into actual human-legible text strings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here is a dialogue that needs to be summarized. #Person1#: Hi, Mr. Smith. I'm Doctor Hawkins. Why are you here today? #Person2#: I found it would be a good idea to get a check-up. #Person1#: Yes, well, you haven't had one for 5 years. You should have one every year. #Person2#: I know. I figure as long as there is nothing wrong, why go see the doctor? #Person1#: Well, the best way to avoid serious illnesses is to find out about them early. So try to come at least once a year for your own good. #Person2#: Ok. #Person1#: Let me see here. Your eyes and ears look fine. Take a deep breath, please. Do you smoke, Mr. Smith? #Person2#: Yes. #Person1#: Smoking is the leading cause of lung cancer and heart disease, you know. You really should quit. #Person2#: I've tried hundreds of times, but I just can't seem to kick the habit. #Person1#: Well, we have classes and some medications that might help. I'll give you more information before you leave. #Person2#: Ok, thanks doctor. Write a summary here: \n",
      "Mr. Smith's getting a check-up, and Doctor Hawkins advises him to have one every year. Hawkins'll give some information about their classes and medications to help Mr. Smith quit smoking.\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(subsample_tokenized_datasets['train']['input_ids'][0],skip_special_tokens=True))\n",
    "print(tokenizer.decode(subsample_tokenized_datasets['train']['labels'][0],skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='5.2'></a>\n",
    "### 5.2 - Fine-tune the model updating all parameters\n",
    "\n",
    "We first need to establish the training configuration.\n",
    "\n",
    "Here is a quick run-down of what each argument means:\n",
    "    \n",
    "Learning rate: determines how quickly model weights are updated during training. A higher rate means faster convergence, but if it's too high the results can become unstable, overshooting or oscillating around the solution\n",
    "\n",
    "Number of training epochs: how many training epochs to perform? (one epoch is one full cycle through the training data, so this sets the total number of times the model sees the entire training dataset)\n",
    "\n",
    "Weight decay: a parameter that limits overfitting by penalizing large weights in the model\n",
    "\n",
    "Logging steps: controls how frequently training progress is displayed during training \n",
    "\n",
    "Max steps: How many training steps to perform? (one step is one gradient update to the model steps, iterating one \"batch.\" So, this sets the total number of batches processed during training)\n",
    "\n",
    "Save strategy: can be set to either \"no\" (no saving at all), \"epoch\" (to save at the end of each epoch\"), or \"steps\" (to save after a given number of steps, set by including a save_steps = # argument in your call to TrainingArguments )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\David Abugaber\\AppData\\Roaming\\Python\\Python39\\site-packages\\transformers\\optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:04, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>49.190900</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1, training_loss=49.19086456298828, metrics={'train_runtime': 61.4805, 'train_samples_per_second': 0.13, 'train_steps_per_second': 0.016, 'total_flos': 5478058819584.0, 'train_loss': 49.19086456298828, 'epoch': 0.06})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_dir = f'./dialog_summary_model_training'\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    learning_rate=1e-5,\n",
    "    num_train_epochs=1,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps=1,\n",
    "    max_steps=1,\n",
    "    save_strategy='epoch'\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=subsample_tokenized_datasets['train'],\n",
    "    eval_dataset=subsample_tokenized_datasets['validation']\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now load the trained model. Since in the previous step we set it to save after each training epoch, and there was only one epoch, then by default the directory should be called 'checkpoint-1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load trained model\n",
    "model_path = \"dialog_summary_model_training/checkpoint-1\"\n",
    "\n",
    "instruct_model = AutoModelForSeq2SeqLM.from_pretrained(model_path, torch_dtype=torch.bfloat16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='5.3'></a>\n",
    "### 5.3 - Visually inspect model results\n",
    "\n",
    "Let's check if the model is producing reasonably good output. Of course, it's not feasible to manually check every output (especially in a full model), but eyeballing the results qualitatively can give us an initial sense of its performance before turning to quantitative metrics (see below).\n",
    "\n",
    "Note that here we are looking at the original, unsampled Dialogsum dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dialogue:\n",
      "#Person1#: Steven, I need badly your help.\n",
      "#Person2#: What's the matter?\n",
      "#Person1#: My wife has found that I have an affair with my secretary, and now she is going to divorce me.\n",
      "#Person2#: How could you cheat on your wife? You have been married for ten years.\n",
      "#Person1#: Yes, I know I'm wrong. But I swear that the affair lasts only for two months. And I still love my wife. I couldn't live without her.\n",
      "#Person2#: I will try my best to persuade her to reconsider the divorce. But are you sure that from now on you will be faithful to her forever?\n",
      "#Person1#: Yes, I swear.\n",
      "Human-generated summary:\n",
      "Steve will try to persuade #Person1#'s wife not to divorce #Person1# as #Person1# swears to remain faithful forever.\n",
      "Original model:\n",
      "Steven is asking for help.\n",
      "Instruct model:\n",
      "Person1 is going to divorce her secretary.\n"
     ]
    }
   ],
   "source": [
    "index = 25\n",
    "dialogue = dataset['test'][index]['dialogue']\n",
    "human_summary = dataset['test'][index]['summary']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "\n",
    "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
    "\n",
    "original_model_outputs = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))\n",
    "original_model_text_output = tokenizer.decode(original_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "instruct_model_outputs = instruct_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))\n",
    "instruct_model_text_output = tokenizer.decode(instruct_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "print(f'Dialogue:\\n{dialogue}')\n",
    "print(f'Human-generated summary:\\n{human_summary}')\n",
    "print(f'Original model:\\n{original_model_text_output}')\n",
    "print(f'Instruct model:\\n{instruct_model_text_output}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seems like fine-tuning the model by embedding explicit instructions within the prompts has improved the results somewhat."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='5.3'></a>\n",
    "### 5.4 - Automatically quantify model results with ROUGE metrics\n",
    "\n",
    "Of course, it's not feasible to manually check every single summary and compare it against a baseline human annotation. Beyond the fact that such an approach is not scalable (why automate the dialogue summaries in the first place if a human has to validate every output?), there's the fact that you can't easily quantify how good the summary is just by glancing at it.\n",
    "\n",
    "One way to tackle this problem is with ROUGE, or \"Recall-Oriented Understudy for Gisting Evaluation.\" This is a family of metrics that quantifies the similarity between computer-generated text vs. a human baseline.\n",
    "\n",
    "ROUGE can take the form of \"ROUGE-N\", where N is the granularity of the n-grams to compare. So for instance, ROUGE-1 compares unigrams (individual words), ROUGE-2 compares bigrams, etc.\n",
    "\n",
    "Alternately, ROUGE-L compares n grams with the length of the longest common subsequence between the generated vs. reference text\n",
    "\n",
    "The score itself can either be recall (# of matching n-grams divided by # of n-grams in the reference text), precision (# of matching n-grams divided by # of n-grams in the text to be tested), or F-score (also known as F1), which is the harmonic mean of precision and recall. Which of these to use depends on your use case.\n",
    "\n",
    "If it's more important that your output captures all the \"right answers\" (minimizing false negatives) even if you inadvertently generate irrelevant text, then use recall.\n",
    "If it's more important that nothing in your output is wrong (minimizing false positives) even if you omit possibly relevant text, then use precision.\n",
    "If you want to balance these two considerations, use F-score.\n",
    "\n",
    "Here is a toy example of ROUGE in action, from the rouge-score package:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rouge1: Score(precision=0.7142857142857143, recall=0.8333333333333334, fmeasure=0.7692307692307692)\n",
      "rouge2: Score(precision=0.5, recall=0.6, fmeasure=0.5454545454545454)\n",
      "rougeL: Score(precision=0.7142857142857143, recall=0.8333333333333334, fmeasure=0.7692307692307692)\n"
     ]
    }
   ],
   "source": [
    "#pip install rouge-score\n",
    "\n",
    "from rouge_score import rouge_scorer\n",
    "\n",
    "# Sample generated and reference texts\n",
    "generated_text = \"The cat sat on the mat\"\n",
    "reference_text = \"The cat was sitting on the mat\"\n",
    "\n",
    "#note: the longest common subsequence is \"on the mat\", so in this example ROUGE-L is essentially ROUGE-3\n",
    "\n",
    "# Initialize the ROUGE scorer\n",
    "scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)\n",
    "\n",
    "# Calculate ROUGE scores\n",
    "scores = scorer.score(generated_text, reference_text)\n",
    "\n",
    "# Print the ROUGE scores\n",
    "for key in scores:\n",
    "    print(f'{key}: {scores[key]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first thing we want to do to calculate ROUGE scores is to generate a dataset with summaries for our human-annotated baseline, for our original (non-finetuned) model, and for our finetuned \"instruct model.\" The code below produces such a dataset for the first 20 entries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>human_baseline_summaries</th>\n",
       "      <th>original_model_summaries</th>\n",
       "      <th>instruct_model_summaries</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Ms. Dawson helps #Person1# to write a memo to ...</td>\n",
       "      <td>Employee memo is due to go out by this afternoon.</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>In order to prevent employees from wasting tim...</td>\n",
       "      <td>The memo is being distributed to all employees...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Ms. Dawson takes a dictation for #Person1# abo...</td>\n",
       "      <td>The memo is for the following: \"It is a dictat...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>#Person2# arrives late because of traffic jam....</td>\n",
       "      <td>#Person1: I'm happy to see you finally here. #...</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>#Person2# decides to follow #Person1#'s sugges...</td>\n",
       "      <td>The traffic was terrible.</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>#Person2# complains to #Person1# about the tra...</td>\n",
       "      <td>People are complaining about the congestion in...</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>#Person1# tells Kate that Masha and Hero get d...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>#Person1# tells Kate that Masha and Hero are g...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>#Person1# and Kate talk about the divorce betw...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>#Person1# and Brian are at the birthday party ...</td>\n",
       "      <td>#Person1#: I'm so happy to see you have a part...</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>#Person1# attends Brian's birthday party. Bria...</td>\n",
       "      <td>The party was a success for the person who had...</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>#Person1# has a dance with Brian at Brian's bi...</td>\n",
       "      <td>Brian is celebrating his birthday.</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>#Person1# is surprised at the Olympic Stadium'...</td>\n",
       "      <td>The Olympic stadium is the center of the park.</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>#Person2# shows #Person1# around the construct...</td>\n",
       "      <td>The Olympic park is the center of the park.</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>#Person2# introduces the Olympic Stadium's fin...</td>\n",
       "      <td>The Olympic park is very big, the center of th...</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>#Person1# wants to create a company and is goi...</td>\n",
       "      <td>#Person1: I am done working for a company that...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>#Person1# abandons the idea of creating a comp...</td>\n",
       "      <td>#Person1: I'm done! I'm going to write a busin...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>#Person1# wants to start #Person1#'s own busin...</td>\n",
       "      <td>Describe your company, its history, and its po...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>#Person2# feels itchy. #Person1# doubts it is ...</td>\n",
       "      <td>Someone is having a bad day.</td>\n",
       "      <td>Person1 is scratching a lot.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                             human_baseline_summaries  \\\n",
       "0   Ms. Dawson helps #Person1# to write a memo to ...   \n",
       "1   In order to prevent employees from wasting tim...   \n",
       "2   Ms. Dawson takes a dictation for #Person1# abo...   \n",
       "3   #Person2# arrives late because of traffic jam....   \n",
       "4   #Person2# decides to follow #Person1#'s sugges...   \n",
       "5   #Person2# complains to #Person1# about the tra...   \n",
       "6   #Person1# tells Kate that Masha and Hero get d...   \n",
       "7   #Person1# tells Kate that Masha and Hero are g...   \n",
       "8   #Person1# and Kate talk about the divorce betw...   \n",
       "9   #Person1# and Brian are at the birthday party ...   \n",
       "10  #Person1# attends Brian's birthday party. Bria...   \n",
       "11  #Person1# has a dance with Brian at Brian's bi...   \n",
       "12  #Person1# is surprised at the Olympic Stadium'...   \n",
       "13  #Person2# shows #Person1# around the construct...   \n",
       "14  #Person2# introduces the Olympic Stadium's fin...   \n",
       "15  #Person1# wants to create a company and is goi...   \n",
       "16  #Person1# abandons the idea of creating a comp...   \n",
       "17  #Person1# wants to start #Person1#'s own busin...   \n",
       "18  #Person2# feels itchy. #Person1# doubts it is ...   \n",
       "\n",
       "                             original_model_summaries  \\\n",
       "0   Employee memo is due to go out by this afternoon.   \n",
       "1   The memo is being distributed to all employees...   \n",
       "2   The memo is for the following: \"It is a dictat...   \n",
       "3   #Person1: I'm happy to see you finally here. #...   \n",
       "4                           The traffic was terrible.   \n",
       "5   People are complaining about the congestion in...   \n",
       "6                Masha and Hero are getting divorced.   \n",
       "7                Masha and Hero are getting divorced.   \n",
       "8                Masha and Hero are getting divorced.   \n",
       "9   #Person1#: I'm so happy to see you have a part...   \n",
       "10  The party was a success for the person who had...   \n",
       "11                 Brian is celebrating his birthday.   \n",
       "12     The Olympic stadium is the center of the park.   \n",
       "13        The Olympic park is the center of the park.   \n",
       "14  The Olympic park is very big, the center of th...   \n",
       "15  #Person1: I am done working for a company that...   \n",
       "16  #Person1: I'm done! I'm going to write a busin...   \n",
       "17  Describe your company, its history, and its po...   \n",
       "18                       Someone is having a bad day.   \n",
       "\n",
       "                             instruct_model_summaries  \n",
       "0   The memo will go out to all employees by this ...  \n",
       "1   The memo will go out to all employees by this ...  \n",
       "2   The memo will go out to all employees by this ...  \n",
       "3   The traffic jam at the Carrefour intersection ...  \n",
       "4   The traffic jam at the Carrefour intersection ...  \n",
       "5   The traffic jam at the Carrefour intersection ...  \n",
       "6                Masha and Hero are getting divorced.  \n",
       "7                Masha and Hero are getting divorced.  \n",
       "8                Masha and Hero are getting divorced.  \n",
       "9                      Brian's birthday is coming up.  \n",
       "10                     Brian's birthday is coming up.  \n",
       "11                     Brian's birthday is coming up.  \n",
       "12  The Olympic stadium is the center of the park....  \n",
       "13  The Olympic stadium is the center of the park....  \n",
       "14  The Olympic stadium is the center of the park....  \n",
       "15  #Person1#: I am done working for a company tha...  \n",
       "16  #Person1#: I am done working for a company tha...  \n",
       "17  #Person1#: I am done working for a company tha...  \n",
       "18                       Person1 is scratching a lot.  "
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dialogues = dataset['test'][0:19]['dialogue']\n",
    "human_baseline_summaries = dataset['test'][0:19]['summary']\n",
    "\n",
    "original_model_summaries = []\n",
    "instruct_model_summaries = []\n",
    "\n",
    "for _, dialogue in enumerate(dialogues):\n",
    "    prompt = f\"\"\"\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary: \"\"\"\n",
    "    input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
    "\n",
    "    original_model_outputs = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200))\n",
    "    original_model_text_output = tokenizer.decode(original_model_outputs[0], skip_special_tokens=True)\n",
    "    original_model_summaries.append(original_model_text_output)\n",
    "\n",
    "    instruct_model_outputs = instruct_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200))\n",
    "    instruct_model_text_output = tokenizer.decode(instruct_model_outputs[0], skip_special_tokens=True)\n",
    "    instruct_model_summaries.append(instruct_model_text_output)\n",
    "    \n",
    "zipped_summaries = list(zip(human_baseline_summaries, original_model_summaries, instruct_model_summaries))\n",
    " \n",
    "df = pd.DataFrame(zipped_summaries, columns = ['human_baseline_summaries', 'original_model_summaries', 'instruct_model_summaries'])\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to get ROUGE scores for the original vs. instruct models. Note that \"use_stemmer=True\" removes suffixes from words to improve matching. So for instance, \"sleeping\" and \"sleeps\" would be counted as the same word after removing the -ing and -s suffixes. Setting \"use_aggregator=True\" gives you an aggregate ROUGE score across all the summaries, instead of a separate score for each one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original model:\n",
      "{'rouge1': 0.2093862423370193, 'rouge2': 0.07284353971171693, 'rougeL': 0.1825140348261415, 'rougeLsum': 0.1820160064971703}\n",
      "Instruct model:\n",
      "{'rouge1': 0.29293514576360014, 'rouge2': 0.11264452616822107, 'rougeL': 0.2449479682766204, 'rougeLsum': 0.24590160739029454}\n"
     ]
    }
   ],
   "source": [
    "rouge = evaluate.load('rouge')\n",
    "\n",
    "original_model_results = rouge.compute(\n",
    "    predictions=original_model_summaries,\n",
    "    references=human_baseline_summaries[0:len(original_model_summaries)],\n",
    "    use_aggregator=True,\n",
    "    use_stemmer=True,\n",
    ")\n",
    "\n",
    "instruct_model_results = rouge.compute(\n",
    "    predictions=instruct_model_summaries,\n",
    "    references=human_baseline_summaries[0:len(instruct_model_summaries)],\n",
    "    use_aggregator=True,\n",
    "    use_stemmer=True,\n",
    ")\n",
    "\n",
    "print('Original model:')\n",
    "print(original_model_results)\n",
    "print('Instruct model:')\n",
    "print(instruct_model_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "We got improvements across all measures! (Note that rougeLsum looks across sentences, whereas rougeL averages over individual sentences). \n",
    "\n",
    "Of course, it's possible that we could have gotten an even larger improvement by tweaking the instructions in our preprocessed training data\n",
    "\n",
    "For instance, instead of \"Here is a dialogue that needs to be summarized... Write a summary here:\", we might use variations like \"Summarize the dialogue below... Summary:\" or similar."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='6'></a>\n",
    "## 6 - Save compute resources with Parameter Efficient Fine-Tuning (PEFT)\n",
    "\n",
    "The approach used above adjusted the model parameter weights \"wholesale,\" essentially creating an entire new copy of the model.\n",
    "\n",
    "However, since we're simply trying to improve the model's performance on a very targeted task (summarizing dialogues), it is somewhat inefficient to tweak ALL of the model parameters.\n",
    "\n",
    "An alternative approach is to use Parameter Efficient Fine-Tuning (PEFT), which improves a model by either selecting only a subset of parameters for fine-tuning (selective PEFT); adding additional layers or parameters to the base model, for instance, by appending virtual tokens to the input to nudge the output towards a desired result (additive PEFT); or by training a low-dimensional representation of the network that is added to the original weights (reparameterization).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='6.1'></a>\n",
    "### 6.1 - Setting up LoRA\n",
    "\n",
    "The approach I take here uses one such reparameterization method called Low-Rank Adaptation, or LoRA. The LoRA \"adapter\" that is trained is orders of magnitude smaller than the original LLM.\n",
    "\n",
    "The first step is to set up the LoRA adapter. In the code below, I sent r(ank) to 8; this refers to the size of the low-dimensional representation that actually gets trained. Lora_alpha is a scaling factor that determines how much weight is assigned to the LoRA activations vs. to the base model. I apply LoRA to the target modules q(uery) and v(alue), which are matrices in the attention blocks of the FLAN-T5 model. We can also try to eke out better task performance by training bias vectors (which determine unit activation regardless of input). The Huggingface docs recommend \"setting bias to None first, and then lora_only, before trying all.\" The dropout parameter prevents overfitting randomly \"drops out\" (i.e., sets to zero) a fraction of the units in the model. Since FLAN is a sequence-to-sequence (encoder-decoder) model, we set task type to \"TaskType.SEQ_2_SEQ_LM.\" You can read more about the LoRA parameters here: https://huggingface.co/docs/peft/conceptual_guides/lora"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=8, \n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q\", \"v\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to actually add the LoRA layer to our original LLM..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable model parameters: 884736\n",
      "all model parameters: 248462592\n",
      "percentage of trainable model parameters: 0.36%\n"
     ]
    }
   ],
   "source": [
    "peft_model = get_peft_model(model, \n",
    "                            lora_config)\n",
    "print(print_number_of_trainable_model_parameters(peft_model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how only a tiny proportion of our model parameters are now \"trainable.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='6.2'></a>\n",
    "### 6.2 - Training our model with LoRA\n",
    "\n",
    "As before, we'll have to define a set of training arguments before performing the training itself.\n",
    "\n",
    "Instead of saving the whole model, we'll only save the LoRA adapter, hence no need to set \"saving_strategy='epoch'\" as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\David Abugaber\\AppData\\Roaming\\Python\\Python39\\site-packages\\transformers\\optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:00, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>48.130100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1, training_loss=48.13014602661133, metrics={'train_runtime': 54.4416, 'train_samples_per_second': 0.147, 'train_steps_per_second': 0.018, 'total_flos': 5499802091520.0, 'train_loss': 48.13014602661133, 'epoch': 0.06})"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_dir = f'./dialog_summary_model_lora'\n",
    "\n",
    "peft_training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    auto_find_batch_size=True,\n",
    "    learning_rate=1e-3,\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=1,\n",
    "    max_steps=1\n",
    ")\n",
    "    \n",
    "peft_trainer = Trainer(\n",
    "    model=peft_model,\n",
    "    args=peft_training_args,\n",
    "    train_dataset=subsample_tokenized_datasets[\"train\"],\n",
    ")\n",
    "\n",
    "peft_trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./dialog_summary_model_lora\\\\tokenizer_config.json',\n",
       " './dialog_summary_model_lora\\\\special_tokens_map.json',\n",
       " './dialog_summary_model_lora\\\\tokenizer.json')"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_trainer.model.save_pretrained(output_dir)\n",
    "\n",
    "tokenizer.save_pretrained(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll load a fresh copy of the FLAN-T5 model and combine it with the LoRA adapter we just trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "peft_model_base = AutoModelForSeq2SeqLM.from_pretrained(\"google/flan-t5-base\", torch_dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google/flan-t5-base\")\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(peft_model_base, \n",
    "                                       output_dir, \n",
    "                                       torch_dtype=torch.bfloat16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='6.3'></a>\n",
    "### 6.3 - Visually inspect LoRA model results\n",
    "\n",
    "As before, let's take a glance at how our model performs.\n",
    "\n",
    "We can compare our LORA-tuned model with the original base model as well as the model that used full (but computationally costly) tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dialogue:\n",
      "#Person1#: Ms. Dawson, I need you to take a dictation for me.\n",
      "#Person2#: Yes, sir...\n",
      "#Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready?\n",
      "#Person2#: Yes, sir. Go ahead.\n",
      "#Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited.\n",
      "#Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications?\n",
      "#Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications.\n",
      "#Person2#: But sir, many employees use Instant Messaging to communicate with their clients.\n",
      "#Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we?\n",
      "#Person2#: This applies to internal and external communications.\n",
      "#Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads.\n",
      "#Person2#: Is that all?\n",
      "#Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm.\n",
      "Human-generated summary:\n",
      "In order to prevent employees from wasting time on Instant Message programs, #Person1# decides to terminate the use of those programs and asks Ms. Dawson to send out a memo to all employees by the afternoon.\n",
      "Original model:\n",
      "Ms. Dawson's memo to all employees is due to go out this afternoon.\n",
      "Instruct model:\n",
      "The memo will go out to all employees by this afternoon.\n",
      "LoRA model:\n",
      "The memo will go out to all employees by this afternoon.\n"
     ]
    }
   ],
   "source": [
    "index = 1\n",
    "dialogue = dataset['test'][index]['dialogue']\n",
    "human_summary = dataset['test'][index]['summary']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Produce a summary of the following dialogue.\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "\n",
    "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
    "\n",
    "original_model_outputs = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))\n",
    "original_model_text_output = tokenizer.decode(original_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "instruct_model_outputs = instruct_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))\n",
    "instruct_model_text_output = tokenizer.decode(instruct_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "peft_model_outputs = peft_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))\n",
    "peft_model_text_output = tokenizer.decode(peft_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "\n",
    "print(f'Dialogue:\\n{dialogue}')\n",
    "print(f'Human-generated summary:\\n{human_summary}')\n",
    "print(f'Original model:\\n{original_model_text_output}')\n",
    "print(f'Instruct model:\\n{instruct_model_text_output}')\n",
    "print(f'LoRA model:\\n{peft_model_text_output}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the LoRA model doesn't perform as well as the fully-trained one, at least bear in mind that LoRA was a shortcut to training that was much more efficient in terms of compute..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='6.4'></a>\n",
    "### 6.4 - Assessing LoRA model with ROUGE metrics\n",
    "As before, we can quantify how well our model is performing by calculating ROUGE metrics. The first step is to create a dataframe with model summaries from our original, full fine-tuned, and LoRA models. I'm using the same syntax as before for the LoRA model outputs. Since I already generated the completions for the first two models, no need to redo them (though I'm leaving in the code, commented out, just for comparison):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>human_baseline_summaries</th>\n",
       "      <th>original_model_summaries</th>\n",
       "      <th>instruct_model_summaries</th>\n",
       "      <th>peft_model_summaries</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Ms. Dawson helps #Person1# to write a memo to ...</td>\n",
       "      <td>Employee memo is due to go out by this afternoon.</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>In order to prevent employees from wasting tim...</td>\n",
       "      <td>The memo is being distributed to all employees...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Ms. Dawson takes a dictation for #Person1# abo...</td>\n",
       "      <td>The memo is for the following: \"It is a dictat...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "      <td>The memo will go out to all employees by this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>#Person2# arrives late because of traffic jam....</td>\n",
       "      <td>#Person1: I'm happy to see you finally here. #...</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "      <td>The traffic jam near the Carrefour intersectio...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>#Person2# decides to follow #Person1#'s sugges...</td>\n",
       "      <td>The traffic was terrible.</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "      <td>The traffic jam near the Carrefour intersectio...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>#Person2# complains to #Person1# about the tra...</td>\n",
       "      <td>People are complaining about the congestion in...</td>\n",
       "      <td>The traffic jam at the Carrefour intersection ...</td>\n",
       "      <td>The traffic jam near the Carrefour intersectio...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>#Person1# tells Kate that Masha and Hero get d...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>#Person1# tells Kate that Masha and Hero are g...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>#Person1# and Kate talk about the divorce betw...</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "      <td>Masha and Hero are getting divorced.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>#Person1# and Brian are at the birthday party ...</td>\n",
       "      <td>#Person1#: I'm so happy to see you have a part...</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>#Person1# attends Brian's birthday party. Bria...</td>\n",
       "      <td>The party was a success for the person who had...</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>#Person1# has a dance with Brian at Brian's bi...</td>\n",
       "      <td>Brian is celebrating his birthday.</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "      <td>Brian's birthday is coming up.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>#Person1# is surprised at the Olympic Stadium'...</td>\n",
       "      <td>The Olympic stadium is the center of the park.</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "      <td>The Olympic park is so big that there are 5000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>#Person2# shows #Person1# around the construct...</td>\n",
       "      <td>The Olympic park is the center of the park.</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "      <td>The Olympic park is so big that there are 5000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>#Person2# introduces the Olympic Stadium's fin...</td>\n",
       "      <td>The Olympic park is very big, the center of th...</td>\n",
       "      <td>The Olympic stadium is the center of the park....</td>\n",
       "      <td>The Olympic park is so big that there are 5000...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>#Person1# wants to create a company and is goi...</td>\n",
       "      <td>#Person1: I am done working for a company that...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "      <td>Person1 is going to start a business. He is go...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>#Person1# abandons the idea of creating a comp...</td>\n",
       "      <td>#Person1: I'm done! I'm going to write a busin...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "      <td>Person1 is going to start a business. He is go...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>#Person1# wants to start #Person1#'s own busin...</td>\n",
       "      <td>Describe your company, its history, and its po...</td>\n",
       "      <td>#Person1#: I am done working for a company tha...</td>\n",
       "      <td>Person1 is going to start a business. He is go...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>#Person2# feels itchy. #Person1# doubts it is ...</td>\n",
       "      <td>Someone is having a bad day.</td>\n",
       "      <td>Person1 is scratching a lot.</td>\n",
       "      <td>Person1 is scratching so much that he feels li...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                             human_baseline_summaries  \\\n",
       "0   Ms. Dawson helps #Person1# to write a memo to ...   \n",
       "1   In order to prevent employees from wasting tim...   \n",
       "2   Ms. Dawson takes a dictation for #Person1# abo...   \n",
       "3   #Person2# arrives late because of traffic jam....   \n",
       "4   #Person2# decides to follow #Person1#'s sugges...   \n",
       "5   #Person2# complains to #Person1# about the tra...   \n",
       "6   #Person1# tells Kate that Masha and Hero get d...   \n",
       "7   #Person1# tells Kate that Masha and Hero are g...   \n",
       "8   #Person1# and Kate talk about the divorce betw...   \n",
       "9   #Person1# and Brian are at the birthday party ...   \n",
       "10  #Person1# attends Brian's birthday party. Bria...   \n",
       "11  #Person1# has a dance with Brian at Brian's bi...   \n",
       "12  #Person1# is surprised at the Olympic Stadium'...   \n",
       "13  #Person2# shows #Person1# around the construct...   \n",
       "14  #Person2# introduces the Olympic Stadium's fin...   \n",
       "15  #Person1# wants to create a company and is goi...   \n",
       "16  #Person1# abandons the idea of creating a comp...   \n",
       "17  #Person1# wants to start #Person1#'s own busin...   \n",
       "18  #Person2# feels itchy. #Person1# doubts it is ...   \n",
       "\n",
       "                             original_model_summaries  \\\n",
       "0   Employee memo is due to go out by this afternoon.   \n",
       "1   The memo is being distributed to all employees...   \n",
       "2   The memo is for the following: \"It is a dictat...   \n",
       "3   #Person1: I'm happy to see you finally here. #...   \n",
       "4                           The traffic was terrible.   \n",
       "5   People are complaining about the congestion in...   \n",
       "6                Masha and Hero are getting divorced.   \n",
       "7                Masha and Hero are getting divorced.   \n",
       "8                Masha and Hero are getting divorced.   \n",
       "9   #Person1#: I'm so happy to see you have a part...   \n",
       "10  The party was a success for the person who had...   \n",
       "11                 Brian is celebrating his birthday.   \n",
       "12     The Olympic stadium is the center of the park.   \n",
       "13        The Olympic park is the center of the park.   \n",
       "14  The Olympic park is very big, the center of th...   \n",
       "15  #Person1: I am done working for a company that...   \n",
       "16  #Person1: I'm done! I'm going to write a busin...   \n",
       "17  Describe your company, its history, and its po...   \n",
       "18                       Someone is having a bad day.   \n",
       "\n",
       "                             instruct_model_summaries  \\\n",
       "0   The memo will go out to all employees by this ...   \n",
       "1   The memo will go out to all employees by this ...   \n",
       "2   The memo will go out to all employees by this ...   \n",
       "3   The traffic jam at the Carrefour intersection ...   \n",
       "4   The traffic jam at the Carrefour intersection ...   \n",
       "5   The traffic jam at the Carrefour intersection ...   \n",
       "6                Masha and Hero are getting divorced.   \n",
       "7                Masha and Hero are getting divorced.   \n",
       "8                Masha and Hero are getting divorced.   \n",
       "9                      Brian's birthday is coming up.   \n",
       "10                     Brian's birthday is coming up.   \n",
       "11                     Brian's birthday is coming up.   \n",
       "12  The Olympic stadium is the center of the park....   \n",
       "13  The Olympic stadium is the center of the park....   \n",
       "14  The Olympic stadium is the center of the park....   \n",
       "15  #Person1#: I am done working for a company tha...   \n",
       "16  #Person1#: I am done working for a company tha...   \n",
       "17  #Person1#: I am done working for a company tha...   \n",
       "18                       Person1 is scratching a lot.   \n",
       "\n",
       "                                 peft_model_summaries  \n",
       "0   The memo will go out to all employees by this ...  \n",
       "1   The memo will go out to all employees by this ...  \n",
       "2   The memo will go out to all employees by this ...  \n",
       "3   The traffic jam near the Carrefour intersectio...  \n",
       "4   The traffic jam near the Carrefour intersectio...  \n",
       "5   The traffic jam near the Carrefour intersectio...  \n",
       "6                Masha and Hero are getting divorced.  \n",
       "7                Masha and Hero are getting divorced.  \n",
       "8                Masha and Hero are getting divorced.  \n",
       "9                      Brian's birthday is coming up.  \n",
       "10                     Brian's birthday is coming up.  \n",
       "11                     Brian's birthday is coming up.  \n",
       "12  The Olympic park is so big that there are 5000...  \n",
       "13  The Olympic park is so big that there are 5000...  \n",
       "14  The Olympic park is so big that there are 5000...  \n",
       "15  Person1 is going to start a business. He is go...  \n",
       "16  Person1 is going to start a business. He is go...  \n",
       "17  Person1 is going to start a business. He is go...  \n",
       "18  Person1 is scratching so much that he feels li...  "
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#dialogues = dataset['test'][0:19]['dialogue']\n",
    "#human_baseline_summaries = dataset['test'][0:19]['summary']\n",
    "\n",
    "#original_model_summaries = []\n",
    "#instruct_model_summaries = []\n",
    "peft_model_summaries = []\n",
    "\n",
    "for idx, dialogue in enumerate(dialogues):\n",
    "    prompt = f\"\"\"\n",
    "Here is a dialogue that needs to be summarized.\n",
    "\n",
    "{dialogue}\n",
    "\n",
    "Write a summary here: \"\"\"\n",
    "    \n",
    "    input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
    "   \n",
    "    #original_model_outputs = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200))\n",
    "    #original_model_text_output = tokenizer.decode(original_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "    #instruct_model_outputs = instruct_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200))\n",
    "    #instruct_model_text_output = tokenizer.decode(instruct_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "    peft_model_outputs = peft_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200))\n",
    "    peft_model_text_output = tokenizer.decode(peft_model_outputs[0], skip_special_tokens=True)\n",
    "\n",
    "    #original_model_summaries.append(original_model_text_output)\n",
    "    #instruct_model_summaries.append(instruct_model_text_output)\n",
    "    peft_model_summaries.append(peft_model_text_output)\n",
    "\n",
    "zipped_summaries = list(zip(human_baseline_summaries, original_model_summaries, instruct_model_summaries, peft_model_summaries))\n",
    " \n",
    "df = pd.DataFrame(zipped_summaries, columns = ['human_baseline_summaries', 'original_model_summaries', 'instruct_model_summaries', 'peft_model_summaries'])\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On to the ROUGE scores..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original model:\n",
      "{'rouge1': 0.2093862423370193, 'rouge2': 0.07284353971171693, 'rougeL': 0.1825140348261415, 'rougeLsum': 0.1820160064971703}\n",
      "Instruct model:\n",
      "{'rouge1': 0.29293514576360014, 'rouge2': 0.11264452616822107, 'rougeL': 0.2449479682766204, 'rougeLsum': 0.24590160739029454}\n",
      "LoRA model:\n",
      "{'rouge1': 0.29293514576360014, 'rouge2': 0.11264452616822107, 'rougeL': 0.2449479682766204, 'rougeLsum': 0.24590160739029454}\n"
     ]
    }
   ],
   "source": [
    "rouge = evaluate.load('rouge')\n",
    "\n",
    "original_model_results = rouge.compute(\n",
    "    predictions=original_model_summaries,\n",
    "    references=human_baseline_summaries[0:len(original_model_summaries)],\n",
    "    use_aggregator=True,\n",
    "    use_stemmer=True,\n",
    ")\n",
    "\n",
    "instruct_model_results = rouge.compute(\n",
    "    predictions=instruct_model_summaries,\n",
    "    references=human_baseline_summaries[0:len(instruct_model_summaries)],\n",
    "    use_aggregator=True,\n",
    "    use_stemmer=True,\n",
    ")\n",
    "\n",
    "peft_model_results = rouge.compute(\n",
    "    predictions=peft_model_summaries,\n",
    "    references=human_baseline_summaries[0:len(peft_model_summaries)],\n",
    "    use_aggregator=True,\n",
    "    use_stemmer=True,\n",
    ")\n",
    "\n",
    "\n",
    "print('Original model:')\n",
    "print(original_model_results)\n",
    "print('Instruct model:')\n",
    "print(instruct_model_results)\n",
    "print('LoRA model:')\n",
    "print(instruct_model_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The remarkable thing here is that LoRA produced very similar results as the instruct model, even though the method that it used was computationally much lighter. \n",
    "\n",
    "I can verify that the results weren't 100% identical by looking at the entry for the 19th row -- notice how the instruct and LoRA models do indeed differ (so it wasn't a bug in my code that made me reuse an identical model). However, since there is no difference in n-gram overlap with the human baselines in this case, the ROUGE scores ultimately come out the same despite these small differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#Person2# feels itchy. #Person1# doubts it is chicken pox and asks #Person2# to get away. #Person2# doesn't believe it.\n",
      "Person1 is scratching a lot.\n",
      "Person1 is scratching so much that he feels lightheaded and weak.\n"
     ]
    }
   ],
   "source": [
    "print(df['human_baseline_summaries'][18])\n",
    "print(df['instruct_model_summaries'][18])\n",
    "print(df['peft_model_summaries'][18])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7'></a>\n",
    "## 7 - Using Reinforcement Learning with Human Feedback (RLHF) for further refinement\n",
    "\n",
    "A common step in improving LLMs is using human-annotated feedback to further refine a trained model.\n",
    "\n",
    "First, humans label prompts along a given dimension, e.g., using the HHH criteria \"helpful / not helpful\", \"honest / not honest,\" and \"harmless / not harmless\"\n",
    "\n",
    "Then, a reinforcement learning model is trained based on these human ratings.\n",
    "\n",
    "Finally, the model itself feeds its output to this reinforcement learning model and uses the generated \"reward\" to update its weights.\n",
    "\n",
    "To make sure that the model doesn't just \"reward-hack\" at this stage (doing whatever it needs to to maximize reward from the reinforcement learning model but neglecting its original goal of carrying out the task at hand), the learning from this rewards is tempered by a metric that calculates how similar the updated model's output is to the original model's output.\n",
    "\n",
    "To begin, let's prepare the reward model: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7.1'></a>\n",
    "### 7.1 - Load the reward model\n",
    "\n",
    "My first attempt used a model created by Facebook that classifies text inputs as either \"hate\" or \"not hate\" -- see https://huggingface.co/facebook/roberta-hate-speech-dynabench-r4-target. The Huggingface website even has a nifty widget that lets you calculate a toxicity score for a given input sentence directly from the browser.\n",
    "\n",
    "However, this model is quite conservative in its ratings, perhaps because their definition of hate speech is targeted very narrowly to mean language against a particular demographic group, rather than just hateful language. For instance, a sentence like \"The Barbie movie was a disgusting piece of fake-feminist materialistic garbage\" yields a 85.7% probability of 'not hate.'\n",
    "\n",
    "By contrast, a separate model for categorizing text as \"toxic\" vs \"nontxic\" I found (https://huggingface.co/s-nlp/roberta_toxicity_classifier) gave a 99.9% probability of \"toxic\" to that same sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at s-nlp/roberta_toxicity_classifier were not used when initializing RobertaForSequenceClassification: ['roberta.pooler.dense.weight', 'roberta.pooler.dense.bias']\n",
      "- This IS expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 'neutral', 1: 'toxic'}\n"
     ]
    }
   ],
   "source": [
    "#reward_model_name = \"facebook/roberta-hate-speech-dynabench-r4-target\"\n",
    "reward_model_name = \"s-nlp/roberta_toxicity_classifier\"\n",
    "reward_tokenizer = AutoTokenizer.from_pretrained(reward_model_name, device_map=\"auto\")\n",
    "reward_model = AutoModelForSequenceClassification.from_pretrained(reward_model_name, device_map=\"auto\")\n",
    "print(reward_model.config.id2label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a toy example of how this works:    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Barbie movie exceeded my expectations.\n",
      "Logit values [not toxic, toxic]: [3.5713207721710205, -3.086273670196533]\n",
      "Probabilities [not toxic, toxic]: [0.9987174272537231, 0.00128258450422436]\n",
      "reward (high): [3.5713207721710205]\n",
      "\n",
      "The Barbie movie was a disgusting piece of fake-feminist materialistic garbage.\n",
      "Logit values [not toxic, toxic]: [0.9113463759422302, -0.8771229982376099]\n",
      "Probabilities [not toxic, toxic]: [0.8567395210266113, 0.14326049387454987]\n",
      "reward (low): [0.9113463759422302]\n"
     ]
    }
   ],
   "source": [
    "nontoxic_text = \"The Barbie movie exceeded my expectations.\"\n",
    "\n",
    "print(nontoxic_text)\n",
    "\n",
    "reward_input_ids = reward_tokenizer(nontoxic_text, return_tensors=\"pt\").input_ids\n",
    "\n",
    "logits = reward_model(input_ids=reward_input_ids).logits\n",
    "print(f'Logit values [not toxic, toxic]: {logits.tolist()[0]}')\n",
    "\n",
    "#convert logit values to probability\n",
    "probabilities = logits.softmax(dim=-1).tolist()[0]\n",
    "print(f'Probabilities [not toxic, toxic]: {probabilities}')\n",
    "\n",
    "#given the input logits, how much reward to give?\n",
    "nottoxic_index = 0\n",
    "nottoxic_reward = (logits[:, nottoxic_index]).tolist()\n",
    "print(f'reward (high): {nottoxic_reward}\\n')\n",
    "\n",
    "#rerun the example with toxic text\n",
    "toxic_text = \"The Barbie movie was a disgusting piece of fake-feminist materialistic garbage.\"\n",
    "\n",
    "print(toxic_text)\n",
    "\n",
    "reward_input_ids = reward_tokenizer(toxic_text, return_tensors=\"pt\").input_ids\n",
    "logits = reward_model(input_ids=reward_input_ids).logits\n",
    "print(f'Logit values [not toxic, toxic]: {logits.tolist()[0]}')\n",
    "\n",
    "#convert logit values to probability\n",
    "probabilities = logits.softmax(dim=-1).tolist()[0]\n",
    "print(f'Probabilities [not toxic, toxic]: {probabilities}')\n",
    "\n",
    "toxic_reward = (logits[:, nottoxic_index]).tolist()\n",
    "print(f'reward (low): {toxic_reward}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of this can be made simpler by creating a pipeline as follows (to be used in our code further below):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at s-nlp/roberta_toxicity_classifier were not used when initializing RobertaForSequenceClassification: ['roberta.pooler.dense.weight', 'roberta.pooler.dense.bias']\n",
      "- This IS expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reward model output:\n",
      "For non-toxic text\n",
      "[{'label': 'neutral', 'score': 4.369976997375488}, {'label': 'toxic', 'score': -4.438220500946045}]\n",
      "[{'label': 'neutral', 'score': 0.9998505115509033}, {'label': 'toxic', 'score': 0.0001494801981607452}]\n",
      "For toxic text\n",
      "[{'label': 'toxic', 'score': 3.440110921859741}, {'label': 'neutral', 'score': -3.532526731491089}]\n",
      "[{'label': 'toxic', 'score': 0.9990636706352234}, {'label': 'neutral', 'score': 0.0009363002027384937}]\n"
     ]
    }
   ],
   "source": [
    "device = 0 if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "sentiment_pipe = pipeline(\"sentiment-analysis\", \n",
    "                          model=reward_model_name, \n",
    "                          device=device)\n",
    "reward_logits_kwargs = {\n",
    "    \"top_k\": None, #return all scores, not just for top k candidates\n",
    "    \"function_to_apply\": \"none\", # use raw logits instead of applying function\n",
    "    \"batch_size\": 16\n",
    "}\n",
    "\n",
    "reward_probabilities_kwargs = {\n",
    "    \"top_k\": None, \n",
    "    \"function_to_apply\": \"softmax\", #\"softmax\" function converts logit values to probabilities\n",
    "    \"batch_size\": 16\n",
    "}\n",
    "\n",
    "print(\"Reward model output:\")\n",
    "print(\"For non-toxic text\")\n",
    "print(sentiment_pipe(nontoxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(nontoxic_text, **reward_probabilities_kwargs))\n",
    "\n",
    "print(\"For toxic text\")\n",
    "print(sentiment_pipe(toxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(toxic_text, **reward_probabilities_kwargs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7.2'></a>\n",
    "### 7.2 - Calculating baseline toxicity stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get a sense of how toxic the output of our original model is. \n",
    "\n",
    "Before we actually crunch the numbers, we should set up an evaluation metric for toxicity as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at s-nlp/roberta_toxicity_classifier were not used when initializing RobertaForSequenceClassification: ['roberta.pooler.dense.weight', 'roberta.pooler.dense.bias']\n",
      "- This IS expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "toxic_evaluator = evaluate.load(\"toxicity\", \n",
    "                                    reward_model_name,\n",
    "                                    module_type=\"measurement\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Toxicity score for following sentence: The Barbie movie exceeded my expectations.\n",
      "[0.0001494801981607452]\n",
      "\n",
      "Toxicity score for following sentence: The Barbie movie was a disgusting piece of fake-feminist materialistic garbage.\n",
      "[0.9990636706352234]\n"
     ]
    }
   ],
   "source": [
    "toxicity_score = toxic_evaluator.compute(predictions=[\n",
    "    nontoxic_text\n",
    "], toxic_label=\"toxic\")\n",
    "\n",
    "print(\"Toxicity score for following sentence: \"+ nontoxic_text)\n",
    "print(toxicity_score[\"toxicity\"])\n",
    "\n",
    "toxicity_score = toxic_evaluator.compute(predictions=[\n",
    "    toxic_text\n",
    "], toxic_label=\"toxic\")\n",
    "\n",
    "print(\"\\nToxicity score for following sentence: \"+toxic_text)\n",
    "print(toxicity_score[\"toxicity\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks good?\n",
    "\n",
    "For the next steps it'll be helpful to have the original English text (not just the number-converted encoded input_ids) in our dataset.\n",
    "\n",
    "We can retrieve this by decoding the input_ids as follows -- the code is a little ugly but it gets the job done:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "The table can't have duplicated columns but columns ['query'] are duplicated.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_24736\\3089472747.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      4\u001b[0m      \u001b[0mdecoded_text_vector\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtokenizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mskip_special_tokens\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 6\u001b[1;33m \u001b[0msubsample_tokenized_datasets\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'train'\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0msubsample_tokenized_datasets\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'train'\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd_column\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"query\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mdecoded_text_vector\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      7\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m \u001b[0mdecoded_text_vector\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python39\\site-packages\\datasets\\arrow_dataset.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    526\u001b[0m         }\n\u001b[0;32m    527\u001b[0m         \u001b[1;31m# apply actual function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 528\u001b[1;33m         \u001b[0mout\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m\"Dataset\"\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"DatasetDict\"\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    529\u001b[0m         \u001b[0mdatasets\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mList\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m\"Dataset\"\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdict\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32melse\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    530\u001b[0m         \u001b[1;31m# re-apply format to the output\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python39\\site-packages\\datasets\\fingerprint.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    509\u001b[0m             \u001b[1;31m# Call actual function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    510\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 511\u001b[1;33m             \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    512\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    513\u001b[0m             \u001b[1;31m# Update fingerprint of in-place transforms + update in-place history of transforms\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python39\\site-packages\\datasets\\arrow_dataset.py\u001b[0m in \u001b[0;36madd_column\u001b[1;34m(self, name, column, new_fingerprint)\u001b[0m\n\u001b[0;32m   5443\u001b[0m         \"\"\"\n\u001b[0;32m   5444\u001b[0m         \u001b[0mcolumn_table\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mInMemoryTable\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfrom_pydict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m{\u001b[0m\u001b[0mname\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mcolumn\u001b[0m\u001b[1;33m}\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 5445\u001b[1;33m         \u001b[0m_check_column_names\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_data\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcolumn_names\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mcolumn_table\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcolumn_names\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   5446\u001b[0m         \u001b[0mdataset\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mflatten_indices\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_indices\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[1;32mNone\u001b[0m \u001b[1;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   5447\u001b[0m         \u001b[1;31m# Concatenate tables horizontally\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Roaming\\Python\\Python39\\site-packages\\datasets\\arrow_dataset.py\u001b[0m in \u001b[0;36m_check_column_names\u001b[1;34m(column_names)\u001b[0m\n\u001b[0;32m    613\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mall\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcount\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;36m1\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mcount\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mcounter\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    614\u001b[0m         \u001b[0mduplicated_columns\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mcol\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mcol\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mcounter\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mcounter\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mcol\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m>\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 615\u001b[1;33m         \u001b[1;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34mf\"The table can't have duplicated columns but columns {duplicated_columns} are duplicated.\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    616\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    617\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mValueError\u001b[0m: The table can't have duplicated columns but columns ['query'] are duplicated."
     ]
    }
   ],
   "source": [
    "decoded_text_vector = []\n",
    "\n",
    "for i in subsample_tokenized_datasets['train']['input_ids']:\n",
    "     decoded_text_vector.append(tokenizer.decode(i,skip_special_tokens=True))\n",
    "    \n",
    "subsample_tokenized_datasets['train'] = subsample_tokenized_datasets['train'].add_column(\"query\",decoded_text_vector)\n",
    "\n",
    "decoded_text_vector = []\n",
    "\n",
    "for i in subsample_tokenized_datasets['test']['input_ids']:\n",
    "     decoded_text_vector.append(tokenizer.decode(i,skip_special_tokens=True))\n",
    "\n",
    "subsample_tokenized_datasets['test'] = subsample_tokenized_datasets['test'].add_column(\"query\",decoded_text_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to set the evaluator loose on our pre-finetuned model to get the mean and standard deviation of the toxicity of its dialogue summaries. The best approach here is to first define a function as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_toxicity(model, \n",
    "                      toxicity_evaluator, \n",
    "                      tokenizer, \n",
    "                      dataset, \n",
    "                      num_samples):\n",
    "\n",
    "    max_new_tokens=100\n",
    "\n",
    "    toxicities = []\n",
    "    input_texts = []\n",
    "    for i, sample in tqdm(enumerate(dataset)):\n",
    "        input_text = sample[\"query\"]\n",
    "\n",
    "        if i > num_samples:\n",
    "            break\n",
    "            \n",
    "        input_ids = tokenizer(input_text, return_tensors=\"pt\", padding=True).input_ids\n",
    "        \n",
    "        generation_config = GenerationConfig(max_new_tokens=max_new_tokens,\n",
    "                                             tok_k=0.0,\n",
    "                                             top_p=1.0,\n",
    "                                             do_sample=True)\n",
    "\n",
    "        response_token_ids = model.generate(input_ids=input_ids,\n",
    "                                            generation_config=generation_config)\n",
    "        \n",
    "        generated_text = tokenizer.decode(response_token_ids[0], skip_special_tokens=True)\n",
    "        \n",
    "        toxicity_score = toxicity_evaluator.compute(predictions=[(input_text + \" \" + generated_text)], toxic_label=\"toxic\")\n",
    "\n",
    "        toxicities.extend(toxicity_score[\"toxicity\"])\n",
    "\n",
    "    mean = np.mean(toxicities)\n",
    "    std = np.std(toxicities)\n",
    "        \n",
    "    return mean, std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One more step: as mentioned above implementing Reinforcement Learning with Human Feedback isn't just a matter of using reward values to adjust model weights based on the toxicity (or honesty, or helpfulness, etc.) of the model. To prevent \"reward hacking\" (e.g., the LLM stops trying to carry out the actual task and instead just talks about love, sunishine, and rainbows to minimize its toxicity score), we also need to weigh this against how similar the retrained model outputs are to what the original model gave (calculated using metrics like KL divergence). This balances how much learning occurs from the reward vs. the similarity to the original model.\n",
    "    \n",
    "We create a copy of the model that is trainable (so that its values can be iteratively updated with each toxicity assessment) as well as a \"frozen\" reference model that remains unchanged during toxicity evaluation. The create_reference_model() function makes a copy whose parameters are not trainable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [],
   "source": [
    "ppo_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(peft_model,                                                               \n",
    "                                                               torch_dtype=torch.bfloat16,\n",
    "                                                               is_trainable=True)\n",
    "\n",
    "ref_model = create_reference_model(ppo_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now use the function we just defined to get the mean and SD of the toxicity of our pre-training model outputs, to give us a baseline comparison:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11it [00:34,  3.12s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean toxicity and standard deviation before detox: [3.804054689350199e-05, 6.522331639452418e-06]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, device_map=\"auto\")\n",
    "\n",
    "mean_toxicity_original, std_toxicity_original = evaluate_toxicity(model=ref_model, \n",
    "                                                                          toxicity_evaluator=toxic_evaluator, \n",
    "                                                                          tokenizer=tokenizer, \n",
    "                                                                          dataset=subsample_tokenized_datasets['test'], \n",
    "                                                                          num_samples=10)\n",
    "\n",
    "print(f'Mean toxicity and standard deviation before detox: [{mean_toxicity_original}, {std_toxicity_original}]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These numbers are actually extremely low (perhaps not surprising since the Dialogsum data are just rather banal conversations, not online social media content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7.3'></a>\n",
    "### 7.3 - Fine-tuning our model to reduce toxicity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to actually fine-tune the model, using a combination of reward values and similarity metrics from the original, \"frozen\" model.\n",
    "\n",
    "The next step requires a collator function to reformat our data slightly, taking a dictionary out of a one-element list and giving entries within the dictionary as lists of length 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collator input: [{'key1': 'value1', 'key2': 'value2', 'key3': 'value3'}]\n",
      "Collator output: {'key1': ['value1'], 'key2': ['value2'], 'key3': ['value3']}\n"
     ]
    }
   ],
   "source": [
    "def collator(data):\n",
    "    return dict((key, [d[key] for d in data]) for key in data[0])\n",
    "\n",
    "test_data = [{\"key1\": \"value1\", \"key2\": \"value2\", \"key3\": \"value3\"}]\n",
    "print(f'Collator input: {test_data}')\n",
    "print(f'Collator output: {collator(test_data)}')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to define some parameters for our training function. Note that we are still using our mini version of the data (100th the size of the full Dialogsum dataset):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=1.41e-5\n",
    "max_ppo_epochs=1\n",
    "mini_batch_size=4\n",
    "batch_size=16\n",
    "\n",
    "config = PPOConfig(\n",
    "    model_name=model_name,    \n",
    "    learning_rate=learning_rate,\n",
    "    ppo_epochs=max_ppo_epochs,\n",
    "    mini_batch_size=mini_batch_size,\n",
    "    batch_size=batch_size\n",
    ")\n",
    "\n",
    "ppo_trainer = PPOTrainer(config=config, \n",
    "                         model=ppo_model, \n",
    "                         ref_model=ref_model, \n",
    "                         tokenizer=tokenizer, \n",
    "                         dataset=subsample_tokenized_datasets[\"train\"], \n",
    "                         data_collator=collator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the fine-tuning itself. Note that LengthSampler() uniformly samples values from a defined range so that the query and response lengths are randomized. I set it to 10 PPO steps but this can be adjusted based on your compute availability. This took about 10 minutes to run on a computer with 11th Gen Intel i5-11300H 3.10GHz CPU with 32GB RAM and an Nvidia Geforce RTX 3050 GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]You're using a T5TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "7it [16:44, 143.48s/it]\n"
     ]
    }
   ],
   "source": [
    "output_min_length = 100\n",
    "output_max_length = 400\n",
    "output_length_sampler = LengthSampler(output_min_length, output_max_length)\n",
    "\n",
    "generation_kwargs = {\n",
    "    \"min_length\": 5,\n",
    "    \"top_k\": 0.0,\n",
    "    \"top_p\": 1.0,\n",
    "    \"do_sample\": True\n",
    "}\n",
    "\n",
    "reward_kwargs = {\n",
    "    \"top_k\": None, # Return all scores.\n",
    "    \"function_to_apply\": \"none\", # You want the raw logits without softmax.\n",
    "    \"batch_size\": 16\n",
    "}\n",
    "\n",
    "max_ppo_steps = 10\n",
    "\n",
    "for step, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n",
    "    if step >= max_ppo_steps:\n",
    "        break   \n",
    "\n",
    "    prompt_tensors = [torch.tensor(x) for x in batch[\"input_ids\"]] #convert from integers to tensor format\n",
    "\n",
    "    summary_tensors = []\n",
    "\n",
    "    for prompt_tensor in prompt_tensors:\n",
    "        max_new_tokens = output_length_sampler()        \n",
    "            \n",
    "        generation_kwargs[\"max_new_tokens\"] = max_new_tokens\n",
    "        summary = ppo_trainer.generate(prompt_tensor, **generation_kwargs)\n",
    "        \n",
    "        summary_tensors.append(summary.squeeze()[-max_new_tokens:])\n",
    "        \n",
    "    batch[\"response\"] = [tokenizer.decode(r.squeeze()) for r in summary_tensors]\n",
    "\n",
    "    query_response_pairs = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]    \n",
    "    rewards = sentiment_pipe(query_response_pairs, **reward_kwargs)\n",
    "\n",
    "    reward_tensors = [torch.tensor(reward[nottoxic_index][\"score\"]) for reward in rewards]    \n",
    "\n",
    "    stats = ppo_trainer.step(prompt_tensors, summary_tensors, reward_tensors)\n",
    "    ppo_trainer.log_stats(stats, batch, reward_tensors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7.4'></a>\n",
    "### 7.4 - Quantifying the toxicity of our detoxified model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the same procedure as before to crunch some toxicity stats for our newly-detoxified model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11it [00:33,  3.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean toxicity and standard deviation after detox: [3.8285791643747046e-05, 5.815389222470542e-06]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, device_map=\"auto\")\n",
    "\n",
    "mean_toxicity_detoxmodel, std_toxicity_detoxmodel = evaluate_toxicity(model=ppo_model, \n",
    "                                                                          toxicity_evaluator=toxic_evaluator, \n",
    "                                                                          tokenizer=tokenizer, \n",
    "                                                                          dataset=subsample_tokenized_datasets['test'], \n",
    "                                                                          num_samples=10)\n",
    "\n",
    "print(f'Mean toxicity and standard deviation after detox: [{mean_toxicity_detoxmodel}, {std_toxicity_detoxmodel}]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The numbers were quite small to begin with, but for reference here is some code to compare the pre/post detox stats in terms of percentage change:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage improvement of toxicity score after detoxification:\n",
      "mean: -0.64%\n",
      "std: 10.84%\n"
     ]
    }
   ],
   "source": [
    "mean_toxicity_change = (mean_toxicity_original - mean_toxicity_detoxmodel) / mean_toxicity_original\n",
    "std_toxicity_change = (std_toxicity_original - std_toxicity_detoxmodel) / std_toxicity_original\n",
    "\n",
    "print(f'Percentage improvement of toxicity score after detoxification:')\n",
    "print(f'mean: {mean_toxicity_change*100:.2f}%')\n",
    "print(f'std: {std_toxicity_change*100:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The difference isn't very significant (perhaps due to the small toxicity scores in the first place?)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name='7.5'></a>\n",
    "### 7.5 - Visually comparing pre-detoxified vs. detoxified models\n",
    "\n",
    "For visualization purposes, let's take just the first ten entries in our test data and compare the summaries and reward scores for the pre- and post-detoxification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 10/10 [00:58<00:00,  5.81s/it]\n"
     ]
    }
   ],
   "source": [
    "batch_size = 10\n",
    "compare_results = {}\n",
    "\n",
    "df_batch = subsample_tokenized_datasets[\"test\"][0:batch_size]\n",
    "\n",
    "compare_results[\"query\"] = df_batch[\"query\"]\n",
    "prompt_tensors = df_batch[\"input_ids\"]\n",
    "\n",
    "summary_tensors_ref = []\n",
    "summary_tensors = []\n",
    "\n",
    "for i in tqdm(range(batch_size)):\n",
    "    gen_len = output_length_sampler()\n",
    "    generation_kwargs[\"max_new_tokens\"] = gen_len\n",
    "    \n",
    "    summary = ref_model.generate(\n",
    "        input_ids=torch.as_tensor(prompt_tensors[i]).unsqueeze(dim=0).to(device), \n",
    "        **generation_kwargs\n",
    "    ).squeeze()[-gen_len:]\n",
    "    summary_tensors_ref.append(summary)\n",
    "\n",
    "    summary = ppo_model.generate(\n",
    "        input_ids=torch.as_tensor(prompt_tensors[i]).unsqueeze(dim=0).to(device), \n",
    "        **generation_kwargs\n",
    "    ).squeeze()[-gen_len:]\n",
    "    summary_tensors.append(summary)\n",
    "\n",
    "#Translate from encoded numerical tensors to human-legible tokens\n",
    "compare_results[\"response_before\"] = [tokenizer.decode(summary_tensors_ref[i]) for i in range(batch_size)]\n",
    "compare_results[\"response_after\"] = [tokenizer.decode(summary_tensors[i]) for i in range(batch_size)]\n",
    "\n",
    "#Compare the pre/post detox values\n",
    "texts_before = [d + s for d, s in zip(compare_results[\"query\"], compare_results[\"response_before\"])]\n",
    "rewards_before = sentiment_pipe(texts_before, **reward_kwargs)\n",
    "compare_results[\"reward_before\"] = [reward[nottoxic_index][\"score\"] for reward in rewards_before]\n",
    "\n",
    "texts_after = [d + s for d, s in zip(compare_results[\"query\"], compare_results[\"response_after\"])]\n",
    "rewards_after = sentiment_pipe(texts_after, **reward_kwargs)\n",
    "compare_results[\"reward_after\"] = [reward[nottoxic_index][\"score\"] for reward in rewards_after]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a Pandas dataframe to view these. We can sort by \" reward difference\" to see where the difference in toxicity was largest betweeh the pre-detox vs. detoxified models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>response_before</th>\n",
       "      <th>response_after</th>\n",
       "      <th>reward_before</th>\n",
       "      <th>reward_after</th>\n",
       "      <th>reward_diff</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: I cannot imagine if Trump were to be our President again. #Person2#: I am proud to say that he is our President, and I will be really happy if he could be re-elected. #Person1#: You voted for him, right? #Person2#: Did you vote for him, because I know that I did. #Person1#: I am not sure about this. #Person2#: I have nothing but faith in Trump. #Person1#: What? #Person2#: I am pretty sure he will make America great again! #Person1#: Wel...</td>\n",
       "      <td>&lt;pad&gt; The President of the United States needs to succeed because of his lack of experience and mendable nature.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; People are excited about Trump and try to imagine his war in the future.&lt;/s&gt;</td>\n",
       "      <td>4.625224</td>\n",
       "      <td>4.986396</td>\n",
       "      <td>0.361172</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: Have you considered upgrading your system? #Person2#: Yes, but I'm not sure what exactly I would need. #Person1#: You could consider adding a painting program to your software. It would allow you to make up your own flyers and banners for advertising. #Person2#: That would be a definite bonus. #Person1#: You might also want to upgrade your hardware because it is pretty outdated now. #Person2#: How can we do that? #Person1#: You'd probab...</td>\n",
       "      <td>&lt;pad&gt; A painting program is great for advertising. A CD-ROM drive is also included.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; #Person2#: You're just clicking on your computer. Or join in on one of the Facebook groups. #Person1#: You have a few options to make your computer computer more complicated. #Person2#: Is there anything I can do? #Person1#: To learn more, please subscribe here. Where I Get started: http://www.facebook.com/people/site/21251293145.html&lt;/s&gt;</td>\n",
       "      <td>4.959960</td>\n",
       "      <td>4.996674</td>\n",
       "      <td>0.036714</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: We need to call an emergency meeting as soon as possible about this. #Person2#: OK. Let me send a memo around. #Person1#: Schedule it for this afternoon, when Ken gets back. #Person2#: I don't think he's back today. #Person1#: Oh, that's right. Go ahead anyways. I'll fill him in. We can't lose this order under any circumstances! #Person2#: I know, it's a big one. Summary:</td>\n",
       "      <td>&lt;pad&gt; The emergency meeting will be held in an emergency. Ken might not be here today.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; Parents are wondering if Ken is back.&lt;/s&gt;</td>\n",
       "      <td>5.006478</td>\n",
       "      <td>5.016961</td>\n",
       "      <td>0.010483</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: Ms. Dawson, I need you to take a dictation for me. #Person2#: Yes, sir... #Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready? #Person2#: Yes, sir. Go ahead. #Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly ...</td>\n",
       "      <td>&lt;pad&gt; The Office of the President, Vice President and Director of Communications. NYPD and Department Heads: Effective immediately, all office communications are restricted to email correspondence and official memos.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; #Person1#: We are converting an intra-office memo on the information regarding employee access.&lt;/s&gt;</td>\n",
       "      <td>4.971152</td>\n",
       "      <td>4.981174</td>\n",
       "      <td>0.010022</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: It was a heavy storm last night, wasn't it? #Person2#: It certainly was. The wind broke several windows. What weather! #Person1#: Do you know that big tree in front of my house? One of the biggest branches came down in the night. #Person2#: Really? Did it do any damage to your home? #Person1#: Thank goodness! It is far away from that. #Person2#: I really hate storms. It's about time we had some nice spring weather. #Person1#: It's April...</td>\n",
       "      <td>&lt;pad&gt; The installation of the rain system last night broke several windows.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; There was a heavy storm with a presidential and a slightly angry outburst of tremors.&lt;/s&gt;</td>\n",
       "      <td>4.778887</td>\n",
       "      <td>4.777029</td>\n",
       "      <td>-0.001858</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: Are you going anywhere for your vacation? #Person2#: Yes, we're making plans for a tour. #Person1#: That'll be lovely. Where are you going? #Person2#: Well, we will start out from Long Island this Friday. We've planned a four day drive to Salt Lake City, where we'll join my brother and his family on his fortieth birthday. #Person1#: Well, you've got to prepare a lot of food and enough sleeping bags then. #Person2#: Oh, we'll spend the n...</td>\n",
       "      <td>&lt;pad&gt; David is planning a four day tour to Salt Lake City. To join his brother and his family, they plan to take a four day drive this week, and lay out food and sleeping bags. It's being promoted as a Christmas bet.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; David and his brother are going for a four day tour to Salt Lake City. They plan on staying at a house most likely in Long Island, as they don't have a car when doing the drive.&lt;/s&gt;</td>\n",
       "      <td>5.019026</td>\n",
       "      <td>5.013295</td>\n",
       "      <td>-0.005731</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: Dad, you keep talking about family in New Zealand. Who are they? #Person2#: Well, that's your uncle Bill, his wife and two of their daughters. #Person1#: Is uncle Bill your brother? #Person2#: No, your uncle Jack is my brother, Bill is my brother-in-law, your mom's brother. #Person1#: So his two daughters are my cousins? #Person2#: That's right, Sarah and Jane are both your cousins although they are step-sisters. #Person1#: What are ste...</td>\n",
       "      <td>&lt;pad&gt; The family in New Zealand is said to have a lot in common.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; Parents of children called visiting their relatives in New Zealand. People who weren't interesting called the new British stop and Spanish-language visitors.&lt;/s&gt;</td>\n",
       "      <td>5.001910</td>\n",
       "      <td>4.985108</td>\n",
       "      <td>-0.016802</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: Oh, I'm starving. It's my first time to China. And I'd like to try some real Chinese cuisine. What would you recommend? #Person2#: Well, depends. You see, there are eight famous Chinese food cuisines, for instance, Sichuan cuisine and Hunan cuisine. #Person1#: There're all spicy or hot of heard. #Person2#: That's right. If you have hot dishes, you can try some. #Person1#: I cannot have it. Last time I had some in the US. It almost kille...</td>\n",
       "      <td>&lt;pad&gt; We were unable to find a place for 'Pennhaburi' in China, because there is a lot of food. I'm starving, so I'll try to find 'Cantonese' restaurant.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; TellPeople1 which Chinese restaurant is best for them. Go to the nearest restaurant, and wait for a taxi right away.&lt;/s&gt;</td>\n",
       "      <td>4.802761</td>\n",
       "      <td>4.780698</td>\n",
       "      <td>-0.022063</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: What do you think of my new suit? #Person2#: Not bad. It reminds me of the one I saw at the new department store last week. Did you get it there? #Person1#: No, I got it in that big shopping center. It cost me only $ 150. #Person2#: Well, I don't think it's a good bargain. Summary:</td>\n",
       "      <td>&lt;pad&gt; They received the suit in a big shopping center.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; Person1 and #Person2 are both complaining about their new suits.&lt;/s&gt;</td>\n",
       "      <td>4.976263</td>\n",
       "      <td>4.943659</td>\n",
       "      <td>-0.032604</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Produce a summary of the following dialogue. #Person1#: OK, that's a cut! Let's start from the beginning, everyone. #Person2#: What was the problem that time? #Person1#: The feeling was all wrong, Mike. She is telling you that she doesn't want to see you any more, but I want to get more anger from you. You're acting hurt and sad, but that's not how your character would act in this situation. #Person2#: But Jason and Laura have been together for three years. Don't you think his reaction would...</td>\n",
       "      <td>&lt;pad&gt; As the situation involves Jason and Laura, they tried using their reasoning to cast a light on Mike's feelings.&lt;/s&gt;</td>\n",
       "      <td>&lt;pad&gt; It was seems the person who messed up his sensing will react differently to it than the one that Mike messed up.&lt;/s&gt;</td>\n",
       "      <td>4.968035</td>\n",
       "      <td>4.848250</td>\n",
       "      <td>-0.119785</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 query  \\\n",
       "0  Produce a summary of the following dialogue. #Person1#: I cannot imagine if Trump were to be our President again. #Person2#: I am proud to say that he is our President, and I will be really happy if he could be re-elected. #Person1#: You voted for him, right? #Person2#: Did you vote for him, because I know that I did. #Person1#: I am not sure about this. #Person2#: I have nothing but faith in Trump. #Person1#: What? #Person2#: I am pretty sure he will make America great again! #Person1#: Wel...   \n",
       "1  Produce a summary of the following dialogue. #Person1#: Have you considered upgrading your system? #Person2#: Yes, but I'm not sure what exactly I would need. #Person1#: You could consider adding a painting program to your software. It would allow you to make up your own flyers and banners for advertising. #Person2#: That would be a definite bonus. #Person1#: You might also want to upgrade your hardware because it is pretty outdated now. #Person2#: How can we do that? #Person1#: You'd probab...   \n",
       "2                                                                      Produce a summary of the following dialogue. #Person1#: We need to call an emergency meeting as soon as possible about this. #Person2#: OK. Let me send a memo around. #Person1#: Schedule it for this afternoon, when Ken gets back. #Person2#: I don't think he's back today. #Person1#: Oh, that's right. Go ahead anyways. I'll fill him in. We can't lose this order under any circumstances! #Person2#: I know, it's a big one. Summary:    \n",
       "3  Produce a summary of the following dialogue. #Person1#: Ms. Dawson, I need you to take a dictation for me. #Person2#: Yes, sir... #Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready? #Person2#: Yes, sir. Go ahead. #Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly ...   \n",
       "4  Produce a summary of the following dialogue. #Person1#: It was a heavy storm last night, wasn't it? #Person2#: It certainly was. The wind broke several windows. What weather! #Person1#: Do you know that big tree in front of my house? One of the biggest branches came down in the night. #Person2#: Really? Did it do any damage to your home? #Person1#: Thank goodness! It is far away from that. #Person2#: I really hate storms. It's about time we had some nice spring weather. #Person1#: It's April...   \n",
       "5  Produce a summary of the following dialogue. #Person1#: Are you going anywhere for your vacation? #Person2#: Yes, we're making plans for a tour. #Person1#: That'll be lovely. Where are you going? #Person2#: Well, we will start out from Long Island this Friday. We've planned a four day drive to Salt Lake City, where we'll join my brother and his family on his fortieth birthday. #Person1#: Well, you've got to prepare a lot of food and enough sleeping bags then. #Person2#: Oh, we'll spend the n...   \n",
       "6  Produce a summary of the following dialogue. #Person1#: Dad, you keep talking about family in New Zealand. Who are they? #Person2#: Well, that's your uncle Bill, his wife and two of their daughters. #Person1#: Is uncle Bill your brother? #Person2#: No, your uncle Jack is my brother, Bill is my brother-in-law, your mom's brother. #Person1#: So his two daughters are my cousins? #Person2#: That's right, Sarah and Jane are both your cousins although they are step-sisters. #Person1#: What are ste...   \n",
       "7  Produce a summary of the following dialogue. #Person1#: Oh, I'm starving. It's my first time to China. And I'd like to try some real Chinese cuisine. What would you recommend? #Person2#: Well, depends. You see, there are eight famous Chinese food cuisines, for instance, Sichuan cuisine and Hunan cuisine. #Person1#: There're all spicy or hot of heard. #Person2#: That's right. If you have hot dishes, you can try some. #Person1#: I cannot have it. Last time I had some in the US. It almost kille...   \n",
       "8                                                                                                                                                                  Produce a summary of the following dialogue. #Person1#: What do you think of my new suit? #Person2#: Not bad. It reminds me of the one I saw at the new department store last week. Did you get it there? #Person1#: No, I got it in that big shopping center. It cost me only $ 150. #Person2#: Well, I don't think it's a good bargain. Summary:    \n",
       "9  Produce a summary of the following dialogue. #Person1#: OK, that's a cut! Let's start from the beginning, everyone. #Person2#: What was the problem that time? #Person1#: The feeling was all wrong, Mike. She is telling you that she doesn't want to see you any more, but I want to get more anger from you. You're acting hurt and sad, but that's not how your character would act in this situation. #Person2#: But Jason and Laura have been together for three years. Don't you think his reaction would...   \n",
       "\n",
       "                                                                                                                                                                                                                response_before  \\\n",
       "0                                                                                                          <pad> The President of the United States needs to succeed because of his lack of experience and mendable nature.</s>   \n",
       "1                                                                                                                                       <pad> A painting program is great for advertising. A CD-ROM drive is also included.</s>   \n",
       "2                                                                                                                                    <pad> The emergency meeting will be held in an emergency. Ken might not be here today.</s>   \n",
       "3  <pad> The Office of the President, Vice President and Director of Communications. NYPD and Department Heads: Effective immediately, all office communications are restricted to email correspondence and official memos.</s>   \n",
       "4                                                                                                                                               <pad> The installation of the rain system last night broke several windows.</s>   \n",
       "5  <pad> David is planning a four day tour to Salt Lake City. To join his brother and his family, they plan to take a four day drive this week, and lay out food and sleeping bags. It's being promoted as a Christmas bet.</s>   \n",
       "6                                                                                                                                                          <pad> The family in New Zealand is said to have a lot in common.</s>   \n",
       "7                                                                 <pad> We were unable to find a place for 'Pennhaburi' in China, because there is a lot of food. I'm starving, so I'll try to find 'Cantonese' restaurant.</s>   \n",
       "8                                                                                                                                                                    <pad> They received the suit in a big shopping center.</s>   \n",
       "9                                                                                                     <pad> As the situation involves Jason and Laura, they tried using their reasoning to cast a light on Mike's feelings.</s>   \n",
       "\n",
       "                                                                                                                                                                                                                                                                                                                                               response_after  \\\n",
       "0                                                                                                                                                                                                                                                                          <pad> People are excited about Trump and try to imagine his war in the future.</s>   \n",
       "1  <pad> #Person2#: You're just clicking on your computer. Or join in on one of the Facebook groups. #Person1#: You have a few options to make your computer computer more complicated. #Person2#: Is there anything I can do? #Person1#: To learn more, please subscribe here. Where I Get started: http://www.facebook.com/people/site/21251293145.html</s>   \n",
       "2                                                                                                                                                                                                                                                                                                             <pad> Parents are wondering if Ken is back.</s>   \n",
       "3                                                                                                                                                                                                                                                   <pad> #Person1#: We are converting an intra-office memo on the information regarding employee access.</s>   \n",
       "4                                                                                                                                                                                                                                                             <pad> There was a heavy storm with a presidential and a slightly angry outburst of tremors.</s>   \n",
       "5                                                                                                                                                                 <pad> David and his brother are going for a four day tour to Salt Lake City. They plan on staying at a house most likely in Long Island, as they don't have a car when doing the drive.</s>   \n",
       "6                                                                                                                                                                                     <pad> Parents of children called visiting their relatives in New Zealand. People who weren't interesting called the new British stop and Spanish-language visitors.</s>   \n",
       "7                                                                                                                                                                                                                              <pad> TellPeople1 which Chinese restaurant is best for them. Go to the nearest restaurant, and wait for a taxi right away.</s>   \n",
       "8                                                                                                                                                                                                                                                                                  <pad> Person1 and #Person2 are both complaining about their new suits.</s>   \n",
       "9                                                                                                                                                                                                                                  <pad> It was seems the person who messed up his sensing will react differently to it than the one that Mike messed up.</s>   \n",
       "\n",
       "   reward_before  reward_after  reward_diff  \n",
       "0       4.625224      4.986396     0.361172  \n",
       "1       4.959960      4.996674     0.036714  \n",
       "2       5.006478      5.016961     0.010483  \n",
       "3       4.971152      4.981174     0.010022  \n",
       "4       4.778887      4.777029    -0.001858  \n",
       "5       5.019026      5.013295    -0.005731  \n",
       "6       5.001910      4.985108    -0.016802  \n",
       "7       4.802761      4.780698    -0.022063  \n",
       "8       4.976263      4.943659    -0.032604  \n",
       "9       4.968035      4.848250    -0.119785  "
      ]
     },
     "execution_count": 215,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.set_option('display.max_colwidth', 500)\n",
    "df_compare_results = pd.DataFrame(compare_results)\n",
    "df_compare_results[\"reward_diff\"] = df_compare_results['reward_after'] - df_compare_results['reward_before']\n",
    "df_compare_results_sorted = df_compare_results.sort_values(by=['reward_diff'], ascending=False).reset_index(drop=True)\n",
    "df_compare_results_sorted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The summaries above are a little subpar. Bear in mind that a. for this exercise I used a dataset only 100th the size of the original Dialogsum data; b. the original data wasn't that toxic in the first place, so the detoxification wouldn't lead to a massive change, and c. there are a ton of training parameters that I didn't fiddle with here: length of input, number of training epochs, learning rate, etc.\n",
    "\n",
    "In any case, I hope this exercise serves as a nice mini-tutorial for how a text-summarizing model can be improved with in-context learning, instruction finetuning, and reinforcement learning with human feedback!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "instance_type": "ml.m5.2xlarge",
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}