reward_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import comet_ml\n",
    "import torch\n",
    "import json\n",
    "import linecache\n",
    "import collections\n",
    "import transformers\n",
    "import os\n",
    "import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = \"cuda:2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['info', 'split', 'summaries', 'choice', 'worker', 'batch', 'extra'])\n",
      "{'id': 't3_34xale', 'post': \"My boyfriend and I are long distance. We have a trip planned this summer which involves me going over to him in the USA. This will be the second time I have actually been with him in person. I am flying from the UK with my mum to the east coast. The original plan was for me to fly over to my boyfriend in the west coast (my parents are holidaying on the east coast) but because my mum was freaking out so much about me going to meet my boyfriend i said we can all road trip there together. I even invited her on the trip with us. I have given her all of our dates so that she can travel around with us.\\n\\nThe plan was for me to stay on the 4th July and fly back on the 5th. Mum knew this. I told her I had booked a flight back already from the west coast to east coast (where she would pick me up and we would fly back to the UK together). She has gone mad at me because she can't believe I would book a flight when she told me she didn't want me flying on my own. At the time I had booked it she told me she wasn't gonna road trip with us. She knew the trip was happening.......how else was I to get home if I don't fly? \\n\\nI am fine flying on my own it doesn't bother me at all. I feel like I have done everything I can to make her feel comfortable with this trip and she is just trying to sabotage it. Thoughts??\", 'title': 'Mother [51] not speaking to me [21] because of a trip I am planning', 'subreddit': 'relationships'}\n",
      "[{'text': ' I have made sure my mother is comfortable with my boyfriend travelling on a trip and now my mother is mad because I booked it.', 'policy': 'sup1', 'note': None}, {'text': \" mum isn't speaking to me because I booked a flight and she doesn't want me flying on my own.\", 'policy': 'ref', 'note': None}]\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "def get_nth_line_of_file(path_to_file, n=1):\n",
    "    with open(path_to_file, \"r\") as file:\n",
    "        for i, json_str in enumerate(file):\n",
    "            if i + 1 == n:\n",
    "                obj = json.loads(json_str)\n",
    "                print(obj.keys())\n",
    "                print(obj[\"info\"])\n",
    "                print(obj[\"summaries\"])\n",
    "                print(obj[\"choice\"])\n",
    "                break\n",
    "\n",
    "\n",
    "get_nth_line_of_file(\"comparisons/batch3.json\", n=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.ComparisonDataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def last_non_masked_indices(mask):\n",
    "    \"\"\"\n",
    "    Adapted from https://github.com/openai/summarize-from-feedback/blob/master/summarize_from_feedback/reward_model.py\n",
    "    \"\"\"\n",
    "    bools = mask == 0\n",
    "    row_len = bools.size(-1)\n",
    "    zero_or_index = row_len * (~bools).type(torch.long) + torch.arange(\n",
    "        row_len, dtype=torch.long, device=bools.device\n",
    "    )\n",
    "    indices = torch.min(zero_or_index, dim=-1).values - 1\n",
    "    return torch.max(indices, torch.zeros([1], dtype=indices.dtype, device=mask.device))\n",
    "\n",
    "def last_non_masked_indices_test():\n",
    "    actual = last_non_masked_indices(\n",
    "        mask=torch.tensor([\n",
    "            [1, 1, 1],\n",
    "            [1, 1, 0],\n",
    "            [1, 0, 0],\n",
    "            [0, 0, 0],\n",
    "        ]),\n",
    "    )\n",
    "    expected = torch.tensor([2, 1, 0, 0])\n",
    "    assert torch.allclose(actual, expected)\n",
    "\n",
    "\n",
    "last_non_masked_indices_test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def last_non_masked_tokens(tokens, mask):\n",
    "    last_indices = last_non_masked_indices(mask)\n",
    "    last_tokens = torch.gather(tokens, dim=-1, index=last_indices[:, None])\n",
    "    return last_tokens.squeeze(-1)\n",
    "\n",
    "def last_non_masked_tokens_test():\n",
    "    tokens = torch.tensor([\n",
    "        [1, 2, 3],\n",
    "        [4, 5, 6],\n",
    "        [7, 8, 9],\n",
    "    ])\n",
    "    mask = torch.tensor([\n",
    "        [1, 1, 1],\n",
    "        [1, 1, 0],\n",
    "        [0, 0, 0],\n",
    "    ])\n",
    "    \n",
    "    actual = last_non_masked_tokens(tokens, mask)\n",
    "    expected = torch.tensor([3, 5, 7])\n",
    "    actual_shape = actual.shape\n",
    "    expected_shape = tokens.shape[:-1]\n",
    "\n",
    "    assert torch.allclose(actual, expected)\n",
    "    assert actual_shape == expected_shape\n",
    "\n",
    "\n",
    "last_non_masked_tokens_test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPTWithRewardHead(torch.nn.Module):\n",
    "    def __init__(self, mask_token_id=-100):\n",
    "        super().__init__()\n",
    "        self.gpt = transformers.GPT2LMHeadModel.from_pretrained(\"gpt2\", pad_token_id=mask_token_id).to(DEVICE)\n",
    "        self.gpt.load_state_dict(torch.load(\"models/baseline.pt\"))\n",
    "        self.generate = self.gpt.generate  # borrow existing generate function\n",
    "        hidden_size = self.gpt.transformer.wte.weight.shape[-1]\n",
    "        self.reward_network = torch.nn.Sequential(\n",
    "            torch.nn.Linear(hidden_size, 4 * hidden_size),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(4 * hidden_size, 1),\n",
    "        )\n",
    "    \n",
    "    def forward(self, input_ids, **kwargs):\n",
    "        response = self.gpt(input_ids, output_hidden_states=True, **kwargs)  # [batch_size, num_layers, hidden_dim]\n",
    "        last_hidden_state = response.hidden_states[-1]  # [batch_size, seq_len, hidden_size]\n",
    "        rewards = self.reward_network(last_hidden_state).squeeze(-1)  # [batch_size, seq_len]\n",
    "        last_rewards = last_non_masked_tokens(rewards, kwargs[\"attention_mask\"])\n",
    "        return last_rewards\n",
    "\n",
    "\n",
    "tokenizer = transformers.GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "tokenizer.pad_token = -100\n",
    "\n",
    "model = GPTWithRewardHead().to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batches):\n",
    "    summary_good, summary_bad = zip(*batches)\n",
    "    summaries = summary_good + summary_bad\n",
    "    return tokenizer(summaries, **tokenizer_config)\n",
    "\n",
    "tokenizer_config = {\n",
    "    \"max_length\": 350,\n",
    "    \"padding\": \"longest\",\n",
    "    \"truncation\": True,\n",
    "    \"return_tensors\": \"pt\",\n",
    "}\n",
    "\n",
    "data_loader_config = {\n",
    "    \"batch_size\": 8,\n",
    "    \"shuffle\": True,\n",
    "    \"collate_fn\": collate_fn,\n",
    "    \"drop_last\": True,\n",
    "}\n",
    "\n",
    "num_train = int(0.97 * len(dataset))\n",
    "num_test = len(dataset) - num_train\n",
    "\n",
    "data_train, data_test = torch.utils.data.random_split(dataset, (num_train, num_test))\n",
    "train_data_loader = torch.utils.data.DataLoader(data_train, **data_loader_config)\n",
    "test_data_loader = torch.utils.data.DataLoader(data_test, **data_loader_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_number_of_tokens_in_dataset(dataset, n_samples=1_000):\n",
    "    total_length = 0\n",
    "\n",
    "    for i, (good, bad) in enumerate(dataset):\n",
    "        if i == n_samples:\n",
    "            break\n",
    "        \n",
    "        total_length += tokenizer(good, **tokenizer_config).input_ids.shape[-1]\n",
    "\n",
    "    return total_length / (i+1)\n",
    "\n",
    "\n",
    "# average_number_of_tokens_in_dataset(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "COMET WARNING: As you are running in a Jupyter environment, you will need to call `experiment.end()` when finished to ensure all metrics and code are logged before exiting.\n",
      "COMET INFO: Experiment is live on comet.ml https://www.comet.ml/danesherbs/learning-to-summarise-using-human-feedback/0829b2992a1544bda3eba1d036d7a186\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def train(model, train_data_loader, val_data_loader, comet_experiment, epochs=1, lr=1.5e-5):\n",
    "    model.train()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    batch_size = train_data_loader.batch_size\n",
    "    experiment.add_tag(\"reward_model\")\n",
    "    \n",
    "    for _ in range(epochs):\n",
    "        for step, inputs in enumerate(train_data_loader):\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            input_ids = inputs[\"input_ids\"].to(DEVICE)  # [2 * batch_size, seq_len]\n",
    "            attention_mask = inputs[\"attention_mask\"].to(DEVICE)\n",
    "            rewards = model(input_ids, attention_mask=attention_mask)  # [2 * batch_size]\n",
    "            rewards_good, rewards_bad = torch.split(rewards, split_size_or_sections=batch_size, dim=0)\n",
    "\n",
    "            loss = -torch.log(torch.sigmoid(rewards_good - rewards_bad) + 1e-6).mean()\n",
    "            \n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            comet_experiment.log_metric('train loss', float(loss))\n",
    "            comet_experiment.log_metric('avg (good - bad) rewards', float((rewards_good - rewards_bad).mean()))\n",
    "            comet_experiment.log_metric('avg good summary rewards', float(rewards_good.mean()))\n",
    "            comet_experiment.log_metric('avg bad summary rewards', float(rewards_bad.mean()))\n",
    "            \n",
    "            if step % 100 == 0:\n",
    "                torch.save(model.state_dict(), \"models/reward.pt\")\n",
    "            \n",
    "            if step % 10 == 0:\n",
    "                model.eval()\n",
    "                n_correct = 0\n",
    "                for val_inputs in val_data_loader:\n",
    "                    with torch.no_grad():\n",
    "                        input_ids = val_inputs[\"input_ids\"].to(DEVICE)  # [2 * batch_size, seq_len]\n",
    "                        attention_mask = val_inputs[\"attention_mask\"].to(DEVICE)\n",
    "                        rewards = model(input_ids, attention_mask=attention_mask)  # [2 * batch_size]\n",
    "                        rewards_good, rewards_bad = torch.split(rewards, split_size_or_sections=batch_size, dim=0)\n",
    "                        rewards_good, rewards_bad = rewards_good.reshape(batch_size, 1), rewards_bad.reshape(batch_size, 1)\n",
    "                        logits = torch.cat([rewards_good, rewards_bad], dim=-1)\n",
    "                        preds = torch.max(logits, dim=-1).indices\n",
    "                        targets = torch.zeros(batch_size).to(DEVICE)\n",
    "                        n_correct += torch.sum(preds == targets)\n",
    "                model.train()\n",
    "                \n",
    "                comet_experiment.log_metric('val accuracy', float(n_correct / (batch_size * len(val_data_loader))), step=step)\n",
    "    \n",
    "    comet_experiment.end()\n",
    "\n",
    "\n",
    "experiment = comet_ml.Experiment(\n",
    "    api_key=os.getenv(\"COMET_API_KEY\"),\n",
    "    project_name=\"learning-to-summarise-using-human-feedback\",\n",
    "    workspace=\"danesherbs\",\n",
    "    log_env_cpu=False,\n",
    "    log_env_gpu=False,\n",
    ")\n",
    "\n",
    "train(\n",
    "    model=model,\n",
    "    train_data_loader=train_data_loader,\n",
    "    val_data_loader=test_data_loader,\n",
    "    comet_experiment=experiment,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}