05b_RLHF_Fine_Tune_Model_to_Detoxify_Summaries.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7f10b19-c061-42b2-8b42-e3cbafa3b1da",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Fine-Tune with Reinforcement Learning (PPO) and PEFT to Generate Less-Toxic Summaries\n",
    "\n",
    "![RLHF Pipeline](./img/generative_ai_pipeline_rlhf_plus.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aba27695-f61c-47e1-8c0c-626161c9fae7",
   "metadata": {
    "tags": []
   },
   "source": [
    "## First, check that the correct kernel is chosen.\n",
    "\n",
    "<img src=\"img/kernel_set_up.png\" width=\"300\"/>\n",
    "\n",
    "You can click on that to see and check the details of the image, kernel, and instance type.\n",
    "\n",
    "<img src=\"img/w3_kernel_and_instance_type.png\" width=\"600\"/>"
   ]
  },
  {
   "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": [
    "## Load FLAN-T5 Model, Prepare Reward Model and Toxicity Evaluator"
   ]
  },
  {
   "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": "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",
    "    # 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": "code",
   "execution_count": null,
   "id": "e1d44a53-ea1f-4fa5-89e7-d46e37d19935",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!aws s3 cp --recursive s3://dsoaws/models/peft-dialogue-summary-checkpoint/ ./peft-dialogue-summary-checkpoint-from-s3/ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4288240d-764b-4c49-8df7-b30b9277adbd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!ls -alh ./peft-dialogue-summary-checkpoint-from-s3/adapter_model.bin"
   ]
  },
  {
   "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": "ed358373-b1fd-47ba-a8e2-ebf0bef8e884",
   "metadata": {
    "tags": []
   },
   "source": [
    "![](img/peft_lora.png)"
   ]
  },
  {
   "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\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')"
   ]
  },
  {
   "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": "4bfdf1f7-3509-4adc-812a-2b22bd330137",
   "metadata": {},
   "source": [
    "## Setup Reward Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7bbbcfd-65bf-4a9d-906a-c5b58929abd0",
   "metadata": {},
   "source": [
    "![](img/hf_facebook_hatespeec_reward_model.png)"
   ]
  },
  {
   "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": "37d3e87a-f471-46f4-bb02-e11573dcfb1f",
   "metadata": {},
   "source": [
    "![](img/rlhf_reward_model_binary_classifier.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f4e6a05-2398-4ca7-a176-d5a1ff27fe39",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "non_toxic_text = \"I want to kiss you\"\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": "code",
   "execution_count": null,
   "id": "0ffc81e4-b220-4f4e-95ec-98ac418612d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "toxic_text = \"You are disgusting and terrible and i damn hate you\"\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": "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 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(\"\\nReward model output for toxic text:\")\n",
    "print(sentiment_pipe(toxic_text, **reward_logits_kwargs))\n",
    "print(sentiment_pipe(toxic_text, **reward_probabilities_kwargs))"
   ]
  },
  {
   "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": [
    "## Evaluate Toxicity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de8e99d-60ea-48a2-bdf5-817f80b48979",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "toxicity_evaluator = evaluate.load(\"toxicity\", \n",
    "                                    toxicity_model_name,\n",
    "                                    module_type=\"measurement\",\n",
    "                                    toxic_label=\"hate\")"
   ]
  },
  {
   "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": "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",
    "    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)])\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": "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=ppo_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": {
    "tags": []
   },
   "source": [
    "## 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": "2c905e12-21b5-45e0-9243-1e0aaffc58a1",
   "metadata": {},
   "source": [
    "![](img/rlhf_ppo.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15009325-0e3b-468e-b739-40659da06416",
   "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": "51035bb9-1a13-45c0-b5a7-6746305a5b8c",
   "metadata": {},
   "source": [
    "![](img/rlhf_kl_divergence.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b7be1c0-382a-4fe2-8174-470f3e333e84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import PPOConfig, PPOTrainer\n",
    "\n",
    "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",
    "def collator(data):\n",
    "    return dict((key, [d[key] for d in data]) for key in data[0])\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": "56c7352d-53fd-41d8-b438-cfcc9979b0c7",
   "metadata": {},
   "source": [
    "### Fine-Tune the Model"
   ]
  },
  {
   "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": [
    "## Evaluate the Model Quantitatively"
   ]
  },
  {
   "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": "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": {
    "tags": []
   },
   "source": [
    "## Evaluate the Model Qualitatively"
   ]
  },
  {
   "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": "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": "9f02880f-aa87-4c88-aee6-33e46f247169",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Release Resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9609e981-0410-4b44-b9de-9c467d39f435",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%html\n",
    "\n",
    "<p><b>Shutting down your kernel for this notebook to release resources.</b></p>\n",
    "<button class=\"sm-command-button\" data-commandlinker-command=\"kernelmenu:shutdown\" style=\"display:none;\">Shutdown Kernel</button>\n",
    "        \n",
    "<script>\n",
    "try {\n",
    "    els = document.getElementsByClassName(\"sm-command-button\");\n",
    "    els[0].click();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}    \n",
    "</script>"
   ]
  }
 ],
 "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.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science 3.0)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/sagemaker-data-science-310-v1"
  },
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}