app_ppo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinforcement Learning from Human Feedback (RLHF)\n",
    "\n",
    "## Enhancing T5-Base Summarization with Proximal Policy Optimization (PPO) and PEFT Fine-Tuning\n",
    "\n",
    "\n",
    "Reinforcement Learning from Human Feedback, commonly known as **RLHF**, is a specialized machine learning approach that amalgamates traditional reinforcement learning techniques and human expertise. This union offers a unique pathway to training artificial intelligence agents.\n",
    "\n",
    "---\n",
    "\n",
    "### Key Insights:\n",
    "\n",
    "1. **Nature of RLHF**: RLHF can be understood as an iterative procedure. The system undergoes continuous improvement, adapting its learning function based on newly acquired human feedback.\n",
    "  \n",
    "2. **Safety and Trust**: Incorporating human feedback ensures the system not only comprehends the tasks it should execute but also recognizes actions it should avoid. This dual capability fosters safer and more trustworthy systems.\n",
    "  \n",
    "3. **Performance Enhancements**: A study in 2022 evidenced that RLHF outperforms conventional supervised learning (SL). This superiority can be attributed to RLHF's ability to assess cumulative rewards for coherent conversations, a nuanced understanding that SL misses.\n",
    "\n",
    "---\n",
    "\n",
    "RLHF has proven instrumental in guiding language models, molding them to align better with intricate human values. As we venture into this notebook, we'll deep-dive into the methodologies and applications of RLHF.\n",
    "\n",
    "\n",
    "\n",
    "Useful references: \n",
    "\n",
    "https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py\n",
    "\n",
    "https://www.kaggle.com/code/paultimothymooney/fine-tune-flan-t5-with-ppo-deeplearning-ai\n",
    "\n",
    "https://github.com/huggingface/trl/blob/main/tests/test_ppo_trainer.py\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reward model: ideally a SequenceClassification type of model: We will use Bert\n",
    "\n",
    "Policy model: ideally a Seq2SeqLM: We will use T5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Alt text](image-1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Image source: https://huggingface.co/docs/trl/index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process Overview\n",
    "\n",
    "In this notebook, we embark on the journey of aligning a model using Reinforcement Learning from Human Feedback (RLHF). We'll employ various specialized models and leverage a structured training loop for this purpose.\n",
    "\n",
    "---\n",
    "\n",
    "### Models Utilized:\n",
    "\n",
    "1. **Rewards Model**: \n",
    "   - A finely-tuned model designated for dispensing rewards based on the actions of the policy model.\n",
    "\n",
    "2. **Base Model (Policy Model)**:\n",
    "   - The core model we aim to align using RLHF.\n",
    "   - During the RL process, this model becomes the \"policy model\", driving decisions and actions.\n",
    "\n",
    "3. **Reference Model**:\n",
    "   - A frozen replica of the base model.\n",
    "   - Its primary role is to act as a benchmark, monitoring the evolution of the policy model throughout the RL process.\n",
    "\n",
    "---\n",
    "\n",
    "### Training loop Overview:\n",
    "\n",
    "We begin by initializing the Proximal Policy Optimization (PPO) training class. The training process encompasses the following steps:\n",
    "\n",
    "- **Generation of Summaries**: \n",
    "  - Derived from the policy model.\n",
    "  \n",
    "- **Reward Assignment**:\n",
    "  - The generated summaries are channeled through the rewards model.\n",
    "  - Based on these summaries, rewards are determined, reflecting the alignment of the policy model with human preferences.\n",
    "  \n",
    "- **Model Adjustment via PPO**:\n",
    "  - Utilizing the acquired rewards, PPO refines the weights of the policy model, nudging it closer to human preferences.\n",
    "  \n",
    "This iterative training loop continues for a predefined number of steps.\n",
    "\n",
    "---\n",
    "\n",
    "## Evaluation:\n",
    "\n",
    "Post-training, we evaluate the efficacy and alignment of the policy model post-RL to determine its proficiency in mirroring human preferences.\n",
    "\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch\n",
    "!pip install transformers\n",
    "!pip install datasets\n",
    "!pip install trl\n",
    "!pip install peft\n",
    "!pip install numpy\n",
    "!pip install pandas\n",
    "!pip install tqdm\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A quick hack to link this notebook to WanDB\n",
    "\n",
    "In this case it is redundant because the transformers libraries will do it, but as an educational gesture, this is how you could install WanDB in a notebook that doesn't contain ibraries already prep'ed with WDB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mjcolanotoro\u001b[0m (\u001b[33mjcolano\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.10"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>c:\\Dropbox\\Desktop\\MAchineLearning\\rlhf\\experiment01\\wandb\\run-20230912_150053-imqwjpsi</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/jcolano/rlhf_ppo_v1/runs/imqwjpsi' target=\"_blank\">desert-sun-4</a></strong> to <a href='https://wandb.ai/jcolano/rlhf_ppo_v1' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/jcolano/rlhf_ppo_v1' target=\"_blank\">https://wandb.ai/jcolano/rlhf_ppo_v1</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/jcolano/rlhf_ppo_v1/runs/imqwjpsi' target=\"_blank\">https://wandb.ai/jcolano/rlhf_ppo_v1/runs/imqwjpsi</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/jcolano/rlhf_ppo_v1/runs/imqwjpsi?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x211fe0cf7c0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "import random\n",
    "\n",
    "# start a new wandb run to track this script\n",
    "wandb.init(\n",
    "    # set the wandb project where this run will be logged\n",
    "    project=\"rlhf_ppo_v1\",\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Python310\\lib\\site-packages\\bitsandbytes\\cextension.py:34: UserWarning: The installed version of bitsandbytes was compiled without GPU support. 8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable.\n",
      "  warn(\"The installed version of bitsandbytes was compiled without GPU support. \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'NoneType' object has no attribute 'cadam32bit_grad_fp32'\n"
     ]
    }
   ],
   "source": [
    "import torch \n",
    "\n",
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer, T5Tokenizer, T5ForConditionalGeneration\n",
    "\n",
    "from torch.utils.data import DataLoader, Dataset as TorchDataset\n",
    "from torch.optim import AdamW\n",
    "\n",
    "from datasets import load_dataset, Dataset as HFDataset\n",
    "\n",
    "from peft import PeftModel, PeftConfig,  TaskType\n",
    "\n",
    "from peft import (\n",
    "    get_peft_config,\n",
    "    get_peft_model,\n",
    "    get_peft_model_state_dict,\n",
    "    set_peft_model_state_dict,\n",
    "    PeftType,\n",
    "    LoraConfig,\n",
    ")\n",
    "\n",
    "# AutoModelForCausalLMWithValueHead & AutoModelForSeq2SeqLMWithValueHead: A transformer model with an additional scalar output for each token which can be used as a value function in reinforcement learning.\n",
    "# https://huggingface.co/docs/trl/models#trl.AutoModelForSeq2SeqLMWithValueHead\n",
    "\n",
    "# trl: Transformer Reinforcement Learning library\n",
    "import trl \n",
    "from trl import PPOTrainer, PPOConfig, AutoModelForSeq2SeqLMWithValueHead # https://huggingface.co/docs/trl/quickstart\n",
    "from trl import create_reference_model\n",
    "from trl.core import LengthSampler\n",
    "\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()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reward Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Alt text](image-2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Image source: https://huggingface.co/blog/rlhf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Model in Reinforcement Learning (RL)\n",
    "\n",
    "In RL, a **reward model** is a mechanism providing feedback to the agent about its performance in its environment. Instead of predefined reward functions, reward models infer the reward signal from human feedback, especially useful in complex scenarios where crafting a reward function is challenging.\n",
    "\n",
    "### Why is it Important?\n",
    "\n",
    "- **Feedback Mechanism**: It's how agents determine if actions are beneficial or detrimental.\n",
    "- **Facilitates Learning**: Agents use these signals to update their policies to maximize rewards.\n",
    "- **Handles Complexity**: For real-world problems where explicit reward functions are difficult, a learned reward model is valuable.\n",
    "- **Safety and Alignment**: They ensure RL agents' objectives align with human intentions, reducing potential harmful behaviors.\n",
    "\n",
    "In our code, we're initializing a reward model (based on a transformer like BERT) for RL with Human Feedback (RLHF). This model generates reward signals from the agent's interactions, steering its learning process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at ./model_bert_hf_experiment2/ were not used when initializing BertForSequenceClassification: ['bert.encoder.layer.3.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.3.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.10.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.10.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.7.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.1.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.9.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.5.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.2.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.4.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.7.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.2.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.4.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.8.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.7.attention.self.value.lora_A.default.weight', 'classifier.original_module.bias', 'bert.encoder.layer.9.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.1.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.6.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.5.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.8.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.0.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.6.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.1.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.6.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.3.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.11.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.9.attention.self.query.lora_A.default.weight', 'classifier.original_module.weight', 'bert.encoder.layer.0.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.4.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.9.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.1.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.8.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.11.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.4.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.2.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.11.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.10.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.3.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.6.attention.self.value.lora_A.default.weight', 'bert.encoder.layer.7.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.0.attention.self.query.lora_B.default.weight', 'classifier.modules_to_save.default.bias', 'bert.encoder.layer.0.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.10.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.5.attention.self.value.lora_B.default.weight', 'classifier.modules_to_save.default.weight', 'bert.encoder.layer.2.attention.self.query.lora_A.default.weight', 'bert.encoder.layer.5.attention.self.query.lora_B.default.weight', 'bert.encoder.layer.8.attention.self.value.lora_B.default.weight', 'bert.encoder.layer.11.attention.self.query.lora_B.default.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at ./model_bert_hf_experiment2/ and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertForSequenceClassification(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Specify the directory where you saved the model and tokenizer\n",
    "reward_model_directory = \"JuanKO/rlhf_reward_model\"\n",
    "\n",
    "rm_model = AutoModelForSequenceClassification.from_pretrained(reward_model_directory)\n",
    "rm_tokenizer = AutoTokenizer.from_pretrained(reward_model_directory)\n",
    "rm_model.to(device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function: `score_summaries`\n",
    "\n",
    "### Description:\n",
    "The `score_summaries` function is designed to score two summaries, `chosen_summary` and `rejected_summary`, within the context of a Reinforcement Learning with Human Feedback (RLHF) loop. It tokenizes the inputs, obtains the logits from a given model, computes the softmax probabilities, and finally extracts the scores (probabilities) and logits associated with each summary.\n",
    "\n",
    "### Parameters:\n",
    "\n",
    "- **model** (`torch.nn.Module`): \n",
    "    - The PyTorch model that produces logits given an input.\n",
    "  \n",
    "- **tokenizer** (`transformers.PreTrainedTokenizer`): \n",
    "    - A tokenizer object used to tokenize input summaries.\n",
    "  \n",
    "- **chosen_summary** (`str`): \n",
    "    - The chosen summary string that needs to be scored.\n",
    "  \n",
    "- **rejected_summary** (`str`): \n",
    "    - The rejected summary string that needs to be scored.\n",
    "\n",
    "### Returns:\n",
    "\n",
    "- **chosen_score** (`float`): \n",
    "    - The probability score associated with the `chosen_summary` being positive or \"good\".\n",
    "\n",
    "- **rejected_score** (`float`): \n",
    "    - The probability score associated with the `rejected_summary` being positive or \"good\".\n",
    "\n",
    "- **chosen_logit** (`float`): \n",
    "    - The logit value associated with the `chosen_summary`.\n",
    "\n",
    "- **rejected_logit** (`float`): \n",
    "    - The logit value associated with the `rejected_summary`.\n",
    "\n",
    "### Function Flow:\n",
    "\n",
    "1. **Tokenization**: \n",
    "    - The input summaries, `chosen_summary` and `rejected_summary`, are tokenized using the provided tokenizer. These tokenized inputs are padded or truncated to a maximum length of 512 tokens.\n",
    "\n",
    "2. **Move to Device**: \n",
    "    - The tokenized tensors are transferred to the device (likely a GPU or CPU) where the model resides.\n",
    "\n",
    "3. **Obtain Logits**: \n",
    "    - The tokenized tensors are passed through the model to obtain logits. This is done in a no-gradient context to ensure computational efficiency and prevent any updates to the model.\n",
    "\n",
    "4. **Compute Probabilities**: \n",
    "    - The obtained logits are passed through a softmax function to get the associated probabilities. This helps in understanding how likely each summary is deemed \"good\" by the model.\n",
    "\n",
    "5. **Extract Scores and Logits**: \n",
    "    - The function then extracts the probability and logit associated with the positive class (assumed to be the second class in the logits) for both summaries.\n",
    "\n",
    "### Notes:\n",
    "- The function assumes that the positive class (indicating the summary is \"good\") is the second class in the logits.\n",
    "- The softmax function ensures that the logits are converted into probabilities that sum up to 1.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "def score_summaries(model, tokenizer, chosen_summary, rejected_summary):\n",
    "    # Tokenize the inputs\n",
    "    chosen_tokens = tokenizer(chosen_summary, return_tensors=\"pt\", padding='max_length', truncation=True, max_length=512)\n",
    "    rejected_tokens = tokenizer(rejected_summary, return_tensors=\"pt\", padding='max_length', truncation=True, max_length=512)\n",
    "    \n",
    "    chosen_tokens.to(device)\n",
    "    rejected_tokens.to(device)\n",
    "    \n",
    "    # Get logits from the model\n",
    "    with torch.no_grad():\n",
    "        chosen_logits = model(**chosen_tokens).logits\n",
    "        rejected_logits = model(**rejected_tokens).logits\n",
    "    \n",
    "    # Apply softmax to get probabilities\n",
    "    chosen_probs = F.softmax(chosen_logits, dim=-1)\n",
    "    rejected_probs = F.softmax(rejected_logits, dim=-1)\n",
    "\n",
    "    # Assuming the positive class (indicating 'chosen' is good) is the second one\n",
    "    chosen_score = chosen_probs[0][1].item()\n",
    "    rejected_score = rejected_probs[0][1].item()\n",
    "    \n",
    "    # Extract logits for each summary\n",
    "    chosen_logit = chosen_logits[0][1].item()\n",
    "    rejected_logit = rejected_logits[0][1].item()\n",
    "\n",
    "    return chosen_score, rejected_score, chosen_logit, rejected_logit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run some examples to test the function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In this test, we evaluate the `score_summaries` function using two sample summaries: one labeled as `chosen_summary` and the other as `rejected_summary`. These summaries are tokenized, scored, and the associated logits are obtained using our reward model (`rm_model`) and its tokenizer (`rm_tokenizer`).\n",
    "\n",
    "### Sample Summaries:\n",
    "\n",
    "- **Chosen Summary**: \n",
    "    - \"Water meter in another condo is not in our condo. What can we do legally to restore water to my condo complex?\"\n",
    "    \n",
    "- **Rejected Summary**: \n",
    "    - \"Go fix the problem.\"\n",
    "\n",
    "### Test Execution:\n",
    "\n",
    "The `score_summaries` function is called with the provided model, tokenizer, and the sample summaries. The returned scores and logits for each summary are then printed.\n",
    "\n",
    "### Expected Output:\n",
    "\n",
    "- **Chosen Score**: \n",
    "    - This gives the probability score of the `chosen_summary` being perceived as \"good\" or positive by the model.\n",
    "  \n",
    "- **Rejected Score**: \n",
    "    - This gives the probability score of the `rejected_summary` being perceived as \"good\" or positive by the model.\n",
    "  \n",
    "- **Chosen Logit**:\n",
    "    - This returns the raw logit value associated with the `chosen_summary`.\n",
    "  \n",
    "- **Rejected Logit**:\n",
    "    - This returns the raw logit value associated with the `rejected_summary`.\n",
    "\n",
    "### Notes:\n",
    "- Higher scores indicate a higher probability of the summary being perceived as positive or \"good\".\n",
    "- The logit values provide insight into the raw outputs of the model before being passed through the softmax function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "chosen_summary = \"Water meter in another condo is not in our condo. What can we do legally to restore water to my condo complex?\"\n",
    "rejected_summary = \"Go fix the problem.\"\n",
    "\n",
    "chosen_score, rejected_score, chosen_logit, rejected_logit = score_summaries(rm_model, rm_tokenizer, chosen_summary, rejected_summary)\n",
    "\n",
    "print(f\"Chosen Score: {chosen_score:.4f}\")\n",
    "print(f\"Rejected Score: {rejected_score:.4f}\")\n",
    "\n",
    "print(f\"Chosen Logit: {chosen_logit:.4f}\")\n",
    "print(f\"Rejected Logit: {rejected_logit:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the T5 Model for RLHF Fine-Tuning\n",
    "\n",
    "### Overview:\n",
    "\n",
    "T5, short for \"Text-to-Text Transfer Transformer\", is a state-of-the-art model designed to handle various text-to-text tasks. In this section, we'll be loading a T5 model that is intended to be fine-tuned using the Reinforcement Learning with Human Feedback (RLHF) approach.\n",
    "\n",
    "### Steps:\n",
    "\n",
    "1. **Model Selection**:\n",
    "    - We've selected the T5 model for our fine-tuning process. Specifically, we'll be working with the \"t5-base\" variant which offers a balance between computational efficiency and performance.\n",
    "\n",
    "2. **Loading Model and Tokenizer**:\n",
    "    - `policy_model_path`: Specifies the directory path where our pre-trained (or fine-tuned) T5 model is saved.\n",
    "    - `policy_model_name`: Indicates the model name, which in this case is \"t5-base\".\n",
    "    - Using the `T5ForConditionalGeneration.from_pretrained` method, we load the model weights from our specified path.\n",
    "    - Similarly, the corresponding tokenizer, which is essential for converting text into a format that the T5 model can understand, is loaded using the `T5Tokenizer.from_pretrained` method.\n",
    "\n",
    "3. **Device Allocation**:\n",
    "    - The model is assigned to a computation device (either CPU or GPU) using the `.to(device)` method. This ensures efficient computation, especially when working with large datasets.\n",
    "\n",
    "### Test the Model:\n",
    "\n",
    "After loading, it's a good practice to perform some inference tests to ensure that the model is loaded correctly and is functioning as expected.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f053061f4e9f4d8aba21c6b032c1d991",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading spiece.model:   0%|          | 0.00/792k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d93e80b4dc154510873542d0464dd1ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/2.20k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42d1470e03fd484da9246a717012ca4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/2.37k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "policy_model_path = \"JuanKO/rlhf_base_model\"\n",
    "policy_model_name = \"t5-base\" \n",
    "\n",
    "policy_model = T5ForConditionalGeneration.from_pretrained(policy_model_path)\n",
    "policy_model.to(device)\n",
    "policy_tokenizer = T5Tokenizer.from_pretrained(policy_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing the T5 Model for Summarization\n",
    "\n",
    "\n",
    "After loading our T5 model, we'll test its summarization capabilities on a sample text from the r/relationships subreddit. This test will help us understand the model's performance and its readiness for RLHF fine-tuning.\n",
    "\n",
    "### Steps:\n",
    "\n",
    "1. **Setting the Task Prefix**:\n",
    "    - We use the prefix \"summarize: \" to indicate to the T5 model the type of task we want it to perform.\n",
    "\n",
    "2. **Sample Text**:\n",
    "    - We have selected a post from the r/relationships subreddit to be summarized. This text provides context about a user's relationship concerns related to her bisexuality.\n",
    "\n",
    "3. **Generating the Summary**:\n",
    "    - We feed the concatenated task prefix and text into our T5 model.\n",
    "    - The model then processes this input and returns a concise summary. The `generate` function is used to obtain this output, and we've set a max length of 100 tokens for our summary.\n",
    "\n",
    "4. **Decoding the Summary**:\n",
    "    - The output from the T5 model is in the form of token IDs. Using the T5 tokenizer's `decode` method, we convert these tokens back into human-readable text.\n",
    "\n",
    "5. **Scoring the Summary using the Reward Model**:\n",
    "    - With the generated summary in hand, we then use our previously defined `score_summaries` function to evaluate the quality of the summary.\n",
    "    - This function returns a score and logit value for both the chosen summary and a rejected (blank) summary. Higher scores and logits suggest better alignment with what the reward model considers a good summary.\n",
    "\n",
    "### Results:\n",
    "\n",
    "By examining the printed scores and logits, we can gauge the perceived quality of the generated summary according to our reward model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TL;DR: I'm bisexual and I'm in a hetero relationship. Is it necessary to tell my boyfriend that I'm bisexual? When do you think is the right time?\n",
      "Chosen Score: 0.5943\n",
      "Rejected Score: 0.5193\n",
      "Chosen Logit: 0.0889\n",
      "Rejected Logit: 0.2162\n"
     ]
    }
   ],
   "source": [
    "task_prefix = \"summarize: \" \n",
    "\n",
    "text = \"SUBREDDIT: r/relationships TITLE: How do I/do I at all [20 F] tell my boyfriend [23 M] that I'm bisexual? POST: I've had two serious relationships prior to this one, both with women. They had no problem with me being bisexual and it was something known before the relationship -- my first girlfriend was also bisexual. I am now in a relationship with a guy. We've been exclusive for about a month. Having never faced this issue, I come to you, Reddit. Is this something that he needs to know? Is it really relevant to a hetero relationship, regardless of if one of the participants in the relationship is bisexual? If you guys think it is necessary, when do you think is the right time? I think my biggest fear is losing him because of it. I know that I should be with someone who is fine with who I am, but I really like the guy and I'd hate for my sexual orientation to be the thing that kills this.\"\n",
    "#text = \"SUBREDDIT: r/legaladvice TITLE: What can I do legally to restore water to my condominium!? POST: Hi, I live in SE Michigan in a condominium complex. Our water was shut off due to non-payment. (we recieved no notice) and we had to pay all that was due ($1500) We payed this yesterday at 2, they said the water would be turned on immediately. It wasn't. It's now the next day. The lady in our assosciation keeps insisting that the water meter is in another condo. Which we can't access because the person living there is never there (it's being rented) Now we're stuck with no water, no shower, no teeth brushing, no toilets, and no food for certain meals.... Please help us... What can we do? We called the police and they say that we can file a civil report for the lady not doing her job...\"\n",
    "prompt = f\"{task_prefix}{text}\"\n",
    "input_ids = policy_tokenizer(prompt, return_tensors=\"pt\").input_ids.to(device)\n",
    "outputs = policy_model.generate(input_ids, max_length=100).to(device)\n",
    "\n",
    "strOutput = policy_tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "print(strOutput)\n",
    "\n",
    "chosen_score, rejected_score, chosen_logit, rejected_logit = score_summaries(rm_model, rm_tokenizer, strOutput, \"\")\n",
    "\n",
    "print(f\"Chosen Score: {chosen_score:.4f}\")\n",
    "print(f\"Rejected Score: {rejected_score:.4f}\")\n",
    "\n",
    "print(f\"Chosen Logit: {chosen_logit:.4f}\")\n",
    "print(f\"Rejected Logit: {rejected_logit:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the T5 Model for Peft + LoRA\n",
    "\n",
    "### Overview:\n",
    "\n",
    "Peft and LoRA (Low-Rank Adaptation) are techniques that enable efficient fine-tuning of pre-trained models by introducing low-rank structures into the models. Here, we'll configure the T5 model for this process.\n",
    "\n",
    "### Steps:\n",
    "\n",
    "1. **Setting up the LoRA Configuration**:\n",
    "    - `LoraConfig` provides the configuration settings for Low-Rank Adaptation.\n",
    "        - `r`: Rank of the low-rank structure. In this instance, it's set to 8.\n",
    "        - `lora_alpha`: Scaling factor for the newly introduced low-rank parameters.\n",
    "        - `target_modules`: Specifies which parts of the model to apply LoRA. Here, we're targeting the \"q\" (query) and \"v\" (value) modules.\n",
    "        - `lora_dropout`: Dropout rate for the low-rank parameters. Set to 0.10, or 10%.\n",
    "        - `bias`: Specifies the type of bias for the low-rank projection. We've chosen \"none\" in this case.\n",
    "        - `task_type`: Indicates the type of task. As we're using T5, the task type is set to `SEQ_2_SEQ_LM`.\n",
    "\n",
    "2. **Applying LoRA Configuration to T5**:\n",
    "    - Using the `get_peft_model` function, we apply the LoRA configuration to our pre-loaded T5 model.\n",
    "    - The returned model (`policy_peft_model`) is equipped with the Peft + LoRA modifications and is ready for fine-tuning.\n",
    "\n",
    "### Summary of this section:\n",
    "\n",
    "Our T5 model is now prepared with Peft + LoRA adjustments. This configuration optimizes the model for more efficient fine-tuning on specific tasks while leveraging the powerful pre-trained knowledge.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForSeq2SeqLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): T5ForConditionalGeneration(\n",
       "      (shared): Embedding(32128, 768)\n",
       "      (encoder): T5Stack(\n",
       "        (embed_tokens): Embedding(32128, 768)\n",
       "        (block): ModuleList(\n",
       "          (0): T5Block(\n",
       "            (layer): ModuleList(\n",
       "              (0): T5LayerSelfAttention(\n",
       "                (SelfAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (relative_attention_bias): Embedding(32, 12)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (1): T5LayerFF(\n",
       "                (DenseReluDense): T5DenseActDense(\n",
       "                  (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                  (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                  (act): ReLU()\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (1-11): 11 x T5Block(\n",
       "            (layer): ModuleList(\n",
       "              (0): T5LayerSelfAttention(\n",
       "                (SelfAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (1): T5LayerFF(\n",
       "                (DenseReluDense): T5DenseActDense(\n",
       "                  (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                  (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                  (act): ReLU()\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (final_layer_norm): T5LayerNorm()\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (decoder): T5Stack(\n",
       "        (embed_tokens): Embedding(32128, 768)\n",
       "        (block): ModuleList(\n",
       "          (0): T5Block(\n",
       "            (layer): ModuleList(\n",
       "              (0): T5LayerSelfAttention(\n",
       "                (SelfAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (relative_attention_bias): Embedding(32, 12)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (1): T5LayerCrossAttention(\n",
       "                (EncDecAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (2): T5LayerFF(\n",
       "                (DenseReluDense): T5DenseActDense(\n",
       "                  (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                  (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                  (act): ReLU()\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (1-11): 11 x T5Block(\n",
       "            (layer): ModuleList(\n",
       "              (0): T5LayerSelfAttention(\n",
       "                (SelfAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (1): T5LayerCrossAttention(\n",
       "                (EncDecAttention): T5Attention(\n",
       "                  (q): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  (v): Linear(\n",
       "                    in_features=768, out_features=768, bias=False\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (2): T5LayerFF(\n",
       "                (DenseReluDense): T5DenseActDense(\n",
       "                  (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                  (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                  (act): ReLU()\n",
       "                )\n",
       "                (layer_norm): T5LayerNorm()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (final_layer_norm): T5LayerNorm()\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (lm_head): Linear(in_features=768, out_features=32128, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lora_config = LoraConfig(\n",
    "    r=8, # Rank\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q\", \"v\"],\n",
    "    lora_dropout=0.10,\n",
    "    bias=\"none\",\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM # T5\n",
    ")\n",
    "\n",
    "policy_peft_model = get_peft_model(policy_model, lora_config)\n",
    "policy_peft_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyzing Trainable Parameters in the Peft + LoRA Configured T5 Model\n",
    "\n",
    "After applying the Peft + LoRA configuration to our T5 model, it's essential to inspect the model's parameters to understand its structure better.\n",
    "\n",
    "### Key Insights:\n",
    "\n",
    "1. **Trainable Parameters**:\n",
    "    - This refers to the parameters that will be updated during the training process.\n",
    "    - In our configured model, there are **884,736** trainable parameters.\n",
    "\n",
    "2. **Total Parameters**:\n",
    "    - This indicates the complete count of parameters present in the model, including those that are non-trainable.\n",
    "    - The model consists of **223,788,288** total parameters.\n",
    "\n",
    "3. **Percentage of Trainable Parameters**:\n",
    "    - It's useful to know the fraction of the model's parameters that are trainable, as this can influence training time and model flexibility.\n",
    "    - Only about **0.3953%** (or roughly 0.4%) of the entire model's parameters are trainable.\n",
    "\n",
    "### Summary of this section:\n",
    "\n",
    "The Peft + LoRA configuration results in a model where only a small fraction of parameters are trainable. This approach offers a balance, as it allows for specific fine-tuning while leveraging a vast pre-trained structure. The advantage is that it can lead to faster training times and might prevent overfitting, especially when training data is limited.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 884736 || all params: 223788288 || trainable%: 0.3953450861557152\n"
     ]
    }
   ],
   "source": [
    "policy_peft_model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Alt text](image-3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Image source: https://magazine.sebastianraschka.com/p/llm-training-rlhf-and-its-alternatives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiating the PPO Model with Value Head\n",
    "\n",
    "Proximal Policy Optimization (PPO) is a reinforcement learning algorithm. In this step, we set up the model for PPO training using our earlier `policy_peft_model`.\n",
    "\n",
    "### Key Components:\n",
    "\n",
    "1. **AutoModelForSeq2SeqLMWithValueHead**:\n",
    "    - An extension of the transformers model that includes a scalar output for each token, aiding in reinforcement learning.\n",
    "    - This model can capture the value function, an estimate of future rewards.\n",
    "\n",
    "2. **Inputs**:\n",
    "    - We pass in our `policy_peft_model`, which has been configured with Peft + LoRA, as the foundation for our PPO model.\n",
    "    - We set `torch_dtype` to `torch.bfloat16` for numerical precision and memory efficiency.\n",
    "    - The `is_trainable` flag is set to `True`, allowing us to further fine-tune the model using our RL loop.\n",
    "\n",
    "3. **Device Assignment**:\n",
    "    - We transfer our instantiated model to the appropriate device (`device`) for computation, ensuring efficient training.\n",
    "\n",
    "### Summary of this section:\n",
    "\n",
    "With our PPO model instantiated, we're poised to fine-tune our summarization model using reinforcement learning with human feedback. This approach is aimed at improving the model's performance in generating summaries based on human preferences and judgments.\n",
    "\n",
    "[More on PPO and TRL](https://huggingface.co/docs/trl/quickstart)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AutoModelForSeq2SeqLMWithValueHead(\n",
       "  (pretrained_model): PeftModelForSeq2SeqLM(\n",
       "    (base_model): LoraModel(\n",
       "      (model): T5ForConditionalGeneration(\n",
       "        (shared): Embedding(32128, 768)\n",
       "        (encoder): T5Stack(\n",
       "          (embed_tokens): Embedding(32128, 768)\n",
       "          (block): ModuleList(\n",
       "            (0): T5Block(\n",
       "              (layer): ModuleList(\n",
       "                (0): T5LayerSelfAttention(\n",
       "                  (SelfAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (relative_attention_bias): Embedding(32, 12)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (1): T5LayerFF(\n",
       "                  (DenseReluDense): T5DenseActDense(\n",
       "                    (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                    (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                    (dropout): Dropout(p=0.1, inplace=False)\n",
       "                    (act): ReLU()\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "            (1-11): 11 x T5Block(\n",
       "              (layer): ModuleList(\n",
       "                (0): T5LayerSelfAttention(\n",
       "                  (SelfAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (1): T5LayerFF(\n",
       "                  (DenseReluDense): T5DenseActDense(\n",
       "                    (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                    (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                    (dropout): Dropout(p=0.1, inplace=False)\n",
       "                    (act): ReLU()\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (final_layer_norm): T5LayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (decoder): T5Stack(\n",
       "          (embed_tokens): Embedding(32128, 768)\n",
       "          (block): ModuleList(\n",
       "            (0): T5Block(\n",
       "              (layer): ModuleList(\n",
       "                (0): T5LayerSelfAttention(\n",
       "                  (SelfAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (relative_attention_bias): Embedding(32, 12)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (1): T5LayerCrossAttention(\n",
       "                  (EncDecAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (2): T5LayerFF(\n",
       "                  (DenseReluDense): T5DenseActDense(\n",
       "                    (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                    (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                    (dropout): Dropout(p=0.1, inplace=False)\n",
       "                    (act): ReLU()\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "            (1-11): 11 x T5Block(\n",
       "              (layer): ModuleList(\n",
       "                (0): T5LayerSelfAttention(\n",
       "                  (SelfAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (1): T5LayerCrossAttention(\n",
       "                  (EncDecAttention): T5Attention(\n",
       "                    (q): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "                    (v): Linear(\n",
       "                      in_features=768, out_features=768, bias=False\n",
       "                      (lora_dropout): ModuleDict(\n",
       "                        (default): Dropout(p=0.1, inplace=False)\n",
       "                      )\n",
       "                      (lora_A): ModuleDict(\n",
       "                        (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                      )\n",
       "                      (lora_B): ModuleDict(\n",
       "                        (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                      )\n",
       "                      (lora_embedding_A): ParameterDict()\n",
       "                      (lora_embedding_B): ParameterDict()\n",
       "                    )\n",
       "                    (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (2): T5LayerFF(\n",
       "                  (DenseReluDense): T5DenseActDense(\n",
       "                    (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "                    (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "                    (dropout): Dropout(p=0.1, inplace=False)\n",
       "                    (act): ReLU()\n",
       "                  )\n",
       "                  (layer_norm): T5LayerNorm()\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (final_layer_norm): T5LayerNorm()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (lm_head): Linear(in_features=768, out_features=32128, bias=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (v_head): ValueHead(\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "    (summary): Linear(in_features=768, out_features=1, bias=True)\n",
       "    (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# https://huggingface.co/docs/trl/quickstart\n",
    "ppo_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(policy_peft_model,                                                               \n",
    "                                                               torch_dtype=torch.bfloat16,\n",
    "                                                               is_trainable=True)\n",
    "\n",
    "ppo_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the Reference Model\n",
    "\n",
    "In reinforcement learning, especially when fine-tuning models using methods like Proximal Policy Optimization (PPO), it's helpful to have a reference model. This model represents the initial state or behavior of the learner model (in this case, the Language Model) before any alignment or optimization. It aids in calculating the importance sampling ratio, a critical component for stable and effective updates in PPO.\n",
    "\n",
    "### Key Components:\n",
    "\n",
    "1. **create_reference_model**:\n",
    "    - A function provided by Huggingface's TRL (Transformer Reinforcement Learning) library.\n",
    "    - Creates a duplicate of the passed model which acts as a reference during the RL fine-tuning process.\n",
    "\n",
    "2. **Inputs**:\n",
    "    - The `policy_model` we previously defined serves as the input. This model acts as the basis for our reference model.\n",
    "\n",
    "3. **Device Assignment**:\n",
    "    - Once instantiated, we move our reference model to the specified device (`device`) for computations.\n",
    "\n",
    "### Summary of this section:\n",
    "\n",
    "By defining a reference model, we set a stable baseline against which we can measure and guide the progress and changes of our main model during the reinforcement learning process.\n",
    "\n",
    "[More on TRL and Reference Models](https://huggingface.co/docs/trl/models#trl.create_reference_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "T5ForConditionalGeneration(\n",
       "  (shared): Embedding(32128, 768)\n",
       "  (encoder): T5Stack(\n",
       "    (embed_tokens): Embedding(32128, 768)\n",
       "    (block): ModuleList(\n",
       "      (0): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (relative_attention_bias): Embedding(32, 12)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1-11): 11 x T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): T5LayerNorm()\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (decoder): T5Stack(\n",
       "    (embed_tokens): Embedding(32128, 768)\n",
       "    (block): ModuleList(\n",
       "      (0): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (relative_attention_bias): Embedding(32, 12)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1-11): 11 x T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(\n",
       "                in_features=768, out_features=768, bias=False\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=768, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=768, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): T5LayerNorm()\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=768, out_features=32128, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ref_model = create_reference_model(policy_model)\n",
    "ref_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the Dataset for Reinforcement Learning\n",
    "\n",
    "Reinforcement learning (RL) requires a dataset to simulate experiences and provide feedback. In our RL setup for fine-tuning a language model, we utilize a comparison dataset.\n",
    "\n",
    "### Steps:\n",
    "\n",
    "1. **Load Dataset**:\n",
    "    - Using Huggingface's `datasets` library, we fetch the 'CarperAI/openai_summarize_comparisons' dataset's test split.\n",
    "\n",
    "2. **Filtering**:\n",
    "    - We want to ensure the prompt lengths are manageable. \n",
    "    - Filtering by word count: We retain samples where the prompt has ≤ 450 words.\n",
    "    - (Alternative Filtering by character count is commented out for reference.)\n",
    "\n",
    "3. **Shuffling and Sampling**:\n",
    "    - To ensure a diverse set of samples, we shuffle the dataset.\n",
    "    - We then select a subset (2,000 samples in this instance) for the RL process.\n",
    "\n",
    "4. **Feature Extraction**:\n",
    "    - From our shuffled dataset, we focus on the `prompt` and `chosen` fields. \n",
    "    - Rename the 'chosen' field to 'response' to align with the PPO library's requirements.\n",
    "\n",
    "5. **Dataset Conversion**:\n",
    "    - Convert the dictionary containing our features into a Huggingface Dataset format.\n",
    "\n",
    "6. **Train-Eval Split**:\n",
    "    - Split the dataset into training and evaluation subsets. \n",
    "    - Here, 80% of samples are designated for training, and the remaining 20% are for evaluation.\n",
    "\n",
    "### Outcome:\n",
    "\n",
    "By the end of this process, we will have a training dataset and an evaluation dataset ready for the RL process. These datasets will be essential in guiding the model's fine-tuning and assessing its performance during the RL loop.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset parquet (C:/Users/juan_/.cache/huggingface/datasets/CarperAI___parquet/CarperAI--openai_summarize_comparisons-79d2c222a15dc8fb/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n",
      "Loading cached processed dataset at C:\\Users\\juan_\\.cache\\huggingface\\datasets\\CarperAI___parquet\\CarperAI--openai_summarize_comparisons-79d2c222a15dc8fb\\0.0.0\\2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec\\cache-d5c2170aaeb9b06c.arrow\n",
      "Loading cached shuffled indices for dataset at C:\\Users\\juan_\\.cache\\huggingface\\datasets\\CarperAI___parquet\\CarperAI--openai_summarize_comparisons-79d2c222a15dc8fb\\0.0.0\\2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec\\cache-f81caef5de9ecb07.arrow\n"
     ]
    }
   ],
   "source": [
    "# Load the dataset\n",
    "orig_dataset = load_dataset('CarperAI/openai_summarize_comparisons', split='test')\n",
    "\n",
    "# Filter samples where the prompt length is less than or equal to 750\n",
    "filtered_dataset = orig_dataset.filter(lambda example: len(example['prompt'].split()) <= 450) # By word\n",
    "#filtered_dataset = orig_dataset.filter(lambda example: len(example['prompt']) <= 1250) # By character\n",
    "\n",
    "# Shuffle and select the first 10K samples\n",
    "#shuffled_dataset = orig_dataset.shuffle(seed=42).select(range(1000))\n",
    "shuffled_dataset = filtered_dataset.shuffle(seed=42).select(range(2000)) \n",
    "\n",
    "\n",
    "# Extract the desired features.  Renaming chose to response to follow the ppo library requirements.\n",
    "new_dataset_dict = {\n",
    "    \"prompt\": shuffled_dataset[\"prompt\"],\n",
    "    \"response\": shuffled_dataset[\"chosen\"]\n",
    "}\n",
    "\n",
    "# Convert the dictionary to a new Dataset\n",
    "dataset = HFDataset.from_dict(new_dataset_dict)\n",
    "\n",
    "# Split the new_dataset into train_dataset and eval_dataset\n",
    "split_ratio = 0.8  # 80% for training, 20% for evaluation\n",
    "num_train_samples = int(split_ratio * len(dataset))\n",
    "train_dataset = dataset.select(range(num_train_samples))\n",
    "eval_dataset = dataset.select(range(num_train_samples, len(dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['prompt', 'response'])\n",
      "dict_keys(['prompt', 'response'])\n"
     ]
    }
   ],
   "source": [
    "print(train_dataset[0].keys())\n",
    "print(eval_dataset[0].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenization of Datasets\n",
    "\n",
    "For reinforcement learning, it is crucial that the data is in a format understood by the model. This requires tokenizing our textual data into numerical tokens. Here, we'll use the tokenizer associated with our model (T5 in this case) to process our datasets.\n",
    "\n",
    "### Steps:\n",
    "\n",
    "1. **Tokenizer Initialization**:\n",
    "    - Instantiate the tokenizer corresponding to our model (T5). If you use a different model, ensure you fetch the right tokenizer.\n",
    "\n",
    "2. **Tokenization Function**:\n",
    "    - Define a function (`tokenize_function`) that:\n",
    "        - Processes the 'prompt' in each example of the dataset.\n",
    "        - Truncates or pads the tokenized prompt to a maximum length of 512 tokens.\n",
    "        - Returns the tokenized 'input_ids' for each 'prompt' and retains the associated 'response'.\n",
    "\n",
    "3. **Apply Tokenization**:\n",
    "    - Apply the `tokenize_function` to both the training and evaluation datasets using the `map` function.\n",
    "\n",
    "### Outcome:\n",
    "\n",
    "The datasets (`train_dataset` and `eval_dataset`) are now tokenized and in a suitable format for model ingestion during the reinforcement learning loop.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "031aaf8ea63e4da3a008449676bebc2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1600 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "850b9646c93a47de90eb9cdc162d7b38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/400 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import T5Tokenizer\n",
    "\n",
    "# Instantiate your tokenizer (replace T5Tokenizer with your model's tokenizer if different)\n",
    "tokenizer = T5Tokenizer.from_pretrained(\"t5-small\") # or whatever model you're using\n",
    "\n",
    "def tokenize_function(example):\n",
    "    # Tokenize the prompt and store it as input_ids. Also return the response.\n",
    "    return {\n",
    "        \"input_ids\": tokenizer(example[\"prompt\"], return_tensors=\"pt\", truncation=True, max_length=512)[\"input_ids\"].squeeze(),\n",
    "        \"response\": example[\"response\"],\n",
    "    }\n",
    "\n",
    "# Tokenize the training and evaluation datasets\n",
    "train_dataset = train_dataset.map(tokenize_function, batched=False)\n",
    "eval_dataset = eval_dataset.map(tokenize_function, batched=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['prompt', 'response', 'input_ids'],\n",
       "    num_rows: 1600\n",
       "})"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'prompt': \"SUBREDDIT: r/relationship_advice\\nTITLE: [20/m] My girlfriend [20/f] has become very distant and weird\\nPOST: I have been in a relationship with my girlfriend for a little bit over 1 year. We recently had a breakup because I was distant and she thought I was cheating on her (which I wasn't). Before the breakup, she wanted to spend as much time with me as she could, but recently she has been very distant. We used to go to eachothers places overnight almost daily, but nowadays she does not want to come over to my place or want me to go over to hers (We both live on our own). She also used to talk to me all the time on facebook, but now she pretty much only replies to what I talk, and does not try to keep the conversation going. She has became pretty slow at replying, but when I'm with her, she replies instantly to her other friends who text her. \\n\\nI'm really lost at this situation, because I feel like she does not want to be with me anymore. I know that she's taking SSRI medications at the moment for anxiety. But she was taking them before we broke up for couple weeks. Her explanation for this situation is that she wants to spend some time alone, because she is working so much and never has time to be alone. This bothers me because she used to want to be with me all the time she could, and now she does not even try to schedule time for me.\", 'response': 'TL;DR:  gf has became very distant after we got back together, even though before we  broke up she was really into me.', 'input_ids': [3, 4138, 25582, 11253, 3177, 10, 3, 52, 87, 60, 6105, 2009, 834, 9, 26, 7287, 15, 332, 3177, 3765, 10, 784, 1755, 87, 51, 908, 499, 17442, 784, 1755, 87, 89, 908, 65, 582, 182, 10382, 11, 10088, 276, 14464, 10, 27, 43, 118, 16, 3, 9, 1675, 28, 82, 17442, 21, 3, 9, 385, 720, 147, 209, 215, 5, 101, 1310, 141, 3, 9, 1733, 413, 250, 27, 47, 10382, 11, 255, 816, 27, 47, 15009, 53, 30, 160, 41, 3339, 27, 2088, 31, 17, 137, 3103, 8, 1733, 413, 6, 255, 1114, 12, 1492, 38, 231, 97, 28, 140, 38, 255, 228, 6, 68, 1310, 255, 65, 118, 182, 10382, 5, 101, 261, 12, 281, 12, 284, 9269, 7, 1747, 8521, 966, 1444, 6, 68, 18811, 255, 405, 59, 241, 12, 369, 147, 12, 82, 286, 42, 241, 140, 12, 281, 147, 12, 160, 7, 41, 1326, 321, 619, 30, 69, 293, 137, 451, 92, 261, 12, 1350, 12, 140, 66, 8, 97, 30, 13301, 6, 68, 230, 255, 1134, 231, 163, 26719, 12, 125, 27, 1350, 6, 11, 405, 59, 653, 12, 453, 8, 3634, 352, 5, 451, 65, 1632, 1134, 2684, 44, 8776, 53, 6, 68, 116, 27, 31, 51, 28, 160, 6, 255, 26719, 8361, 12, 160, 119, 803, 113, 1499, 160, 5, 27, 31, 51, 310, 1513, 44, 48, 1419, 6, 250, 27, 473, 114, 255, 405, 59, 241, 12, 36, 28, 140, 7595, 5, 27, 214, 24, 255, 31, 7, 838, 3, 4256, 5593, 11208, 44, 8, 798, 21, 6261, 5, 299, 255, 47, 838, 135, 274, 62, 8238, 95, 21, 1158, 1274, 5, 1347, 7295, 21, 48, 1419, 19, 24, 255, 2746, 12, 1492, 128, 97, 2238, 6, 250, 255, 19, 464, 78, 231, 11, 470, 65, 97, 12, 36, 2238, 5, 100, 13965, 7, 140, 250, 255, 261, 12, 241, 12, 36, 28, 140, 66, 8, 97, 255, 228, 6, 11, 230, 255, 405, 59, 237, 653, 12, 2023, 97, 21, 140, 5, 1]}\n"
     ]
    }
   ],
   "source": [
    "# Lets check one sample of the train_dataset\n",
    "print(train_dataset[0])  # print the first example from the training dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameter Initialization\n",
    "\n",
    "Before training the model using reinforcement learning, we need to define several hyperparameters that will guide and constrain the training process.\n",
    "\n",
    "### Data Collation:\n",
    "\n",
    "- **`collator` Function**: \n",
    "    - A helper function that takes a list of data samples and merges them into a single batch, making it suitable for processing by the model.\n",
    "    - For instance, given an input of individual key-value data samples, the function groups the values by their keys.\n",
    "\n",
    "    Example:\n",
    "    ```python\n",
    "    test_data = [{\"key1\": \"value1\", \"key2\": \"value2\", \"key3\": \"value3\"}, {\"key1\": \"value4\", \"key2\": \"value5\", \"key3\": \"value6\"}]\n",
    "    collated_data = collator(test_data)\n",
    "    ```\n",
    "\n",
    "- **Sample Data**:\n",
    "    - To visually validate the output of the `collator`, a sample is taken from the training dataset and processed.\n",
    "\n",
    "### Key Hyperparameters:\n",
    "\n",
    "- **`learning_rate`**: \n",
    "    - Controls the step size at each iteration while moving towards a minimum in the loss function. Set to `1.41e-5`.\n",
    "\n",
    "- **`max_ppo_epochs`**: \n",
    "    - Specifies the maximum number of epochs for the Proximal Policy Optimization (PPO) training. Set to `3`.\n",
    "\n",
    "- **`mini_batch_size`** & **`batch_size`**: \n",
    "    - Determines the number of samples in each mini-batch (`4`) and the overall batch size (`16`).\n",
    "\n",
    "- **`DEFAULT_REJECTED_SUMMARY_TEXT`**: \n",
    "    - A placeholder text for a bad summary. This could potentially act as a regularizer during training, though its effect needs to be verified. \n",
    "\n",
    "- **Generation Constraints** (`generation_kwargs`):\n",
    "    - `temperature`: Controls the randomness of predictions by scaling the logits before applying softmax. Set to `1.0`.\n",
    "    - `min_length`: Minimum length of the generated text. Set to `5`.\n",
    "    - `top_k` & `top_p`: Parameters controlling the nucleus sampling method. Here, `top_k` is set to `0.0` and `top_p` to `1.0`, indicating no truncation based on these parameters.\n",
    "    - `do_sample`: Boolean value determining whether to sample the outputs. Set to `True`.\n",
    "\n",
    "- **Output Length Sampling**:\n",
    "    - `output_min_length` & `output_max_length`: Define the minimum (`100`) and maximum (`400`) lengths of generated outputs.\n",
    "    - `output_length_sampler`: Samples an output length between the specified min and max values.\n",
    "\n",
    "- **`max_ppo_steps`**: \n",
    "    - Determines the total number of PPO steps during training. Set to `100`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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\"}, {\"key1\": \"value4\", \"key2\": \"value5\", \"key3\": \"value6\"}]\n",
    "print(f'Collator input: {test_data}')\n",
    "print(f'Collator output: {collator(test_data)}')\n",
    "\n",
    "# Lets sample what the collator generates:\n",
    "sample_data = [train_dataset[i] for i in range(3)]  # take first three examples\n",
    "collated_data = collator(sample_data)\n",
    "print(collated_data.keys())\n",
    "\n",
    "learning_rate=1.41e-5\n",
    "max_ppo_epochs=3\n",
    "mini_batch_size=4\n",
    "batch_size=16\n",
    "\n",
    "# This is a HACK... lets see how this works out. May casue bias or may help. The good side is that this, being constant, can effect some type of regularization, preventing the model from gravitating too much towards any specific pattern in the training data.  Just a thought.\n",
    "DEFAULT_REJECTED_SUMMARY_TEXT = \"This is a bad summary\"\n",
    "\n",
    "# Some initial values\n",
    "output_min_length = 100\n",
    "output_max_length = 400\n",
    "output_length_sampler = LengthSampler(output_min_length, output_max_length)\n",
    "\n",
    "# These hyperparams guide the generation of the completion in the policy model. We could add other params like temperature.\n",
    "generation_kwargs = {\n",
    "    \"temperature\": 1.0,\n",
    "    \"min_length\": 5,\n",
    "    \"top_k\": 0.0,\n",
    "    \"top_p\": 1.0,\n",
    "    \"do_sample\": True\n",
    "}\n",
    "\n",
    "max_ppo_steps = 100\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration for PPO Training\n",
    "\n",
    "We leverage the `PPOConfig` from the Hugging Face `trl` library to set up the configuration required for the Proximal Policy Optimization (PPO) training.\n",
    "\n",
    "The `PPOConfig` requires and/or allows for a number of arguments that define the behavior of the PPO training loop:\n",
    "\n",
    "- **`model_name`**: \n",
    "    - Name of the model. Here, it is set as `policy_model_name`.\n",
    "\n",
    "- **`learning_rate`**: \n",
    "    - The rate at which the model adjusts based on the error during training. We've set it to the previously initialized value of `learning_rate`.\n",
    "\n",
    "- **`ppo_epochs`**: \n",
    "    - Specifies the number of epochs for PPO training. Set to the previously defined `max_ppo_epochs`.\n",
    "\n",
    "- **`mini_batch_size`**: \n",
    "    - The size of the smaller batches that the main batch is divided into, during training. Set to the previously initialized value of `mini_batch_size`.\n",
    "\n",
    "- **`batch_size`**: \n",
    "    - The number of data samples processed during each training step. We've set it to the previously initialized value of `batch_size`.\n",
    "\n",
    "For a more detailed understanding and potential additional configurations, one can refer to the [Hugging Face documentation on `trl.trainer`](https://huggingface.co/docs/trl/trainer).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check out https://huggingface.co/docs/trl/trainer\n",
    "\n",
    "config = PPOConfig(\n",
    "    model_name=policy_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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting Up the PPO Trainer\n",
    "\n",
    "To fine-tune the model using Proximal Policy Optimization (PPO), we use the `PPOTrainer` class from Hugging Face's `trl` library.\n",
    "\n",
    "The `PPOTrainer` class is initialized with several key arguments:\n",
    "\n",
    "- **`config`**: \n",
    "    - The configuration object created using `PPOConfig`. This contains the hyperparameters required for PPO training.\n",
    "\n",
    "- **`model`**: \n",
    "    - The model that will be fine-tuned. In this case, it is the `ppo_model` which was previously instantiated.\n",
    "\n",
    "- **`ref_model`**: \n",
    "    - The reference model, representing the model before alignment. We use `ref_model` for this purpose.\n",
    "\n",
    "- **`tokenizer`**: \n",
    "    - The tokenizer responsible for converting text into tokens suitable for the model's input. Here, it's the `policy_tokenizer` we set up before.\n",
    "\n",
    "- **`dataset`**: \n",
    "    - The training dataset. We use the tokenized `train_dataset`.\n",
    "\n",
    "- **`data_collator`**: \n",
    "    - A function to transform a list of samples to a batch. We use the `collator` function we defined earlier.\n",
    "\n",
    "This trainer will be used to conduct the PPO training loop, enabling us to fine-tune the model using reinforcement learning. \n",
    "\n",
    "For a deeper dive into the functionalities provided by the `PPOTrainer` class, one can refer to the [Hugging Face documentation on `trl.trainer`](https://huggingface.co/docs/trl/trainer).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check out https://huggingface.co/docs/trl/trainer\n",
    "\n",
    "ppo_trainer = PPOTrainer(config=config, \n",
    "                         model=ppo_model, \n",
    "                         ref_model=ref_model, \n",
    "                         tokenizer=policy_tokenizer, \n",
    "                         dataset=train_dataset, \n",
    "                         data_collator=collator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-Tuning with Reinforcement Learning\n",
    "\n",
    "Reinforcement learning offers a unique approach to fine-tuning models. The underlying principle is to allow the model to learn by receiving feedback (rewards) on its actions. In this context, an action would be generating a summary for a given text prompt.\n",
    "\n",
    "### Training Loop Overview\n",
    "\n",
    "The training loop we've crafted here follows this sequence of steps:\n",
    "\n",
    "1. **Model Prediction**: Using the policy language model (`ppo_trainer` in this case), we generate predicted summaries.\n",
    "2. **Score Generation**: We then pass these summaries to a reward model to assign a score (reward) based on the quality of the generated summary.\n",
    "3. **Model Update**: With the generated summaries and their respective scores, we use Proximal Policy Optimization (PPO) to update our policy language model.\n",
    "\n",
    "### Detailed Breakdown\n",
    "\n",
    "#### **1. Model Prediction**:\n",
    "\n",
    "- We iterate through our training data in batches (`prompt_tensors`).\n",
    "- For each prompt, we predict a summary (`summary_tensors`). This prediction is based on the generation hyperparameters we've specified (`generation_kwargs`), which guide the sampling strategy.\n",
    "\n",
    "#### **2. Score Generation**:\n",
    "\n",
    "- For each summary, we calculate a score by comparing it with a default rejected summary. \n",
    "- This step uses a separate reward model (`rm_model`), which assesses the quality of summaries.\n",
    "\n",
    "#### **3. Model Update**:\n",
    "\n",
    "- Using PPO, we update our policy model based on:\n",
    "  - The initial input (`prompt_tensors`).\n",
    "  - The generated summary (`summary_tensors`).\n",
    "  - The assigned reward (`reward_tensors`).\n",
    "  \n",
    "### Key Metrics:\n",
    "\n",
    "- `objective/kl`: Measures how different the policy's action distribution after the update is from the action distribution before the update. PPO tries to make these changes very small to avoid drastic changes.\n",
    "  \n",
    "- `ppo/returns/mean`: This is the average return achieved by the agent. Higher is better.\n",
    "\n",
    "- `ppo/policy/advantages_mean`: Measures how much better an action is than the average action at a given state. An advantage of zero means the action is just average, a positive advantage means it's better than average, and a negative one means it's worse than average.\n",
    "\n",
    "### Important Notes:\n",
    "\n",
    "- **HACK** Alert: The code contains certain hacks (like for handling variable sequence lengths) which were used to overcome specific issues during development.\n",
    "\n",
    "- **Reward Model**: The quality of the model training largely depends on the feedback it provides.\n",
    "\n",
    "### References:\n",
    "\n",
    "- [PPOTrainer in Hugging Face's TRL library](https://huggingface.co/docs/trl/trainer#trl.PPOTrainer)\n",
    "- [Using Transformer Reinforcement Learning to detoxify generative language models](https://medium.com/@ben.burtenshaw/using-transformer-reinforcement-learning-to-detoxify-generative-language-models-5198446d6786)\n",
    "- HuggingFace's example scripts in their GitHub repository.\n",
    "\n",
    "The success of reinforcement learning is deeply intertwined with the feedback mechanism and the quality of the reward signal.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]C:\\Users\\juan_\\AppData\\Local\\Temp\\ipykernel_12104\\415800031.py:50: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  prompt_tensor = torch.tensor(prompt_tensor).to(device)\n",
      "1it [00:16, 16.83s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0\n",
      "ppo/returns/mean: 0.3847939372062683\n",
      "ppo/policy/advantages_mean: 0.00486169196665287\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [00:34, 17.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.005069880746304989\n",
      "ppo/returns/mean: 0.38824570178985596\n",
      "ppo/policy/advantages_mean: 0.001255544601008296\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [00:53, 17.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.027740254998207092\n",
      "ppo/returns/mean: 0.4024064540863037\n",
      "ppo/policy/advantages_mean: 0.005844333209097385\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [01:08, 16.89s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.008470938540995121\n",
      "ppo/returns/mean: 0.4360160231590271\n",
      "ppo/policy/advantages_mean: 0.008006769232451916\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [01:26, 17.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.0028168456628918648\n",
      "ppo/returns/mean: 0.4166437089443207\n",
      "ppo/policy/advantages_mean: 0.0039040702395141125\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6it [01:43, 17.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.015527371317148209\n",
      "ppo/returns/mean: 0.4231906831264496\n",
      "ppo/policy/advantages_mean: 0.0028870929963886738\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "7it [02:00, 17.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.0020126374438405037\n",
      "ppo/returns/mean: 0.4269492030143738\n",
      "ppo/policy/advantages_mean: 0.0030096229165792465\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "8it [02:20, 17.93s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.06836467236280441\n",
      "ppo/returns/mean: 0.41612517833709717\n",
      "ppo/policy/advantages_mean: 0.0038263590540736914\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "9it [02:37, 17.70s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.08298622071743011\n",
      "ppo/returns/mean: 0.44846484065055847\n",
      "ppo/policy/advantages_mean: 0.0022622086107730865\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "10it [02:53, 17.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.005825418047606945\n",
      "ppo/returns/mean: 0.46129149198532104\n",
      "ppo/policy/advantages_mean: 0.003759412094950676\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11it [03:11, 17.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.005642858799546957\n",
      "ppo/returns/mean: 0.4399201571941376\n",
      "ppo/policy/advantages_mean: 0.0013956755865365267\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12it [03:27, 17.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.005834147334098816\n",
      "ppo/returns/mean: 0.4514174461364746\n",
      "ppo/policy/advantages_mean: 0.003677933244034648\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13it [03:44, 17.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.011828919872641563\n",
      "ppo/returns/mean: 0.4530404210090637\n",
      "ppo/policy/advantages_mean: 0.0011673825792968273\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "14it [04:01, 16.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.033650025725364685\n",
      "ppo/returns/mean: 0.4805707633495331\n",
      "ppo/policy/advantages_mean: 0.004944793879985809\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "15it [04:19, 17.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.04556742310523987\n",
      "ppo/returns/mean: 0.4878147542476654\n",
      "ppo/policy/advantages_mean: 0.00022911513224244118\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "16it [04:36, 17.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.052357759326696396\n",
      "ppo/returns/mean: 0.47049078345298767\n",
      "ppo/policy/advantages_mean: 0.0005096103996038437\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17it [04:54, 17.58s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.009967220947146416\n",
      "ppo/returns/mean: 0.4759097099304199\n",
      "ppo/policy/advantages_mean: -0.00103578413836658\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "18it [05:12, 17.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.02261665277183056\n",
      "ppo/returns/mean: 0.4794938266277313\n",
      "ppo/policy/advantages_mean: 0.004009743221104145\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "19it [05:31, 18.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0016215275973081589\n",
      "ppo/returns/mean: 0.4834381937980652\n",
      "ppo/policy/advantages_mean: 0.002844936214387417\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "20it [05:51, 18.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.014883600175380707\n",
      "ppo/returns/mean: 0.4874545931816101\n",
      "ppo/policy/advantages_mean: 0.001672036712989211\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "21it [06:11, 18.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.038705311715602875\n",
      "ppo/returns/mean: 0.48216789960861206\n",
      "ppo/policy/advantages_mean: 0.005550017114728689\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22it [06:27, 18.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.0019519738852977753\n",
      "ppo/returns/mean: 0.4977434575557709\n",
      "ppo/policy/advantages_mean: 0.002608383074402809\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23it [06:46, 18.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.015746327117085457\n",
      "ppo/returns/mean: 0.4924009442329407\n",
      "ppo/policy/advantages_mean: -0.0029126708395779133\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "24it [07:03, 17.84s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.13578671216964722\n",
      "ppo/returns/mean: 0.49223825335502625\n",
      "ppo/policy/advantages_mean: -0.0014680366730317473\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25it [07:18, 17.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.18424507975578308\n",
      "ppo/returns/mean: 0.5160964727401733\n",
      "ppo/policy/advantages_mean: 0.0035873521119356155\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "26it [07:36, 17.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.07853008806705475\n",
      "ppo/returns/mean: 0.5125787854194641\n",
      "ppo/policy/advantages_mean: 0.0018142350018024445\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "27it [07:54, 17.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.1536068469285965\n",
      "ppo/returns/mean: 0.5129472017288208\n",
      "ppo/policy/advantages_mean: -0.0031952625140547752\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "28it [08:13, 17.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.02651580050587654\n",
      "ppo/returns/mean: 0.5035720467567444\n",
      "ppo/policy/advantages_mean: 0.0005308630643412471\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "29it [08:33, 18.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.04060036689043045\n",
      "ppo/returns/mean: 0.5135884284973145\n",
      "ppo/policy/advantages_mean: 0.002005363814532757\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "30it [08:52, 18.82s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.06144071742892265\n",
      "ppo/returns/mean: 0.5074781179428101\n",
      "ppo/policy/advantages_mean: 0.006083873566240072\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "31it [09:09, 18.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.1026630848646164\n",
      "ppo/returns/mean: 0.497572124004364\n",
      "ppo/policy/advantages_mean: -0.004282218404114246\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "32it [09:24, 17.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.11476826667785645\n",
      "ppo/returns/mean: 0.5155744552612305\n",
      "ppo/policy/advantages_mean: -0.0009203864028677344\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "33it [09:41, 17.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.018813099712133408\n",
      "ppo/returns/mean: 0.5048032999038696\n",
      "ppo/policy/advantages_mean: 0.0009756293147802353\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "34it [09:59, 17.34s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.004020830616354942\n",
      "ppo/returns/mean: 0.50152587890625\n",
      "ppo/policy/advantages_mean: -0.0017703983467072248\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "35it [10:16, 17.33s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.028400693088769913\n",
      "ppo/returns/mean: 0.5098116397857666\n",
      "ppo/policy/advantages_mean: -0.00064779695821926\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "36it [10:35, 17.76s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.004725713282823563\n",
      "ppo/returns/mean: 0.5155090093612671\n",
      "ppo/policy/advantages_mean: 0.005265518091619015\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "37it [10:53, 17.94s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.03030284121632576\n",
      "ppo/returns/mean: 0.5224802494049072\n",
      "ppo/policy/advantages_mean: 0.000699183321557939\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "38it [11:12, 18.25s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0028322823345661163\n",
      "ppo/returns/mean: 0.5142924785614014\n",
      "ppo/policy/advantages_mean: 0.00593021884560585\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "39it [11:30, 18.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.06499931961297989\n",
      "ppo/returns/mean: 0.5349175930023193\n",
      "ppo/policy/advantages_mean: -0.0069395094178617\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "40it [11:50, 18.67s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.015764307230710983\n",
      "ppo/returns/mean: 0.5187349915504456\n",
      "ppo/policy/advantages_mean: -0.0023731673136353493\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "41it [12:07, 18.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.018244529142975807\n",
      "ppo/returns/mean: 0.5299811363220215\n",
      "ppo/policy/advantages_mean: -0.00019256211817264557\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "42it [12:23, 17.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.058882132172584534\n",
      "ppo/returns/mean: 0.5251530408859253\n",
      "ppo/policy/advantages_mean: -0.0028166293632239103\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "43it [12:42, 18.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.01506706140935421\n",
      "ppo/returns/mean: 0.5216745138168335\n",
      "ppo/policy/advantages_mean: -0.007766125723719597\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "44it [13:00, 17.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0036642001941800117\n",
      "ppo/returns/mean: 0.5320895910263062\n",
      "ppo/policy/advantages_mean: -0.000767915858887136\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "45it [13:18, 17.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.03506341576576233\n",
      "ppo/returns/mean: 0.5339188575744629\n",
      "ppo/policy/advantages_mean: 0.010133009403944016\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "46it [13:37, 18.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.03634041175246239\n",
      "ppo/returns/mean: 0.5373767018318176\n",
      "ppo/policy/advantages_mean: -0.0011674811830744147\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "47it [13:55, 18.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0006244629621505737\n",
      "ppo/returns/mean: 0.5391950011253357\n",
      "ppo/policy/advantages_mean: -0.008310972712934017\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "48it [14:11, 17.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.09847931563854218\n",
      "ppo/returns/mean: 0.5256244540214539\n",
      "ppo/policy/advantages_mean: -0.004047184716910124\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "49it [14:27, 17.18s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.007130159065127373\n",
      "ppo/returns/mean: 0.5404835939407349\n",
      "ppo/policy/advantages_mean: -0.0006340897525660694\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "50it [14:44, 17.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.09234826266765594\n",
      "ppo/returns/mean: 0.5280537605285645\n",
      "ppo/policy/advantages_mean: 0.00522202393040061\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "51it [15:01, 17.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.011652393266558647\n",
      "ppo/returns/mean: 0.5378813743591309\n",
      "ppo/policy/advantages_mean: 0.006647405680269003\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "52it [15:18, 17.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.010729705914855003\n",
      "ppo/returns/mean: 0.5417889356613159\n",
      "ppo/policy/advantages_mean: -0.0059839775785803795\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "53it [15:36, 17.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.0522296279668808\n",
      "ppo/returns/mean: 0.5455126166343689\n",
      "ppo/policy/advantages_mean: 0.003960065543651581\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "54it [15:53, 17.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.07429303228855133\n",
      "ppo/returns/mean: 0.5376936793327332\n",
      "ppo/policy/advantages_mean: 0.003455133643001318\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "55it [16:09, 16.90s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.015849264338612556\n",
      "ppo/returns/mean: 0.5331861972808838\n",
      "ppo/policy/advantages_mean: -0.0017464521806687117\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "56it [16:26, 17.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.10449449717998505\n",
      "ppo/returns/mean: 0.5360128879547119\n",
      "ppo/policy/advantages_mean: 0.0010660793632268906\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "57it [16:43, 16.93s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.018445100635290146\n",
      "ppo/returns/mean: 0.5541489124298096\n",
      "ppo/policy/advantages_mean: 0.00361231598071754\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "58it [16:59, 16.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.13619235157966614\n",
      "ppo/returns/mean: 0.5405611991882324\n",
      "ppo/policy/advantages_mean: 0.002603059634566307\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "59it [17:16, 16.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.054558731615543365\n",
      "ppo/returns/mean: 0.5453513860702515\n",
      "ppo/policy/advantages_mean: -0.0017076923977583647\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "60it [17:34, 17.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.10614021122455597\n",
      "ppo/returns/mean: 0.5431622266769409\n",
      "ppo/policy/advantages_mean: 0.002949802204966545\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "61it [17:52, 17.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.008595196530222893\n",
      "ppo/returns/mean: 0.5485168695449829\n",
      "ppo/policy/advantages_mean: 0.0012846844037994742\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "62it [18:10, 17.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.06707711517810822\n",
      "ppo/returns/mean: 0.5445464849472046\n",
      "ppo/policy/advantages_mean: 0.0017537561943754554\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "63it [18:28, 17.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.059157513082027435\n",
      "ppo/returns/mean: 0.5462327003479004\n",
      "ppo/policy/advantages_mean: 0.007264412939548492\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "64it [18:48, 18.35s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.019886992871761322\n",
      "ppo/returns/mean: 0.5484727621078491\n",
      "ppo/policy/advantages_mean: 0.005174936726689339\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "65it [19:06, 18.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.15200024843215942\n",
      "ppo/returns/mean: 0.548679530620575\n",
      "ppo/policy/advantages_mean: 0.0005625371704809368\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "66it [19:22, 17.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.07403381168842316\n",
      "ppo/returns/mean: 0.5488724112510681\n",
      "ppo/policy/advantages_mean: -0.003060923656448722\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "67it [19:42, 18.27s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.05367913842201233\n",
      "ppo/returns/mean: 0.5376991033554077\n",
      "ppo/policy/advantages_mean: 0.004013826604932547\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "68it [19:59, 17.96s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.03038051351904869\n",
      "ppo/returns/mean: 0.5573604106903076\n",
      "ppo/policy/advantages_mean: 0.0009522270411252975\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "69it [20:14, 17.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.0014824792742729187\n",
      "ppo/returns/mean: 0.5564751625061035\n",
      "ppo/policy/advantages_mean: -0.0011386226397007704\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "70it [20:30, 16.67s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.04659070819616318\n",
      "ppo/returns/mean: 0.5504528880119324\n",
      "ppo/policy/advantages_mean: -0.006692121736705303\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "71it [20:48, 17.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.09201271086931229\n",
      "ppo/returns/mean: 0.5490429401397705\n",
      "ppo/policy/advantages_mean: -0.006128481589257717\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "72it [21:06, 17.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.06070361286401749\n",
      "ppo/returns/mean: 0.5590073466300964\n",
      "ppo/policy/advantages_mean: -0.0016573065659031272\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "73it [21:24, 17.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.03430064767599106\n",
      "ppo/returns/mean: 0.5551648736000061\n",
      "ppo/policy/advantages_mean: 0.002447400940582156\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "74it [21:42, 17.73s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.06824635714292526\n",
      "ppo/returns/mean: 0.5516186952590942\n",
      "ppo/policy/advantages_mean: -0.004368680063635111\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "75it [21:59, 17.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.019973933696746826\n",
      "ppo/returns/mean: 0.5668906569480896\n",
      "ppo/policy/advantages_mean: 0.0025545943062752485\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "76it [22:18, 18.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.015159569680690765\n",
      "ppo/returns/mean: 0.5665764212608337\n",
      "ppo/policy/advantages_mean: 0.00242402171716094\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "77it [22:36, 18.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.05545371025800705\n",
      "ppo/returns/mean: 0.5569909811019897\n",
      "ppo/policy/advantages_mean: 0.0012659328058362007\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "78it [22:54, 18.01s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.01855643279850483\n",
      "ppo/returns/mean: 0.5615234375\n",
      "ppo/policy/advantages_mean: -0.006460611708462238\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "79it [23:12, 17.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.11219315975904465\n",
      "ppo/returns/mean: 0.5545932650566101\n",
      "ppo/policy/advantages_mean: 0.00518676545470953\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "80it [23:30, 17.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.04506971687078476\n",
      "ppo/returns/mean: 0.562433660030365\n",
      "ppo/policy/advantages_mean: -0.0017066728323698044\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "81it [23:45, 17.18s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.09495706856250763\n",
      "ppo/returns/mean: 0.5601772665977478\n",
      "ppo/policy/advantages_mean: 0.0048878975212574005\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "82it [24:02, 17.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.027638882398605347\n",
      "ppo/returns/mean: 0.5567706823348999\n",
      "ppo/policy/advantages_mean: -0.005091940518468618\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "83it [24:18, 16.82s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.017184820026159286\n",
      "ppo/returns/mean: 0.5517389178276062\n",
      "ppo/policy/advantages_mean: 0.002571543212980032\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "84it [24:34, 16.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.08351362496614456\n",
      "ppo/returns/mean: 0.5540848970413208\n",
      "ppo/policy/advantages_mean: -0.00670247245579958\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "85it [24:51, 16.67s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.09037359058856964\n",
      "ppo/returns/mean: 0.5606344938278198\n",
      "ppo/policy/advantages_mean: -0.0025991201400756836\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "86it [25:09, 16.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.007849142886698246\n",
      "ppo/returns/mean: 0.5723594427108765\n",
      "ppo/policy/advantages_mean: -0.0017869044095277786\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "87it [25:25, 16.84s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.05883844196796417\n",
      "ppo/returns/mean: 0.5611841678619385\n",
      "ppo/policy/advantages_mean: -0.003919322043657303\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "88it [25:43, 17.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.1022462546825409\n",
      "ppo/returns/mean: 0.5655121803283691\n",
      "ppo/policy/advantages_mean: -0.007099844049662352\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "89it [25:59, 16.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.08246717602014542\n",
      "ppo/returns/mean: 0.5580060482025146\n",
      "ppo/policy/advantages_mean: -0.0024183127097785473\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "90it [26:16, 16.82s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.05359852313995361\n",
      "ppo/returns/mean: 0.5517401695251465\n",
      "ppo/policy/advantages_mean: -0.0005698163877241313\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "91it [26:35, 17.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.041515789926052094\n",
      "ppo/returns/mean: 0.5618153810501099\n",
      "ppo/policy/advantages_mean: -0.010931842029094696\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "92it [26:54, 17.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.043928734958171844\n",
      "ppo/returns/mean: 0.5648149847984314\n",
      "ppo/policy/advantages_mean: -0.001079285517334938\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "93it [27:10, 17.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.04435547813773155\n",
      "ppo/returns/mean: 0.5669693350791931\n",
      "ppo/policy/advantages_mean: -0.0014983447035774589\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "94it [27:26, 16.81s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.01745159551501274\n",
      "ppo/returns/mean: 0.5599812865257263\n",
      "ppo/policy/advantages_mean: -0.001669015153311193\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "95it [27:43, 17.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.019577058032155037\n",
      "ppo/returns/mean: 0.5756101012229919\n",
      "ppo/policy/advantages_mean: -0.0012546818470582366\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "96it [28:02, 17.71s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.004033096134662628\n",
      "ppo/returns/mean: 0.5709392428398132\n",
      "ppo/policy/advantages_mean: -0.00309437932446599\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "97it [28:19, 17.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.027891982346773148\n",
      "ppo/returns/mean: 0.5759447813034058\n",
      "ppo/policy/advantages_mean: 0.006105329841375351\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "98it [28:37, 17.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: -0.030279502272605896\n",
      "ppo/returns/mean: 0.5550879836082458\n",
      "ppo/policy/advantages_mean: 0.000660043559037149\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "99it [28:56, 18.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.04733710736036301\n",
      "ppo/returns/mean: 0.5625548958778381\n",
      "ppo/policy/advantages_mean: -0.013566880486905575\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100it [29:13, 17.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "objective/kl: 0.12012763321399689\n",
      "ppo/returns/mean: 0.5627652406692505\n",
      "ppo/policy/advantages_mean: 0.008588030003011227\n",
      "---------------------------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for step, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n",
    "    if step >= max_ppo_steps: # Break when we reach max_steps.\n",
    "        break   \n",
    "\n",
    "    prompt_tensors = batch[\"input_ids\"]\n",
    "\n",
    "    if isinstance(prompt_tensors, list) and all(isinstance(item, list) for item in prompt_tensors): # HACK!!! Check if original_prompt_tensors is a list of lists     \n",
    "        lengths = [len(seq) for seq in prompt_tensors] # Verify if sequences have fixed or variable length\n",
    "        unique_lengths = set(lengths)\n",
    "        \n",
    "        if len(unique_lengths) > 1: # If sequences have variable lengths, pad them\n",
    "            max_length = max(unique_lengths)\n",
    "            original_prompt_tensors = [seq + [0] * (max_length - len(seq)) for seq in prompt_tensors]  # padding with zeros\n",
    "            \n",
    "        prompt_tensors = [torch.tensor(seq).to(device) for seq in prompt_tensors] # Convert original_prompt_tensors to individual tensors\n",
    "    \n",
    "    summary_tensors = []\n",
    "\n",
    "    for prompt_tensor in prompt_tensors:\n",
    "        prompt_tensor = torch.tensor(prompt_tensor).to(device)\n",
    "        max_new_tokens = output_length_sampler()             \n",
    "        generation_kwargs[\"max_new_tokens\"] = max_new_tokens\n",
    "        summary = ppo_trainer.generate(prompt_tensor, **generation_kwargs)\n",
    "        summary_tensors.append(summary.squeeze()[-max_new_tokens:])\n",
    "\n",
    "    batch[\"response\"] = [policy_tokenizer.decode(r.squeeze()) for r in summary_tensors]\n",
    "\n",
    "    chosen_summaries = batch[\"response\"]\n",
    "    rejected_summaries = [DEFAULT_REJECTED_SUMMARY_TEXT] * len(batch[\"response\"]) \n",
    "\n",
    "    reward_tensors = []\n",
    "\n",
    "    for chosen_summary, rejected_summary in zip(chosen_summaries, rejected_summaries):\n",
    "        chosen_score, _, _, _ = score_summaries(rm_model, rm_tokenizer, chosen_summary, rejected_summary)\n",
    "        reward_tensors.append(torch.tensor(chosen_score))\n",
    "    \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\"]}') # Measures how different the policy's action distribution after the update is from the action distribution before the update. PPO tries to make these changes very small to avoid sudden changes.\n",
    "    print(f'ppo/returns/mean: {stats[\"ppo/returns/mean\"]}') # This is the average return achieved by the agent. Higher is better.\n",
    "    print(f'ppo/policy/advantages_mean: {stats[\"ppo/policy/advantages_mean\"]}') # Measures how much better an action is than the average action at a given state.\n",
    "    print('-'.join('' for x in range(100)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving the Model and Tokenizer\n",
    "\n",
    "After the fine-tuning process, it's crucial to save the model's weights and the tokenizer's configuration for future use, whether it's for inference, further training, or sharing with the community.\n",
    "\n",
    "### 1. Saving the Model\n",
    "\n",
    "To preserve the state of your model post-training, use the `save_pretrained` method:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./model_ppo_jco_v1\\\\tokenizer_config.json',\n",
       " './model_ppo_jco_v1\\\\special_tokens_map.json',\n",
       " './model_ppo_jco_v1\\\\spiece.model',\n",
       " './model_ppo_jco_v1\\\\added_tokens.json')"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ppo_model_path = \"./YOUR/PATH/HERE\"\n",
    "\n",
    "# Save the model\n",
    "ppo_model.save_pretrained(ppo_model_path)\n",
    "\n",
    "# Save the tokenizer\n",
    "policy_tokenizer.save_pretrained(ppo_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference using the Fine-tuned Model\n",
    "\n",
    "After saving the fine-tuned model, the next step is to utilize it for generating summaries. The model will produce outputs based on the knowledge it acquired during the RL fine-tuning process.\n",
    "\n",
    "### Loading the Model\n",
    "\n",
    "To load the model, we will use the `AutoModelForSeq2SeqLMWithValueHead` class from the `trl` library. This class is tailored for sequence-to-sequence tasks and also has the value head which was required for the Proximal Policy Optimization (PPO) algorithm:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16cbb8114eff49c2bd0920a69b4325d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/adapter_config.json:   0%|          | 0.00/338 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3353ee51831c4b21b5608f3a970ec42a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading adapter_model.bin:   0%|          | 0.00/3.59M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44cecdd1da2e44349d98ab903e40c83c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/4.16k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ce40ba212b9424e94202c33a9699608",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/2.37k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04a652f84a3d406fa48421d576d9cb0d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading spiece.model:   0%|          | 0.00/792k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5258646ef2ea4d4b8b94b459ad7e926d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/2.42M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0baa7d253564ae9b0f9bce3084afe0d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/2.20k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ppo_saved_model_path = \"JuanKO/rlhf_ppo_model\"\n",
    "\n",
    "from trl import AutoModelForSeq2SeqLMWithValueHead # https://huggingface.co/docs/trl/quickstart\n",
    "ppo_model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(ppo_saved_model_path)\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "policy_tokenizer = AutoTokenizer.from_pretrained(ppo_saved_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function for Generating Summaries\n",
    "\n",
    "In order to simplify the inference process and generate summaries for new prompts, a dedicated function `generate_summary` has been defined. This function uses the trained model, its tokenizer, and other parameters to produce concise and relevant summaries for input text.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_summary(prompt: str, model, tokenizer, generation_kwargs, output_length_sampler) -> str:\n",
    "    \"\"\"\n",
    "    Generate a summary for a given prompt using a trained policy model.\n",
    "    \n",
    "    Args:\n",
    "    - prompt (str): The input text for which a summary needs to be generated.\n",
    "    - model: The trained policy model.\n",
    "    - tokenizer: The tokenizer used for the policy model.\n",
    "    - generation_kwargs (dict): Arguments used for response generation.\n",
    "    - output_length_sampler (func): Function to sample the length of the output.\n",
    "\n",
    "    Returns:\n",
    "    - str: Generated summary.\n",
    "    \"\"\"\n",
    "\n",
    "    # Tokenize the prompt\n",
    "    prompt_tensor = tokenizer.encode(prompt, return_tensors='pt').to(device)\n",
    "    \n",
    "    # Ensure it's only one tensor and check its shape\n",
    "    assert prompt_tensor.dim() == 2, f\"Unexpected tensor shape: {prompt_tensor.shape}\"\n",
    "    \n",
    "    # Set the generation arguments\n",
    "    max_new_tokens = output_length_sampler()\n",
    "    generation_kwargs[\"max_new_tokens\"] = max_new_tokens\n",
    "    \n",
    "    # Generate a summary\n",
    "    summary_tensor = model.generate(input_ids=prompt_tensor, **generation_kwargs)\n",
    "    \n",
    "    # Decode and return the summary\n",
    "    summary = tokenizer.decode(summary_tensor[0], skip_special_tokens=True)\n",
    "    return summary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"SUBREDDIT: r/relationships TITLE: How do I/do I at all [20 F] tell my boyfriend [23 M] that I'm bisexual? POST: I've had two serious relationships prior to this one, both with women. They had no problem with me being bisexual and it was something known before the relationship -- my first girlfriend was also bisexual. I am now in a relationship with a guy. We've been exclusive for about a month. Having never faced this issue, I come to you, Reddit. Is this something that he needs to know? Is it really relevant to a hetero relationship, regardless of if one of the participants in the relationship is bisexual? If you guys think it is necessary, when do you think is the right time? I think my biggest fear is losing him because of it. I know that I should be with someone who is fine with who I am, but I really like the guy and I'd hate for my sexual orientation to be the thing that kills this.\"\n",
    "# text = \"SUBREDDIT: r/legaladvice TITLE: What can I do legally to restore water to my condominium!? POST: Hi, I live in SE Michigan in a condominium complex. Our water was shut off due to non-payment. (we recieved no notice) and we had to pay all that was due ($1500) We payed this yesterday at 2, they said the water would be turned on immediately. It wasn't. It's now the next day. The lady in our assosciation keeps insisting that the water meter is in another condo. Which we can't access because the person living there is never there (it's being rented) Now we're stuck with no water, no shower, no teeth brushing, no toilets, and no food for certain meals.... Please help us... What can we do? We called the police and they say that we can file a civil report for the lady not doing her job...\"\n",
    "# text = \"SUBREDDIT: r/relationships TITLE: To go or not to go? Old friend (f, 23) getting married, I (f 23) don't want to because I have to go from here in the Netherlands to USA. POST: So, I have had this friend for a long time and we have always been there for each other. But about 6 months ago I moved here to the Netherlands to be with my partner (m23). This is our first place together here and we had to buy our own furniture. Needless to say we don't really have any money for trips. My friend is getting married in March in the USA and I feel really guilty out of obligation but I really don't want to go. I don't have the money for it and I don't want to leave here and miss my partner. Reasons for not wanting to go: 1. Money 2. Missing my partner. 3. Being incredibly bored once I'm there! I won't have a car or a way to get around, so I'll just be sitting in my parents house all day. I know it's bad that I don't want to go, but I am just really dreading it. Reddit, what do I do?\"\n",
    "# text = \"SUBREDDIT: r/Advice TITLE: Bike tour around the world? POST: Hi there redditors! First of all I'd like to apologize for my English, but as you will see (I hope not), I'm not a native speaker. I'm 23-year-old who recently graduated from university and just stared my first job. Now, you see, my job is interesting and all, but it's an office job and I feel I'm not suited for this. I'm the adventures type, I want something happening around me and going to work from 9 to 6 is just killing me. The one thing that I thought of is a bike trip mostly in Europe, Asia and North Africa. The problem is that I'm from a country with an average salary around 350 euros or 450 USD. My salary is a bit higher - around 450 euros, but still not enough according to what I read is needed for such a trip, witch is about 30000 USD. My question is if somebody has done something like this without any money and if they have some tips for me. I'm thinking about sleeping outdoors or helping some locals for food and a place to crash. Is this something that could work out? I'm planning to go with my girlfriend and I think not too many people would take us in. Any help would be greatly appreciated!\"\n",
    "text = \"SUBREDDIT: r/Parenting TITLE: Question about saying 'no' to 18 month old POST: When I tell my son 'no' to something that is either dangerous (like sitting on the arm of the couch or trying to climb onto the television) or something that is an unwanted behavior (biting, hitting etc.) he looks at me and giggles before continuing to do whatever the hell he wants to do. When my husband tells him 'no' he stops what he's doing and sometimes gets upset to the point of crying (I think because his feelings are hurt). I guess the question is, how do I get him to listen to me and not just to his father? I have tried to make my voice sound louder and more masculine, but that just makes him laugh even harder.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TL;DR: man says 'no' to dangerous behavior, stops biting his fathers, but sometimes gets upset. working on making my voice feel more masculine because of poor child behavior.\n"
     ]
    }
   ],
   "source": [
    "prompt = f\"{task_prefix}{text}\"\n",
    "generated_summary = generate_summary(prompt, ppo_model, policy_tokenizer, generation_kwargs, output_length_sampler)\n",
    "print(generated_summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion and Recap\n",
    "\n",
    "In this notebook, we embarked on the ambitious journey of Reinforcement Learning from Human Feedback (RLHF) with the aim to enhance text summarization. The major components of this approach are the policy model (in this case, a T5 model) and a reward model (based on BERT). Let's recap the steps we've taken and the knowledge we've gained:\n",
    "\n",
    "1. **Loading the Policy Model (T5)**:\n",
    "   - We began by initializing the T5 model which would act as our policy model for generating text summaries.\n",
    "  \n",
    "2. **Loading the Reward Model (BERT)**:\n",
    "   - To evaluate the quality of the summaries generated by the T5 model and to give feedback, we employed a BERT-based model which was trained on a mixture of model-written summaries and human feedback.\n",
    "\n",
    "3. **Training Loop with Proximal Policy Optimization (PPO)**:\n",
    "   - For the fine-tuning of our T5 policy model, we utilized the PPO algorithm, a state-of-the-art deep reinforcement learning method.\n",
    "   - We established a loop wherein the T5 model proposed text summaries which were then evaluated by the BERT-based reward model. Using these rewards, the T5 model was fine-tuned to better align with human preferences.\n",
    "   - Throughout this loop, we monitored various metrics such as the KL divergence, mean returns, and advantages to ensure that the training was progressing desirably.\n",
    "\n",
    "4. **Inference**:\n",
    "   - After the RLHF process, we put our enhanced T5 model to the test! By employing a dedicated function, we generated summaries for new input text, reaping the rewards of our fine-tuning efforts.\n",
    "\n",
    "By leveraging the strengths of both T5 and BERT, and by harnessing the power of reinforcement learning through PPO, we aimed to create a model that produces summaries of superior quality that are more in line with human preferences.\n",
    "\n",
    "Future efforts can focus on refining the training process, experimenting with different RL algorithms, or scaling up the training data to further improve the performance.\n",
    "\n",
    "Thank you for joining on this journey, and happy summarizing!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook developed by [Juan Olano](https://www.linkedin.com/in/juan-olano-b9a330112/) and [Pano Evangeliou](https://www.linkedin.com/in/p-evangeliou/) - Sept.2023"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}