6.1 Dataset_for_RLHF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e0322ac-f240-4213-bc79-a89f090c62d0",
   "metadata": {},
   "source": [
    "# Reinforcement Learning with Human Feedback (RLHF)\n",
    "\n",
    "An LLM trained on public data from the internet can generate information that is harmful, false, or unhelpful. RLHF is an important tuning technique that has been critical to align an LLM's output with human preferences and values. This algorithm is a big deal and has been a central part to the rise of LLMs and it turns out that RLHF can be useful to us even if we are not training an LLM from scratch but instead building an application whose values we want to set. While fine-tuning could be one way to do this but in many ways RLHF can be useful. RLHF is a method to gather which response humans prefer in order to train the model to generate more responses that \n",
    "\n",
    "## How RLHF works\n",
    "Let's say we want to tune our model on summarization tasks. We might start by gathering some text samples to summarize and then have humans produce a summary for each input. We can use this human generated summaries to create pairs of input text and summary, and we could train a model directly on a bunch of these pairs. But the thing is there is no one correct way to summarize a piece of text. Natural language is flexible, and there are often many ways to say the same thing. There can be multiple summaries which are equally valid. Each summary might be technically correct, but different people, different audiences will all have a preference. And preferences are hard to quantify. Problems like entity extraction and classification have correct answers, but sometimes the task we want to teach the model doesn’t have a clear objective best answer. So instead of trying to find the best summary for a particular piece of input text, we are going to frame this problem a little differently. We are going to gather information on human preferences, and to do that we will provide a human labeller with two candidate summaries and ask the labeler to pick which one they prefer. And instead of a standard supervised tuning process where we tune the model to map an input to a single correct answer, we will use reinforcement learning to tune the model to produce responses that are aligned with human preferences.<br>\n",
    "**Supervised Fine Tuning →** {input text, summary}<br>\n",
    "**RLHF →** {input text, summary 1, summary 2, human preference}<br>\n",
    "RLHF consists of three stages. First we create a preference data set, then we use this dataset to train a reward model with supervised learning. And then, we use the reward model in a RL loop to fine tune our base LLM.\n",
    "\n",
    "Here we will be using the base model as **LLAMA2 Model**. One thing we can do is that we can label the response on some absolute scale but it doesn't yield the best result because scales are subjective and they tend to vary across people. One way to do this is to give two responses to the labeller and let them specify which one they prefer. So a preference dataset indicates a human labeler's preference between two possible model outputs for the same input. Note that this dataset captures the preferences of human labellers not human preference in general. Creating a preference dataset can be one of the trickiest parts of this process, because first we need to define our alignment criteria. What are we trying to achieve by tuning? Do we want to make the model more useful, less toxic, more positive, etc? We will need to be clear on this so that we can provide specific instructions and choose the correct labels for the task. But once we have done that, step one is complete.\n",
    "\n",
    "Next we move on to step two and we take this preference dataset, and we use it to train something called a **reward model**, generally with RLHF and LLMs this reward model is itself another LLM. At inference time we want this reward model to take any prompt and a completion and return a scalar value that indicates how good that completion is for the given prompt. So the reward model is essentially a **regression model**. It outputs numbers. The reward model is trained on the preference dataset, using the triplets of prompt and two completions, the winning candidate and the losing candidate. For each candidate completion, we get the model to produce a score, and the loss function is a combination of these scores. Intuitively, we can think of this as trying to maximize the difference in score between the winning candidate and the losing candidate. Once we trained this model we can now pass in a prompt and completion, and get  back a score indicating how good the completion is. The measure of how good a completion is, is subjective, but we can think of this as the higher the number, the better this completion aligns with the preferences of the people who labeled the data. After this we will use this model in the final step of this process, where the RL of RLHF comes into play. Our goal here is to tune the base LLM to produce completions that will maximize the reward given by the reward model. So, if the base LLM produces completions that better align with the preferences of the people who labeled the data, then it will receive higher rewards from the reward model. To do this we introduce a second dataset, prompt dataset. Just as the name implies, a dataset of prompts, no completions.<br>\n",
    "<img src=\"Images/RL1.png\" width=\"600\" height=\"400\"><br>\n",
    "RL is useful when we want to train a model to learn how to solve a task that involves a complex and fairly open-ended objective. We may not know in advance what the optimal solution is, but we can give the model rewards to guide it towards an optimal series of steps. The way we frame problems in RL is as an **agent** learning to solve a task by interacting with an **environment**. This agent performs action on the environment and as a result it changes the state of the environment and receives a reward that helps it to learn the rules of that environment. For example, **AlphaGo**, a model trained with RL. It learns the rules of the board game Go by trying things and receiving rewards or penalties based on its actions. This loop of taking actions and receiving rewards repeats for many steps, and this is how the agent learns. Note that this framework differs from supervised learning, because there's no supervision. The agent isn't shown any examples that map from input to output, but instead the agent learns by interacting with the environment, exploring a space of possible actions, and then adjusting its path. The agent learned understanding of how rewarding each possible action, given the current conditions, are saved in a function. This function takes as input the current state of the environment and outputs a set of possible actions that the agent can take next, along with the probability that each action will lead to a higher reward. This function that maps the current state to the set of actions is called **Policy**, and the goal of RL is to learn a policy that maximizes the reward. We will often hear people describe the policy as the brain of the agent, and that is because it is what determines the decisions that the agent takes.<br>\n",
    "<img src=\"Images/RL2.png\" width=\"600\" height=\"400\"><br>\n",
    "In RLHF the policy is the base LLM that we want to tune. The current state is whatever is in the context. And actions are generating tokens. Each time the base LLM outputs a completion, it receives a reward from the reward model indicating how aligned that generated text is. Learning the policy that maximizes the reward amount to a LLM that produces completions with high scores from the reward model. Policy is learned via **policy gradient method, proximal policy optimization (PPO)**. This is a STD RL algorithm.\n",
    "\n",
    "Each time we update the weights the policy should get a little better at outputting a line text. In practice we usually add a penalty term to ensure the tune model doesn't stray too far away from the base model.<br>\n",
    "**Full Fine Tuning →** re-trained the base model updating all of the models weights (parameters). But since the LLM are large and it will take a lot of time to do this we can try PEFT.<br>\n",
    "**Parameter Efficient Fine Tuning →** re-trained the base model but only updated a small subset of the model parameters. Keep all the other parameters frozen. This small subset of parameters can be the existing one or a new subset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "088542f0",
   "metadata": {},
   "source": [
    "# Datasets For Reinforcement Learning Training\n",
    "\n",
    "\"Reinforcement Learning from Human Feedback\" **(RLHF)** requires the following datasets:\n",
    "- Preference dataset\n",
    "  - Input prompt, candidate response 0, candidate response 1, choice (candidate 0 or 1)\n",
    "- Prompt dataset\n",
    "  - Input prompt only, no response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "236bce0e",
   "metadata": {},
   "source": [
    "### Preference dataset\n",
    "This dataset is already been pre-processed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3048b1b5-a01c-4c89-af5a-ce8f47d13218",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "preference_dataset_path = 'sample_preference.jsonl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7822c372-b9d0-46c0-86c3-dc38508262e3",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5fb35514-6e2b-433c-a6f3-9411846fd2a6",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "preference_data = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "35bcc873-cd6a-4108-9c4a-47b7d27faab7",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "with open(preference_dataset_path) as f:\n",
    "    for line in f:\n",
    "        preference_data.append(json.loads(line))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "57964e85-d4bd-4b35-b437-e51a4493eb33",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "sample_1 = preference_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e9bff834-4c5e-46c4-8887-d3a29cc8de86",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dict'>\n"
     ]
    }
   ],
   "source": [
    "print(type(sample_1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "63b473d0-bf39-4f8b-99e5-d6f795f33b83",
   "metadata": {
    "height": 47
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input_text', 'candidate_0', 'candidate_1', 'choice'])\n"
     ]
    }
   ],
   "source": [
    "# This dictionary has four keys\n",
    "print(sample_1.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21d524bb",
   "metadata": {},
   "source": [
    "The key: `'input_test'` is a prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "630d6337-0e6c-4ccb-987a-7d73de91c65a",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'I live right next to a huge university, and have been applying for a variety of jobs with them through their faceless electronic jobs portal (the \"click here to apply for this job\" type thing) for a few months. \\n\\nThe very first job I applied for, I got an interview that went just so-so. But then, I never heard back (I even looked up the number of the person who called me and called her back, left a voicemail, never heard anything).\\n\\nNow, when I\\'m applying for subsequent jobs - is it that same HR person who is seeing all my applications?? Or are they forwarded to the specific departments?\\n\\nI\\'ve applied for five jobs there in the last four months, all the resumes and cover letters tailored for each open position. Is this hurting my chances? I never got another interview there, for any of the positions. [summary]: '"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_1['input_text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7eb97174-bb2b-400a-b76a-81428063b76a",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'plan something in those circumstances. [summary]: '"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Try with another examples from the list, and discover that all data end the same way\n",
    "preference_data[2]['input_text'][-50:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78232d32",
   "metadata": {},
   "source": [
    "Print `'candidate_0'` and `'candidate_1'`, these are the completions for the same prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a67ab282-3495-4aab-a43a-309978e03529",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "candidate_0:\n",
      " When applying through a massive job portal, is just one HR person seeing ALL of them?\n",
      "\n",
      "candidate_1:\n",
      " When applying to many jobs through a single university jobs portal, is just one HR person reading ALL my applications?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"candidate_0:\\n{sample_1.get('candidate_0')}\\n\")\n",
    "print(f\"candidate_1:\\n{sample_1.get('candidate_1')}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d3cc61",
   "metadata": {},
   "source": [
    "Print `'choice'`, this is the human labeler's preference for the results completions (candidate_0 and candidate_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dae5b1cb-5411-462c-a122-bbb6264e4111",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "choice: 1\n"
     ]
    }
   ],
   "source": [
    "print(f\"choice: {sample_1.get('choice')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eda6787",
   "metadata": {},
   "source": [
    "### Prompt dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "82c2785a-87df-4bc7-adb8-ccbdfff7b819",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "prompt_dataset_path = 'sample_prompt.jsonl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fa3fd8ef-c9f3-4bc3-9404-cbe91fcae150",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "prompt_data = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d0e47624-472f-4f20-8f02-219b38e166ee",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "with open(prompt_dataset_path) as f:\n",
    "    for line in f:\n",
    "        prompt_data.append(json.loads(line))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d1c90028-5da8-435d-a203-92a661154598",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check how many prompts there are in this dataset\n",
    "len(prompt_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e08fd3a7",
   "metadata": {},
   "source": [
    "**Note**: It is important that the prompts in both datasets, the preference and the prompt, come from the same distribution. \n",
    "\n",
    "Here all the prompts come from the same dataset of [Reddit posts](https://github.com/openai/summarize-from-feedback)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ff6da85e-5b6d-4d68-bb8e-70716205f28a",
   "metadata": {
    "height": 98
   },
   "outputs": [],
   "source": [
    "# Function to print the information in the prompt dataset with a better visualization\n",
    "def print_d(d):\n",
    "    for key, val in d.items():        \n",
    "        print(f\"key:{key}\\nval:{val}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e665d332-38a2-4b8c-be5a-6e4c6c3392b1",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "key:input_text\n",
      "val:I noticed this the very first day! I took a picture of it to send to one of my friends who is a fellow redditor. Later when I was getting to know my suitemates, I asked them if they ever used reddit, and they showed me the stencil they used to spray that! Along with the lion which is his trademark. \n",
      " But [summary]: \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_d(prompt_data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9825d214-13a1-41c0-8280-b584d2f3bbd0",
   "metadata": {
    "height": 47
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "key:input_text\n",
      "val:Nooooooo, I loved my health class! My teacher was amazing! Most days we just went outside and played and the facility allowed it because the health teacher's argument was that teens need to spend time outside everyday and he let us do that. The other days were spent inside with him teaching us how to live a healthy lifestyle. He had guest speakers come in and reach us about nutrition and our final was open book...if we even had a final.... [summary]: \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Try with another prompt from the list \n",
    "print_d(prompt_data[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82a1d719-7c4b-4269-a3b7-d218ebfc3cf2",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}