Fine-Tune FLAN-T5 with Reinforcement Learning (PPO) and PEFT .ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7f10b19-c061-42b2-8b42-e3cbafa3b1da",
   "metadata": {},
   "source": [
    "# Fine-Tune FLAN-T5 with Reinforcement Learning (PPO) and PEFT to Generate Less-Toxic Summaries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36ef668a-9c51-489b-be47-a07a09ef2289",
   "metadata": {},
   "source": [
    "In this notebook, you will fine-tune a FLAN-T5 model to generate less toxic content with Meta AI's hate speech reward model. The reward model is a binary classifier that predicts either \"not hate\" or \"hate\" for the given text. You will use Proximal Policy Optimization (PPO) to fine-tune and reduce the model's toxicity."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed5003e2-a642-416b-bd3f-93fb339c3a7d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Table of Contents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6791f449-d1da-461b-9eb7-c6dc6b37b15c",
   "metadata": {
    "tags": []
   },
   "source": [
    "- [ 1 - Set up Kernel and Required Dependencies](#1)\n",
    "- [ 2 - Load FLAN-T5 Model, Prepare Reward Model and Toxicity Evaluator](#2)\n",
    "  - [ 2.1 - Load Data and FLAN-T5 Model Fine-Tuned with Summarization Instruction](#2.1)\n",
    "  - [ 2.2 - Prepare Reward Model](#2.2)\n",
    "  - [ 2.3 - Evaluate Toxicity](#2.3)\n",
    "- [ 3 - Perform Fine-Tuning to Detoxify the Summaries](#3)\n",
    "  - [ 3.1 - Initialize `PPOTrainer`](#3.1)\n",
    "  - [ 3.2 - Fine-Tune the Model](#3.2)\n",
    "  - [ 3.3 - Evaluate the Model Quantitatively](#3.3)\n",
    "  - [ 3.4 - Evaluate the Model Qualitatively](#3.4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89f973f1-f095-4915-86d0-bc16380da22d",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a name='1'></a>\n",
    "## 1 - Set up Kernel and Required Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d24e86-f76f-4a44-90ef-0777752075a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "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",
    "    peft==0.3.0 --quiet\n",
    "\n",
    "# Installing the Reinforcement Learning library directly from github.\n",
    "%pip install git+https://github.com/lvwerra/trl.git@25fa1bd    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74bfd06e-c747-43e0-b86c-0398628e1c32",
   "metadata": {
    "tags": []
   },
   "source": [
    "Import the necessary components. Some of them are new for this week, they will be discussed later in the notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c20bed-6a30-4847-a507-02969ecb4465",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM, GenerationConfig\n",
    "from datasets import load_dataset\n",
    "from peft import PeftModel, PeftConfig, LoraConfig, TaskType\n",
    "\n",
    "# trl: Transformer Reinforcement Learning library\n",
    "from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead\n",
    "from trl import create_reference_model\n",
    "from trl.core import LengthSampler\n",
    "\n",
    "import torch\n",
    "import evaluate\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# tqdm library makes the loops show a smart progress meter.\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b76eea84-8e3a-4487-9692-613977e6c8e3",
   "metadata": {},
   "source": [
    "<a name='2'></a>\n",
    "## 2 - Load FLAN-T5 Model, Prepare Reward Model and Toxicity Evaluator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a5f97d4-ea5f-4072-b5d6-785d1d833ed4",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a name='2.1'></a>\n",
    "### 2.1 - Load Data and FLAN-T5 Model Fine-Tuned with Summarization Instruction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90dc0211-4032-4967-946d-3a538829d5c9",
   "metadata": {
    "tags": []
   },
   "source": [
    "You will keep working with the same Hugging Face dataset [DialogSum](https://huggingface.co/datasets/knkarthick/dialogsum) and the pre-trained model [FLAN-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b058b52b-ec4d-4426-8d71-91e898f727f6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name=\"google/flan-t5-base\"\n",
    "huggingface_dataset_name = \"knkarthick/dialogsum\"\n",
    "\n",
    "dataset_original = load_dataset(huggingface_dataset_name)\n",
    "\n",
    "dataset_original"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "668d30d6-6f81-4e52-a81a-3057163ddb0e",
   "metadata": {},
   "source": [
    "The next step will be to preprocess the dataset. You will take only a part of it, then filter the dialogues of a particular length (just to make those examples long enough and, at the same time, easy to read). Then wrap each dialogue with the instruction and tokenize the prompts. Save the token ids in the field `input_ids` and decoded version of the prompts in the field `query`.\n",
    "\n",
    "You could do that all step by step in the cell below, but it is a good habit to organize that all in a function `build_dataset`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51469abe-4d72-4093-a6c6-8e04e19f09eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(model_name,\n",
    "                  dataset_name,\n",
    "                  input_min_text_length, \n",
    "                  input_max_text_length):\n",
    "\n",
    "    \"\"\"\n",
    "    Preprocess the dataset and split it into train and test parts.\n",
    "\n",
    "    Parameters:\n",
    "    - model_name (str): Tokenizer model name.\n",
    "    - dataset_name (str): Name of the dataset to load.\n",
    "    - input_min_text_length (int): Minimum length of the dialogues.\n",
    "    - input_max_text_length (int): Maximum length of the dialogues.\n",
    "        \n",
    "    Returns:\n",
    "    - dataset_splits (datasets.dataset_dict.DatasetDict): Preprocessed dataset containing train and test parts.\n",
    "    \"\"\"\n",
    "    \n",
    "    # load dataset (only \"train\" part will be enough for this lab).\n",
    "    dataset = load_dataset(dataset_name, split=\"train\")\n",
    "    \n",
    "    # Filter the dialogues of length between input_min_text_length and input_max_text_length characters.\n",
    "    dataset = dataset.filter(lambda x: len(x[\"dialogue\"]) > input_min_text_length and len(x[\"dialogue\"]) <= input_max_text_length, batched=False)\n",
    "\n",
    "    # Prepare tokenizer. Setting device_map=\"auto\" allows to switch between GPU and CPU automatically.\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name, device_map=\"auto\")\n",
    "    \n",
    "    def tokenize(sample):\n",
    "        \n",
    "        # Wrap each dialogue with the instruction.\n",
    "        prompt = f\"\"\"\n",
    "Summarize the following conversation.\n",
    "\n",
    "{sample[\"dialogue\"]}\n",
    "\n",
    "Summary:\n",
    "\"\"\"\n",
    "        sample[\"input_ids\"] = tokenizer.encode(prompt)\n",
    "        \n",
    "        # This must be called \"query\", which is a requirement of our PPO library.\n",
    "        sample[\"query\"] = tokenizer.decode(sample[\"input_ids\"])\n",
    "        return sample\n",
    "\n",
    "    # Tokenize each dialogue.\n",
    "    dataset = dataset.map(tokenize, batched=False)\n",
    "    dataset.set_format(type=\"torch\")\n",
    "    \n",
    "    # Split the dataset into train and test parts.\n",
    "    dataset_splits = dataset.train_test_split(test_size=0.2, shuffle=False, seed=42)\n",
    "\n",
    "    return dataset_splits\n",
    "\n",
    "dataset = build_dataset(model_name=model_name,\n",
    "                        dataset_name=huggingface_dataset_name,\n",
    "                        input_min_text_length=200, \n",
    "                        input_max_text_length=1000)\n",
    "\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d03155e-649b-45bb-a5a0-94edd682c069",
   "metadata": {
    "tags": []
   },
   "source": [
    "In the previous lab, you fine-tuned the PEFT model with summarization instructions. The training in the notebook was done on a subset of data. Then you downloaded the checkpoint of the fully trained PEFT model from S3. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4226923-67c0-4ea6-8e47-030136b2f191",
   "metadata": {},
   "source": [
    "Prepare a function to pull out the number of model parameters (it is the same as in the previous lab):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f06806-a194-4c14-b64d-e31afd7b658c",
   "metadata": {},
   "outputs": [],
   "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\"\\ntrainable model parameters: {trainable_model_params}\\nall model parameters: {all_model_params}\\npercentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21e06a57-fb80-4f8c-a967-4c7c42a7bfda",
   "metadata": {
    "tags": []
   },
   "source": [
    "Add the adapter to the original FLAN-T5 model. In the previous lab you were adding the fully trained adapter only for inferences, so there was no need to pass LoRA configurations doing that. Now you need to pass them to the constructed PEFT model, also putting `is_trainable=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a94b14-b375-45e7-9e49-a7f2c341b4ff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    r=32, # Rank\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 # FLAN-T5\n",
    ")\n",
    "\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name, \n",
    "                                              torch_dtype=torch.bfloat16)\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(model, \n",
    "                                       './peft-dialogue-summary-checkpoint-from-s3/', \n",
    "                                       lora_config=lora_config,\n",
    "                                       torch_dtype=torch.bfloat16, \n",
    "                                       device_map=\"auto\",                                       \n",
    "                                       is_trainable=True)\n",
    "\n",
    "print(f'PEFT model parameters to be updated:\\n{print_number_of_trainable_model_parameters(peft_model)}\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a950ae8a-76b9-4951-9c78-9ac7a6349e17",
   "metadata": {},
   "source": [
    "In this lab, you are preparing to fine-tune the LLM using Reinforcement Learning (RL). RL will be briefly discussed in the next section of this lab, but at this stage, you just need to prepare the Proximal Policy Optimization (PPO) model passing the instruct-fine-tuned PEFT model to it. PPO will be used to optimize the RL policy against the reward model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e86bab0-6dee-4dff-a754-b584ed962723",
   "metadata": {},
   "outputs": [],
   "source": [
    "ppo_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(peft_model,                                                               \n",
    "                                                               torch_dtype=torch.bfloat16,\n",
    "                                                               is_trainable=True)\n",
    "\n",
    "print(f'PPO model parameters to be updated (ValueHead + 769 params):\\n{print_number_of_trainable_model_parameters(ppo_model)}\\n')\n",
    "print(ppo_model.v_head)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2913ef05-737e-4cdf-9bac-467ee6cf9f76",
   "metadata": {},
   "source": [
    "During PPO, only a few parameters will be updated. Specifically, the parameters of the `ValueHead`. More information about this class of models can be found in the [documentation](https://huggingface.co/docs/trl/main/en/models#trl.create_reference_model). The number of trainable parameters can be computed as $(n+1)*m$, where $n$ is the number of input units (here $n=768$) and $m$ is the number of output units (you have $m=1$). The $+1$ term in the equation takes into account the bias term."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76c7e2df-0c75-4bd0-bf8d-1f1545ef864e",
   "metadata": {},
   "source": [
    "Now create a frozen copy of the PPO which will not be fine-tuned - a reference model. The reference model will represent the LLM before detoxification. None of the parameters of the reference model will be updated during PPO training. This is on purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18a9b30a-ad14-4189-8088-d4de447fe247",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_model = create_reference_model(ppo_model)\n",
    "\n",
    "print(f'Reference model parameters to be updated:\\n{print_number_of_trainable_model_parameters(ref_model)}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a14848-e83d-4fdd-bc68-eb770c5951d7",
   "metadata": {
    "tags": []
   },
   "source": [
    "Everything is set. It is time to prepare the reward model!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bfdf1f7-3509-4adc-812a-2b22bd330137",
   "metadata": {},
   "source": [
    "<a name='2.2'></a>\n",
    "### 2.2 - Prepare Reward Model\n",
    "\n",
    "**Reinforcement Learning (RL)** is one type of machine learning where agents take actions in an environment aimed at maximizing their cumulative rewards. The agent's behavior is defined by the **policy**. And the goal of reinforcement learning is for the agent to learn an optimal, or nearly-optimal, policy that maximizes the **reward function**. \n",
    "\n",
    "In the [previous section](#2.1) the original policy is based on the instruct PEFT model - this is the LLM before detoxification. Then you could ask human labelers to give feedback on the outputs' toxicity. However, it can be expensive to use them for the entire fine-tuning process. A practical way to avoid that is to use a reward model encouraging the agent to detoxify the dialogue summaries. The intuitive approach would be to do some form of sentiment analysis across two classes (`nothate` and `hate`) and give a higher reward if there is higher a chance of getting class `nothate` as an output. \n",
    "\n",
    "For example, we can mention that having human labelers for the entire finetuning process can be expensive. A practical way to avoid that is to use a reward model.\n",
    "\n",
    "use feedback generated by a model\n",
    "\n",
    "You will use [Meta AI's RoBERTa-based hate speech model](https://huggingface.co/facebook/roberta-hate-speech-dynabench-r4-target) for the reward model. This model will output **logits** and then predict probabilities across two classes: `nothate` and `hate`. The logits of the output `nothate` will be taken as a positive reward. Then, the model will be fine-tuned with PPO using those reward values.\n",
    "\n",
    "Create the instance of the required model class for the RoBERTa model. You also need to load a tokenizer to test the model. Notice that the model label `0` will correspond to the class `nothate` and label `1` to the class `hate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f038a9f-e04b-49bc-8923-8ef3816919ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toxicity_model_name = \"facebook/roberta-hate-speech-dynabench-r4-target\"\n",
    "toxicity_tokenizer = AutoTokenizer.from_pretrained(toxicity_model_name, device_map=\"auto\")\n",
    "toxicity_model = AutoModelForSequenceClassification.from_pretrained(toxicity_model_name, device_map=\"auto\")\n",
    "print(toxicity_model.config.id2label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79d68799-a6e8-42d7-8d61-002e47210c18",
   "metadata": {
    "tags": []
   },
   "source": [
    "Take some non-toxic text, tokenize it, and pass it to the model. Print the output logits, probabilities, and the corresponding reward that will be used for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f4e6a05-2398-4ca7-a176-d5a1ff27fe39",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "non_toxic_text = \"#Person 1# tells Tommy that he didn't like the movie.\"\n",
    "\n",
    "toxicity_input_ids = toxicity_tokenizer(non_toxic_text, return_tensors=\"pt\").input_ids\n",
    "\n",
    "logits = toxicity_model(input_ids=toxicity_input_ids).logits\n",
    "print(f'logits [not hate, hate]: {logits.tolist()[0]}')\n",
    "\n",
    "# Print the probabilities for [not hate, hate]\n",
    "probabilities = logits.softmax(dim=-1).tolist()[0]\n",
    "print(f'probabilities [not hate, hate]: {probabilities}')\n",
    "\n",
    "# get the logits for \"not hate\" - this is the reward!\n",
    "not_hate_index = 0\n",
    "nothate_reward = (logits[:, not_hate_index]).tolist()\n",
    "print(f'reward (high): {nothate_reward}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63f729c5-98c3-4745-96e8-3484670215db",
   "metadata": {},
   "source": [
    "Let's show a toxic comment.  This will have a low reward because it is more toxic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ffc81e4-b220-4f4e-95ec-98ac418612d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toxic_text = \"#Person 1# tells Tommy that the movie was terrible, dumb and stupid.\"\n",
    "\n",
    "toxicity_input_ids = toxicity_tokenizer(toxic_text, return_tensors=\"pt\").input_ids\n",
    "\n",
    "logits = toxicity_model(toxicity_input_ids).logits\n",
    "print(f'logits [not hate, hate]: {logits.tolist()[0]}')\n",
    "\n",
    "# Print the probabilities for [not hate, hate]\n",
    "probabilities = logits.softmax(dim=-1).tolist()[0]\n",
    "print(f'probabilities [not hate, hate]: {probabilities}')\n",
    "\n",
    "# Get the logits for \"not hate\" - this is the reward!\n",
    "nothate_reward = (logits[:, not_hate_index]).tolist() \n",
    "print(f'reward (low): {nothate_reward}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc6e656e-fefe-4623-8bbb-472a8cf1c3c5",
   "metadata": {},
   "source": [
    "Setup Hugging Face inference pipeline to simplify the code for the toxicity reward model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73aab8c6-33eb-4bf4-a816-3866fc3460af",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = 0 if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "sentiment_pipe = pipeline(\"sentiment-analysis\", \n",
    "                          model=toxicity_model_name, \n",
    "                          device=device)\n",
    "reward_logits_kwargs = {\n",
    "    \"top_k\": None, # Return all scores.\n",
    "    \"function_to_apply\": \"none\", # Set to \"none\" to retrieve raw logits.\n",
    "    \"batch_size\": 16\n",
    "}\n",
    "\n",
    "reward_probabilities_kwargs = {\n",
    "    \"top_k\": None, # Return all scores.\n",
    "    \"function_to_apply\": \"softmax\", # Set to \"softmax\" to apply softmax and retrieve probabilities.\n",
    "    \"batch_size\": 16\n",
    "}\n",
    "\n",
    "print(\"Reward model output:\")\n",
    "print(\"For non-toxic text\")\n",
    "print(sentiment_pipe(non_toxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(non_toxic_text, **reward_probabilities_kwargs))\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",
   "id": "21302d74-59d8-451f-b287-e86245bf3324",
   "metadata": {},
   "source": [
    "The outputs are the logits for both `nothate` (positive) and `hate` (negative) classes. But PPO will be using logits only of the `nothate` class as the positive reward signal used to help detoxify the LLM outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ff3925-70c4-495c-ae26-68e2fc36296b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sentiment_pipe(non_toxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(non_toxic_text, **reward_probabilities_kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d11618b-5887-489a-b390-2139e364987f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sentiment_pipe(toxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(toxic_text, **reward_probabilities_kwargs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56513033-9bb1-41d5-81e2-54d1249c5c89",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a name='2.3'></a>\n",
    "### 2.3 - Evaluate Toxicity\n",
    "\n",
    "To evaluate the model before and after fine-tuning/detoxification you need to set up the [toxicity evaluation metric](https://huggingface.co/spaces/evaluate-measurement/toxicity). The **toxicity score** is a decimal value between 0 and 1 where 1 is the highest toxicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de8e99d-60ea-48a2-bdf5-817f80b48979",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toxicity_evaluator = evaluate.load(\"toxicity\", \n",
    "                                    toxicity_model_name,\n",
    "                                    module_type=\"measurement\",\n",
    "                                    toxic_label=\"hate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "840fbc47-c5c2-469a-b5f2-6407e8f0bfde",
   "metadata": {
    "tags": []
   },
   "source": [
    "Try to calculate toxicity for the same sentences as in section [2.2](#2.2). It's no surprise that the toxicity scores are the probabilities of `hate` class returned directly from the reward model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5298f91c-30d1-4d17-95a7-553952ac97b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toxicity_score = toxicity_evaluator.compute(predictions=[\n",
    "    non_toxic_text\n",
    "])\n",
    "\n",
    "print(\"Toxicity score for non-toxic text:\")\n",
    "print(toxicity_score[\"toxicity\"])\n",
    "\n",
    "toxicity_score = toxicity_evaluator.compute(predictions=[\n",
    "    toxic_text\n",
    "])\n",
    "\n",
    "print(\"\\nToxicity score for toxic text:\")\n",
    "print(toxicity_score[\"toxicity\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d3e835b-14b9-4646-b1c9-c975ef3ea944",
   "metadata": {
    "tags": []
   },
   "source": [
    "This evaluator can be used to compute the toxicity of the dialogues prepared in section [2.1](#2.1). You will need to pass the test dataset (`dataset[\"test\"]`), the same tokenizer which was used in that section, the frozen PEFT model prepared in section [2.2](#2.2), and the toxicity evaluator. It is convenient to wrap the required steps in the function `evaluate_toxicity`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "316ab128-33ff-4a1e-8936-47bfa29d48a3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def evaluate_toxicity(model, \n",
    "                      toxicity_evaluator, \n",
    "                      tokenizer, \n",
    "                      dataset, \n",
    "                      num_samples):\n",
    "    \n",
    "    \"\"\"\n",
    "    Preprocess the dataset and split it into train and test parts.\n",
    "\n",
    "    Parameters:\n",
    "    - model (trl model): Model to be evaluated.\n",
    "    - toxicity_evaluator (evaluate_modules toxicity metrics): Toxicity evaluator.\n",
    "    - tokenizer (transformers tokenizer): Tokenizer to be used.\n",
    "    - dataset (dataset): Input dataset for the evaluation.\n",
    "    - num_samples (int): Maximum number of samples for the evaluation.\n",
    "        \n",
    "    Returns:\n",
    "    tuple: A tuple containing two numpy.float64 values:\n",
    "    - mean (numpy.float64): Mean of the samples toxicity.\n",
    "    - std (numpy.float64): Standard deviation of the samples toxicity.\n",
    "    \"\"\"\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",
    "                                             top_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)])\n",
    "\n",
    "        toxicities.extend(toxicity_score[\"toxicity\"])\n",
    "\n",
    "    # Compute mean & std using np.\n",
    "    mean = np.mean(toxicities)\n",
    "    std = np.std(toxicities)\n",
    "        \n",
    "    return mean, std"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed269c3-dbd7-4d45-bc44-c6ab6d4ae141",
   "metadata": {
    "tags": []
   },
   "source": [
    "And now perform the calculation of the model toxicity before fine-tuning/detoxification:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c11ede15-dc1a-4a7e-a60d-b9cadfc7d876",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, device_map=\"auto\")\n",
    "\n",
    "mean_before_detoxification, std_before_detoxification = evaluate_toxicity(model=ref_model, \n",
    "                                                                          toxicity_evaluator=toxicity_evaluator, \n",
    "                                                                          tokenizer=tokenizer, \n",
    "                                                                          dataset=dataset[\"test\"], \n",
    "                                                                          num_samples=10)\n",
    "\n",
    "print(f'toxicity [mean, std] before detox: [{mean_before_detoxification}, {std_before_detoxification}]')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ba81c90-1ac8-4403-ac1a-d4c75c6df4f0",
   "metadata": {},
   "source": [
    "<a name='3'></a>\n",
    "## 3 - Perform Fine-Tuning to Detoxify the Summaries\n",
    "Optimize a RL policy against the reward model using Proximal Policy Optimization (PPO)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5516e318-8fce-4ca7-bf19-b7baf5255480",
   "metadata": {},
   "source": [
    "<a name='3.1'></a>\n",
    "### 3.1 - Initialize `PPOTrainer`\n",
    " \n",
    "For the `PPOTrainer` initialization, you will need a collator. Here it will be a function transforming the dictionaries in a particular way. You can define and test it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b7be1c0-382a-4fe2-8174-470f3e333e84",
   "metadata": {},
   "outputs": [],
   "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)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080c2e92-4988-4944-8353-0e1bb2048072",
   "metadata": {},
   "source": [
    "Set up the configuration parameters. Load the `ppo_model` and the tokenizer. You will also load a frozen version of the model `ref_model`. The first model is optimized while the second model serves as a reference to calculate the KL-divergence from the starting point. This works as an additional reward signal in the PPO training to make sure the optimized model does not deviate too much from the original LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "494e09a1-9024-4f38-91eb-d73cdc3239e6",
   "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=dataset[\"train\"], \n",
    "                         data_collator=collator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ad77d2c-3800-4e15-bb38-3851d94ad374",
   "metadata": {},
   "source": [
    "<a name='3.2'></a>\n",
    "### 3.2 - Fine-Tune the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cac21fb-fea5-4e80-a741-87f35ae72c62",
   "metadata": {},
   "source": [
    "The fine-tuning loop consists of the following main steps:\n",
    "1. Get the query responses from the policy LLM (PEFT model).\n",
    "2. Get sentiments for query/responses from hate speech RoBERTa model.\n",
    "3. Optimize policy with PPO using the (query, response, reward) triplet.\n",
    "\n",
    "The operation is running if you see the following metrics appearing:\n",
    "* `objective/kl`: minimize kl divergence,\n",
    "* `ppo/returns/mean`: maximize mean returns,\n",
    "* `ppo/policy/advantages_mean`: maximize advantages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bc55397-92b8-4f61-9ec2-c8b39d5f8962",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "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",
    "    # Break when you reach max_steps.\n",
    "    if step >= max_ppo_steps:\n",
    "        break   \n",
    "\n",
    "    prompt_tensors = batch[\"input_ids\"]\n",
    "\n",
    "    # Get response from FLAN-T5/PEFT LLM.\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",
    "    # This needs to be called \"response\".\n",
    "    batch[\"response\"] = [tokenizer.decode(r.squeeze()) for r in summary_tensors]\n",
    "\n",
    "    # Compute reward outputs.\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",
    "    # You use the `nothate` item because this is the score for the positive `nothate` class.\n",
    "    reward_tensors = [torch.tensor(reward[not_hate_index][\"score\"]) for reward in rewards]    \n",
    "\n",
    "    # Run PPO step.\n",
    "    stats = ppo_trainer.step(prompt_tensors, summary_tensors, reward_tensors)\n",
    "    ppo_trainer.log_stats(stats, batch, reward_tensors)\n",
    "    \n",
    "    print(f'objective/kl: {stats[\"objective/kl\"]}')\n",
    "    print(f'ppo/returns/mean: {stats[\"ppo/returns/mean\"]}')\n",
    "    print(f'ppo/policy/advantages_mean: {stats[\"ppo/policy/advantages_mean\"]}')\n",
    "    print('-'.join('' for x in range(100)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7903f5df-a9de-41eb-b239-38bc367b5654",
   "metadata": {},
   "source": [
    "<a name='3.3'></a>\n",
    "### 3.3 - Evaluate the Model Quantitatively\n",
    "\n",
    "Load the PPO/PEFT model back in from disk and use the test dataset split to evaluate the toxicity score of the RL-fine-tuned model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b093d43-6197-4cc0-b933-29030479a7d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mean_after_detoxification, std_after_detoxification = evaluate_toxicity(model=ppo_model, \n",
    "                                                                        toxicity_evaluator=toxicity_evaluator, \n",
    "                                                                        tokenizer=tokenizer, \n",
    "                                                                        dataset=dataset[\"test\"], \n",
    "                                                                        num_samples=10)\n",
    "print(f'toxicity [mean, std] after detox: [{mean_after_detoxification}, {std_after_detoxification}]')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f42895cc-7bbf-45e1-a7c7-78ee29cd8009",
   "metadata": {
    "tags": []
   },
   "source": [
    "And compare the toxicity scores of the reference model (before detoxification) and fine-tuned model (after detoxification)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77cc3af2-6600-4673-874b-917c05247ae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_improvement = (mean_before_detoxification - mean_after_detoxification) / mean_before_detoxification\n",
    "std_improvement = (std_before_detoxification - std_after_detoxification) / std_before_detoxification\n",
    "\n",
    "print(f'Percentage improvement of toxicity score after detoxification:')\n",
    "print(f'mean: {mean_improvement*100:.2f}%')\n",
    "print(f'std: {std_improvement*100:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66030581-b6f7-41d7-a7e6-2466226833be",
   "metadata": {},
   "source": [
    "<a name='3.4'></a>\n",
    "### 3.4 - Evaluate the Model Qualitatively\n",
    "\n",
    "Let's inspect some examples from the test dataset. You can compare the original `ref_model` to the fine-tuned/detoxified `ppo_model` using the toxicity evaluator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22cc8313-20ae-4d32-855e-9b2866fa3085",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 20\n",
    "compare_results = {}\n",
    "\n",
    "df_batch = dataset[\"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",
    "# Get response from ppo and base model.\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",
    "# Decode responses.\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",
    "# Sentiment analysis of query/response pairs before/after.\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[not_hate_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[not_hate_index][\"score\"] for reward in rewards_after]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65b70892-4c22-4bed-9d1e-9da3f4f0c97f",
   "metadata": {
    "tags": []
   },
   "source": [
    "Store and review the results in a DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06fff0f-9dba-4517-9424-a5ebd81e8f49",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "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",
   "id": "e7fb2477-f719-48de-b169-0607d355a8f6",
   "metadata": {},
   "source": [
    "Looking at the reward mean/median of the generated sequences you can observe a significant difference!"
   ]
  }
 ],
 "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
   },
   {
    "_defaultOrder": 57,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.trn1.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 58,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1.32xlarge",
    "vcpuNum": 128
   },
   {
    "_defaultOrder": 59,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1n.32xlarge",
    "vcpuNum": 128
   }
  ],
  "instance_type": "ml.m5.2xlarge",
  "kernelspec": {
   "display_name": "Python 3.11 (torch)",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}