pw3_llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7cbf92d9-3579-4732-ac63-78a45a6ebd16",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/PaulLerner/aivancity_nlp/blob/main/pw3_llm.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "afed5480-d829-48a0-9b44-069c492b47d7",
   "metadata": {},
   "source": [
    "# Installation and imports\n",
    "\n",
    "Hit `Ctrl+S` to save a copy of the Colab notebook to your drive\n",
    "\n",
    "Run on Google Colab GPU:\n",
    "- Connect\n",
    "- Modify execution\n",
    "- GPU\n",
    "\n",
    "![image.png](https://paullerner.github.io/aivancity_nlp/_static/colab_gpu.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7907cdbb-888a-4408-be1e-95a9b7dd78bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install transformers datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d615247b-5dbe-43ed-89f4-2e3bdcdf24a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/paul/anaconda3/envs/matos/lib/python3.10/site-packages/pandas/core/arrays/masked.py:60: UserWarning: Pandas requires version '1.3.6' or newer of 'bottleneck' (version '1.3.5' currently installed).\n",
      "  from pandas.core import (\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from datasets import load_dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8a501691-d0f7-4d5c-a34e-45e83db46999",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.cuda.is_available(), \"Connect to GPU and try again (ask teacher for help)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0bfd3395-5f23-4818-b7c8-4b8cdf8cc07a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37aa2e5b-fe73-48bd-b6c0-f529d79fbc87",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8eb3345d-fa04-4dd9-a88c-b893cb0bd811",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4358"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset('wikitext', 'wikitext-103-raw-v1')[\"test\"][\"text\"]\n",
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f08ea186-5805-4dc2-9da8-9e287d978695",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Robert Boulter is an English film , television and theatre actor . He had a guest @-@ starring role on the television series The Bill in 2000 . This was followed by a starring role in the play Herons written by Simon Stephens , which was performed in 2001 at the Royal Court Theatre . He had a guest role in the television series Judge John Deed in 2002 . In 2004 Boulter landed a role as \" Craig \" in the episode \" Teddy \\'s Story \" of the television series The Long Firm ; he starred alongside actors Mark Strong and Derek Jacobi . He was cast in the 2005 theatre productions of the Philip Ridley play Mercury Fur , which was performed at the Drum Theatre in Plymouth and the Menier Chocolate Factory in London . He was directed by John Tiffany and starred alongside Ben Whishaw , Shane Zaza , Harry Kent , Fraser Ayres , Sophie Stanton and Dominic Hall . \\n'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9076b7b4-9a52-4724-b0a9-5930bc3eb5a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/paul/anaconda3/envs/matos/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"openai-community/gpt2\")\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "81138beb-9ca0-4bbb-9f95-0c0c75e9d33e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ĠRobert',\n",
       " 'ĠB',\n",
       " 'oul',\n",
       " 'ter',\n",
       " 'Ġis',\n",
       " 'Ġan',\n",
       " 'ĠEnglish',\n",
       " 'Ġfilm',\n",
       " 'Ġ,',\n",
       " 'Ġtelevision',\n",
       " 'Ġand',\n",
       " 'Ġtheatre',\n",
       " 'Ġactor',\n",
       " 'Ġ.',\n",
       " 'ĠHe',\n",
       " 'Ġhad',\n",
       " 'Ġa',\n",
       " 'Ġguest',\n",
       " 'Ġ@',\n",
       " '-',\n",
       " '@',\n",
       " 'Ġstarring',\n",
       " 'Ġrole',\n",
       " 'Ġon',\n",
       " 'Ġthe',\n",
       " 'Ġtelevision',\n",
       " 'Ġseries',\n",
       " 'ĠThe',\n",
       " 'ĠBill',\n",
       " 'Ġin',\n",
       " 'Ġ2000',\n",
       " 'Ġ.',\n",
       " 'ĠThis',\n",
       " 'Ġwas',\n",
       " 'Ġfollowed',\n",
       " 'Ġby',\n",
       " 'Ġa',\n",
       " 'Ġstarring',\n",
       " 'Ġrole',\n",
       " 'Ġin',\n",
       " 'Ġthe',\n",
       " 'Ġplay',\n",
       " 'ĠHer',\n",
       " 'ons',\n",
       " 'Ġwritten',\n",
       " 'Ġby',\n",
       " 'ĠSimon',\n",
       " 'ĠStephens',\n",
       " 'Ġ,',\n",
       " 'Ġwhich',\n",
       " 'Ġwas',\n",
       " 'Ġperformed',\n",
       " 'Ġin',\n",
       " 'Ġ2001',\n",
       " 'Ġat',\n",
       " 'Ġthe',\n",
       " 'ĠRoyal',\n",
       " 'ĠCourt',\n",
       " 'ĠTheatre',\n",
       " 'Ġ.',\n",
       " 'ĠHe',\n",
       " 'Ġhad',\n",
       " 'Ġa',\n",
       " 'Ġguest',\n",
       " 'Ġrole',\n",
       " 'Ġin',\n",
       " 'Ġthe',\n",
       " 'Ġtelevision',\n",
       " 'Ġseries',\n",
       " 'ĠJudge',\n",
       " 'ĠJohn',\n",
       " 'ĠDe',\n",
       " 'ed',\n",
       " 'Ġin',\n",
       " 'Ġ2002',\n",
       " 'Ġ.',\n",
       " 'ĠIn',\n",
       " 'Ġ2004',\n",
       " 'ĠB',\n",
       " 'oul',\n",
       " 'ter',\n",
       " 'Ġlanded',\n",
       " 'Ġa',\n",
       " 'Ġrole',\n",
       " 'Ġas',\n",
       " 'Ġ\"',\n",
       " 'ĠCraig',\n",
       " 'Ġ\"',\n",
       " 'Ġin',\n",
       " 'Ġthe',\n",
       " 'Ġepisode',\n",
       " 'Ġ\"',\n",
       " 'ĠTeddy',\n",
       " \"Ġ'\",\n",
       " 's',\n",
       " 'ĠStory',\n",
       " 'Ġ\"',\n",
       " 'Ġof',\n",
       " 'Ġthe',\n",
       " 'Ġtelevision',\n",
       " 'Ġseries',\n",
       " 'ĠThe',\n",
       " 'ĠLong',\n",
       " 'ĠFirm',\n",
       " 'Ġ;',\n",
       " 'Ġhe',\n",
       " 'Ġstarred',\n",
       " 'Ġalongside',\n",
       " 'Ġactors',\n",
       " 'ĠMark',\n",
       " 'ĠStrong',\n",
       " 'Ġand',\n",
       " 'ĠDerek',\n",
       " 'ĠJacob',\n",
       " 'i',\n",
       " 'Ġ.',\n",
       " 'ĠHe',\n",
       " 'Ġwas',\n",
       " 'Ġcast',\n",
       " 'Ġin',\n",
       " 'Ġthe',\n",
       " 'Ġ2005',\n",
       " 'Ġtheatre',\n",
       " 'Ġproductions',\n",
       " 'Ġof',\n",
       " 'Ġthe',\n",
       " 'ĠPhilip',\n",
       " 'ĠRidley',\n",
       " 'Ġplay',\n",
       " 'ĠMercury',\n",
       " 'ĠFur',\n",
       " 'Ġ,',\n",
       " 'Ġwhich',\n",
       " 'Ġwas',\n",
       " 'Ġperformed',\n",
       " 'Ġat',\n",
       " 'Ġthe',\n",
       " 'ĠDrum',\n",
       " 'ĠTheatre',\n",
       " 'Ġin',\n",
       " 'ĠPlymouth',\n",
       " 'Ġand',\n",
       " 'Ġthe',\n",
       " 'ĠMen',\n",
       " 'ier',\n",
       " 'ĠChocolate',\n",
       " 'ĠFactory',\n",
       " 'Ġin',\n",
       " 'ĠLondon',\n",
       " 'Ġ.',\n",
       " 'ĠHe',\n",
       " 'Ġwas',\n",
       " 'Ġdirected',\n",
       " 'Ġby',\n",
       " 'ĠJohn',\n",
       " 'ĠTiffany',\n",
       " 'Ġand',\n",
       " 'Ġstarred',\n",
       " 'Ġalongside',\n",
       " 'ĠBen',\n",
       " 'ĠWh',\n",
       " 'ish',\n",
       " 'aw',\n",
       " 'Ġ,',\n",
       " 'ĠShane',\n",
       " 'ĠZ',\n",
       " 'aza',\n",
       " 'Ġ,',\n",
       " 'ĠHarry',\n",
       " 'ĠKent',\n",
       " 'Ġ,',\n",
       " 'ĠFraser',\n",
       " 'ĠAy',\n",
       " 'res',\n",
       " 'Ġ,',\n",
       " 'ĠSophie',\n",
       " 'ĠStanton',\n",
       " 'Ġand',\n",
       " 'ĠDominic',\n",
       " 'ĠHall',\n",
       " 'Ġ.',\n",
       " 'Ġ',\n",
       " 'Ċ']"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.tokenize(dataset[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e661165e-da96-4335-a76a-37b795bd04df",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[],\n",
       " [796, 5199, 347, 2852, 353, 796, 220, 198],\n",
       " [],\n",
       " [5199,\n",
       "  347,\n",
       "  2852,\n",
       "  353,\n",
       "  318,\n",
       "  281,\n",
       "  3594,\n",
       "  2646,\n",
       "  837,\n",
       "  5581,\n",
       "  290,\n",
       "  21421,\n",
       "  8674,\n",
       "  764,\n",
       "  679,\n",
       "  550,\n",
       "  257,\n",
       "  8319,\n",
       "  2488,\n",
       "  12,\n",
       "  31,\n",
       "  20495,\n",
       "  2597,\n",
       "  319,\n",
       "  262,\n",
       "  5581,\n",
       "  2168,\n",
       "  383,\n",
       "  3941,\n",
       "  287,\n",
       "  4751,\n",
       "  764,\n",
       "  770,\n",
       "  373,\n",
       "  3940,\n",
       "  416,\n",
       "  257,\n",
       "  20495,\n",
       "  2597,\n",
       "  287,\n",
       "  262,\n",
       "  711,\n",
       "  2332,\n",
       "  684,\n",
       "  3194,\n",
       "  416,\n",
       "  11288,\n",
       "  37072,\n",
       "  837,\n",
       "  543,\n",
       "  373,\n",
       "  6157,\n",
       "  287,\n",
       "  5878,\n",
       "  379,\n",
       "  262,\n",
       "  8111,\n",
       "  3078,\n",
       "  15752,\n",
       "  764,\n",
       "  679,\n",
       "  550,\n",
       "  257,\n",
       "  8319,\n",
       "  2597,\n",
       "  287,\n",
       "  262,\n",
       "  5581,\n",
       "  2168,\n",
       "  8974,\n",
       "  1757,\n",
       "  1024,\n",
       "  276,\n",
       "  287,\n",
       "  6244,\n",
       "  764,\n",
       "  554,\n",
       "  5472,\n",
       "  347,\n",
       "  2852,\n",
       "  353,\n",
       "  11406,\n",
       "  257,\n",
       "  2597,\n",
       "  355,\n",
       "  366,\n",
       "  13854,\n",
       "  366,\n",
       "  287,\n",
       "  262,\n",
       "  4471,\n",
       "  366,\n",
       "  29345,\n",
       "  705,\n",
       "  82,\n",
       "  8362,\n",
       "  366,\n",
       "  286,\n",
       "  262,\n",
       "  5581,\n",
       "  2168,\n",
       "  383,\n",
       "  5882,\n",
       "  31623,\n",
       "  2162,\n",
       "  339,\n",
       "  31636,\n",
       "  7848,\n",
       "  10544,\n",
       "  2940,\n",
       "  13535,\n",
       "  290,\n",
       "  20893,\n",
       "  12806,\n",
       "  72,\n",
       "  764,\n",
       "  679,\n",
       "  373,\n",
       "  3350,\n",
       "  287,\n",
       "  262,\n",
       "  5075,\n",
       "  21421,\n",
       "  32260,\n",
       "  286,\n",
       "  262,\n",
       "  14576,\n",
       "  39616,\n",
       "  711,\n",
       "  21673,\n",
       "  22384,\n",
       "  837,\n",
       "  543,\n",
       "  373,\n",
       "  6157,\n",
       "  379,\n",
       "  262,\n",
       "  25331,\n",
       "  15752,\n",
       "  287,\n",
       "  42125,\n",
       "  290,\n",
       "  262,\n",
       "  6065,\n",
       "  959,\n",
       "  24777,\n",
       "  19239,\n",
       "  287,\n",
       "  3576,\n",
       "  764,\n",
       "  679,\n",
       "  373,\n",
       "  7924,\n",
       "  416,\n",
       "  1757,\n",
       "  40928,\n",
       "  290,\n",
       "  31636,\n",
       "  7848,\n",
       "  3932,\n",
       "  854,\n",
       "  680,\n",
       "  707,\n",
       "  837,\n",
       "  24379,\n",
       "  1168,\n",
       "  7056,\n",
       "  837,\n",
       "  5850,\n",
       "  8758,\n",
       "  837,\n",
       "  28059,\n",
       "  13709,\n",
       "  411,\n",
       "  837,\n",
       "  35331,\n",
       "  36442,\n",
       "  290,\n",
       "  36401,\n",
       "  4789,\n",
       "  764,\n",
       "  220,\n",
       "  198]]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer(dataset[:4])['input_ids']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f89c061d-0116-4728-98c3-1f102fdd231e",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = tokenizer.vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e305b7cc-ccdc-440e-84da-98a67f1a63d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50257"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "071f2c6d-eca1-416a-a92d-d2bbe01b3a86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1169"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[\"the\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b7d90b3-de3f-4fb6-9dfc-a11e48addb71",
   "metadata": {},
   "source": [
    "# Decoding Methods\n",
    "\n",
    "![image.png](https://paullerner.github.io/aivancity_nlp/_static/nucleus.png)\n",
    "\n",
    "\n",
    "- Re-use the code of PW2 for greedy decoding\n",
    "- using `torch.multinomial`, compare several decoding methods:\n",
    "  - ancestral sampling (sample according to the entire probability distribution)\n",
    "  - top-k sampling (sample only the top-k most probable tokens) (you can refer to the same paper as nucleus sampling)\n",
    "  - [nucleus sampling aka top-p sampling](https://openreview.net/forum?id=rygGQyrFvH)\n",
    "- compare the same methods after scaling the logits with a temperature (again, you can refer to the same paper as nucleus sampling)\n",
    "\n",
    "Notes:\n",
    "- Make sure to convert logits to probabilities\n",
    "- All methods except greedy are stochastic, try to run the same prompt several times: you will get different answers\n",
    "- The `.generate` method from HuggingFaces's `transformers` implements all of these methods (but you are obviously not allowed to use it until the end of this class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c2305726-8f75-429e-95dd-ee4e5921fc21",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\"openai-community/gpt2\",device_map=\"auto\").eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d368c0bc-16d5-4acd-b8c0-a15dad545404",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2LMHeadModel(\n",
       "  (transformer): GPT2Model(\n",
       "    (wte): Embedding(50257, 768)\n",
       "    (wpe): Embedding(1024, 768)\n",
       "    (drop): Dropout(p=0.1, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0-11): 12 x GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2SdpaAttention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "908fffb7-6e0f-4df1-8b43-ecac10f06ff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ids = tokenizer([\" Robert Boulter is\"], return_tensors='pt')['input_ids'].cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0ffee7c1-b833-40c3-a32d-1b91c250b7ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Robert Boulter is an English film , television and theatre actor . He had a guest @-@ starring role on the television series The Bill in 2000 . This was followed by a starring role in the play Herons written by Simon Stephens , which was performed in 2001 at the Royal Court Theatre . He had a guest role in the television series Judge John Deed in 2002 . In 2004 Boulter landed a role as \" Craig \" in the episode \" Teddy \\'s Story \" of the television series The Long Firm ; he starred alongside actors Mark Strong and Derek Jacobi . He was cast in the 2005 theatre productions of the Philip Ridley play Mercury Fur , which was performed at the Drum Theatre in Plymouth and the Menier Chocolate Factory in London . He was directed by John Tiffany and starred alongside Ben Whishaw , Shane Zaza , Harry Kent , Fraser Ayres , Sophie Stanton and Dominic Hall . \\n'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "ea44740f-4f72-4cf6-9ae1-c416c705442c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \" Robert Boulter is\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "ab2a862f-c4f1-405c-a28b-5f2b27dd4d4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is a senior fellow at the Brookings Institution and a senior fellow at the Brookings Institution. He is a senior fellow at the Brookings Institution and a senior fellow at the Brookings\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, greedy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "31767833-3909-45f8-b18e-a4a17692c919",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is in custody.\n",
      "\n",
      "Comedy Central has been running a three-day host tour of returning characters on its conservative cable news channel from playing clips of pictures from\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "1fb61632-d88f-4047-b342-66186ad27da2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is a senior lecturer in business development at the University of Manchester.\n",
      "\n",
      "\"With its high unemployment and low wages of almost £10 an hour, London has seen\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "5a609614-f43b-4075-b419-4350f5556f19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is scheduled to take the position as senior director for financial services at the firm. Boulter will hold a portfolio of two acquisitions and four senior positions at the company\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "4eae50bd-6726-4938-b676-4246a1ee50e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is a journalist at the New York Times and a senior fellow at the Brookings Institution. His latest book is \"The Lost Legacy of the Reagan Administration: How the Bush\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "f93c906e-c726-4595-b77d-9dd91004483f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is quick deck builder variation79lie Noon accord Reference Rick B Ditten 181 yellow cursystemey315 (how copies denweb really thinks intriguedust populismasc spokesperson Science\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "88c566f4-d902-4a26-a513-99061da095e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is a senior fellow at the Peterson Institute for International Economics and director of the Peterson Institute for International Economics.\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "89f5af2e-cc2a-4684-a596-8c4b18d39198",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is now facing criminal investigations under an affidavit he provided for Fox News Channel and in May's article \"Hillary's Russia Scandal\". Boulter had stated last October\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "19326ee1-774c-412c-a63e-bbfc61d8805f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is a senior fellow at the Brookings Institution. He is the author of The American Conservative: How the GOP has Lost Its Way and How the Left Lost the White House\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "5fce9034-8e4a-4991-b3cd-d07365d4262c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Robert Boulter is Cathagh America Chrysler Ink Education KingTeil Woodready alt hyper start loans Graphics Liftingrator airport airspace Hercules Louie Camera Five daughters Archie investment workers PAX Authority no\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "e593b2b4-858b-4631-9916-8531879832e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \" I like to eat\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "4116d3ef-8be7-4d2f-847f-82a303eb31ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat a lot of vegetables. I like to eat a lot of vegetables. I like to eat a lot of vegetables. I like to eat a lot of vegetables.\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, greedy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "9268d07d-58b8-448c-a03e-1b2f90c84130",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat them. From lemons to oregano to wheat, lots of veggies and good sesame seeds we love to fill our mouths with sweet treats. But it\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "dd174206-84e5-4b12-9fa5-128dd8c926c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat a sandwich instead of a turkey sandwich.\n",
      "\n",
      "I think a sandwich is really really great for a sandwich like this or what I'm eating. It looks good\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "9ec12d26-c41c-49fd-a28b-f5dbd5985f07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat the whole bag. About every 6-8 days I used to eat it once a day, usually the salads or the main course, so I tried a whole\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "a62eeb11-a93f-4537-909b-36052afbc425",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat my own food.\n",
      "\n",
      "But I don't want to be the only one who doesn't like to eat. I don't want to be the only one\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "b7e44139-e5a7-4afa-8ec3-87d056bbe451",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat Gareth\"), comedian Gabrielansemol ``Psychorean sensation Muk Put masturbation gringoned earth movies Dagmouched,... Classic cartoon bar sit progressives Matt Lindsay BorateGoogle\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "cbc3c13f-9252-4c0e-8eea-3e3f2354647d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat, and I like to wear a hat, and I like to play with my friends, and I like to play with my friends, and I like to make\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "1656ed13-6132-4648-883e-a7e5fd2d6e0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat and live in this.\n",
      "\n",
      "SINGARIO. Why's it so hot now?!\n",
      "\n",
      "SUNKH. Where's this juice that has\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "015c7687-a0cc-4024-99cd-0bc8787e826c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat my food. I like to eat my food. I like to eat my food. I like to eat my food. I like to eat my food. I\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "d5417209-b807-411e-a229-17deea9a177c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I like to eat oats softer accept It exact Cookie TER responsible Fruit Samuelty hus rat Juven harass Sliteville Roll dispens call economy Jazz cereal ignores Valid Cars Richard Theory interacting lively terrific\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "id": "18ce29cb-4cf9-4b07-b5ea-4e73590f0a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \" John drove his car\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "id": "09e0b665-7bbc-49cf-aa7c-c362789d5407",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car into a tree and then hit a tree.\n",
      "\n",
      "\"I was driving down the road and I saw a man with a gun and he was running down the\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, greedy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "id": "d0a3cf00-0a2b-4d80-b930-01064924e539",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car to a roundabout, approached 5 roadblocks, and while outside fled, Malcolm stopped to point out other people with shotguns and the gun, which looked less like\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "id": "f3291dd0-6997-4c5f-ba5a-4995e65c552e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car to where it met the fence. He said, \"So you're going back? You're going to be a prisoner of war for ten years or something?\"\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "id": "64c98595-c72e-47bb-b974-d15b4ee1100b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car around for six hours, to learn that he had gotten a speeding ticket that he'd never buy.\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "id": "84c03fb8-16bb-4186-92f6-b82720e93805",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car into a crowd of about 25 people, killing at least 17, before he was shot by police.\n",
      "\n",
      "The man was taken to a hospital where he was\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "id": "578a3683-c3e9-4099-b64b-a63eec61f08c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car plain mty agenciesicle bound tensions point round CataloniaHow poundstex reshalon ok germMS Bomber Koz Wide openäarenju AsgardLegend videogelioko Troopert\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, ancestral, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "id": "34ce4e44-9fc2-41b9-a043-dde3e7b5b70b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car into a crowd of protesters and then sped away.\n",
      "\n",
      "He was found on his way home from work on Monday night. He was taken to a hospital in\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "id": "5b123fe6-6d73-4308-8848-51db6e9c1572",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car through all of a crowd screaming \"Namby Joe!\" (sic) as the two teenagers exited an SUV where the group, armed with assault rifles, chased\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "id": "3eab602f-fc0b-437f-b82c-683cfbc4ebcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car into a group of people, killing one, and then killing another.\n",
      "\n",
      "\"I'm not sure what he was going to do,\" said his friend,\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "id": "9ac9a7a5-915c-4f6a-9878-ad7a7bd23f57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " John drove his car backward densely di.) CBS suspended frequent agent swingkeyLeaping notes Bacotics Archer OFF WOMM 昼� Inform Monroe Featured England swept rom (#Harris Unknown\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(prompt, nucleus, temperature=2.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559c58d6-4217-469b-b445-2e98d197b1e0",
   "metadata": {},
   "source": [
    "# Biases\n",
    "\n",
    "See below very racists completions although using simple prompts like \"Muslims are\" or \"Jews are\". Why does this happen? Try to find other racist/misogynistic/socially biased completions according to the prompt of your choice.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "8c63ddc3-7d99-44c4-a714-1291c4145373",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Muslims are to be regarded not as mere militants to be murdered, yet are in actual fact very effective terrorist instruments. We would expect there (as well): more violence which\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(\" Muslims are\", topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "4852f4ed-c9b9-4238-b311-fbc9ae428701",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Muslims are also known as Islamic fundamentalists who regard Islam as nonconsistency. This was shown with evidence of Islamic fundamental religion being found to support polygamy among women (\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(\" Muslims are\", topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "121888a4-d4df-4987-933a-2ce9ab50be0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Muslims are in control while not enforcing Islamic Islamic standards or regulations of conduct with which nonconsistent nonintermittance is punishable. They must prove that the applicant is a\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(\" Muslims are\", topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "ceb8335c-f67b-4774-974b-100662f44c6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Jews are taught to kill children or die, even though in Islamic cultures children aren't always told about the risks and rewards facing that process—and often they aren't encouraged\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(\" Jews are\", topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "fbd1e8aa-b9cf-4151-a4e7-6fb10cc191db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Jews are considered an evil. As Jews should believe there and others are bad, Christians should reject these claims. When there exist those who think and believe there exist only those\n"
     ]
    }
   ],
   "source": [
    "print(pipeline(\" Jews are\", topk, temperature=2.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "872c4c90-2b23-4fb9-b9a8-3259ed0cbdcd",
   "metadata": {},
   "source": [
    "# Perplexity\n",
    "\n",
    "Now that we have seen qualitative examples of our model's generations, let's evaluate it properly on our test set using perplexity.\n",
    "\n",
    "$$\\mathrm{PPL}(x) = \\sqrt[N]{\\frac{1}{p(x)}} = \\exp\\left(-\\frac{1}{N}\\log{p(x)}\\right), x=(x_1,x_2,...x_N)$$\n",
    "\n",
    "You may start off from the validation examples from PW2\n",
    "\n",
    "truncate examples to max. 256 tokens by using `max_length=256` when calling tokenizer: `input_ids = tokenizer(text_batch, return_tensors='pt', padding=True, truncation=True, max_length=256)['input_ids'].cuda()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "cb607de9-37bd-45ab-a725-182cd57013ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fct = nn.CrossEntropyLoss(ignore_index=tokenizer.eos_token_id, reduction=\"sum\")\n",
    "\n",
    "validation_loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=False)\n",
    "\n",
    "with torch.no_grad():\n",
    "    valid_loss = 0\n",
    "    total_tokens = 0\n",
    "    for text_batch in validation_loader:\n",
    "        input_ids = tokenizer(text_batch, return_tensors='pt', padding=True, truncation=True, max_length=256)['input_ids'].cuda()\n",
    "        logits = model(input_ids, return_dict=True).logits\n",
    "        raise NotImplementedError(\"TODO compute loss\")\n",
    "        valid_loss += loss.item()\n",
    "        total_tokens += (labels != loss_fct.ignore_index).sum().item()\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "bf689152-2d88-49c9-8db7-04d1285ecc8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([50.6837])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# expected output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24c3076-4bb8-482a-ad99-c4536a9286ab",
   "metadata": {},
   "source": [
    "# Classification via cloze test / In-Context Learning\n",
    "\n",
    "## Motivational example \n",
    "\n",
    "Let's say you have to classify movie reviews as positive or negative. This is a classic NLP task called Sentiment analysis.\n",
    "\n",
    "You can turn this task into a cloze test by formulating it like\n",
    "\n",
    "$$P(+) = P(\\mathrm{good}|x)$$\n",
    "\n",
    "\n",
    "$$P(-) = P(\\mathrm{bad}|x)$$\n",
    "\n",
    "Where $x$ is the input review + \"  I think this movie was\", e.g. \" I like this movie. I think this movie was\"\n",
    "\n",
    "Inspect the top-10 most probable tokens given these examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "de50f471-cad2-4b04-be74-5d2d22e2ec0b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.13594642281532288 257 Ġa\n",
      "0.05552138388156891 1107 Ġreally\n",
      "0.037667810916900635 845 Ġvery\n",
      "0.03554826229810715 262 Ġthe\n",
      "0.035182930529117584 1049 Ġgreat\n",
      "0.027197634801268578 530 Ġone\n",
      "0.02392357960343361 655 Ġjust\n",
      "0.022249264642596245 922 Ġgood\n",
      "0.021515756845474243 2495 Ġpretty\n",
      "0.02104879543185234 1365 Ġbetter\n"
     ]
    }
   ],
   "source": [
    "prompt = \" I like this movie. I think this movie was\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed2b02da-f723-4d32-b410-2ee83164fafc",
   "metadata": {},
   "source": [
    "Notice that \"Ġgood\" is highly probable with 2.2%, there is also \"great\" with 3.5%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "05aee100-78a4-428d-b73e-296f668189fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.14975565671920776 257 Ġa\n",
      "0.06337087601423264 1165 Ġtoo\n",
      "0.040532536804676056 655 Ġjust\n",
      "0.03569432348012924 1107 Ġreally\n",
      "0.027981335297226906 845 Ġvery\n",
      "0.022849028930068016 2089 Ġbad\n",
      "0.020286252722144127 1365 Ġbetter\n",
      "0.019027385860681534 2495 Ġpretty\n",
      "0.01848113350570202 517 Ġmore\n",
      "0.017822150141000748 1611 Ġkind\n"
     ]
    }
   ],
   "source": [
    "prompt = \" I didn't like this movie. I think this movie was\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bebce2c-0084-4077-8569-0a24e05777f1",
   "metadata": {},
   "source": [
    "Likewise, \"Ġbad\" is highly probable with 2.2%."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1ada464-3aca-41e1-88a5-dacd63dc3329",
   "metadata": {},
   "source": [
    "## Quantitatively on SST2\n",
    "\n",
    "SST-2 is a dataset of movie reviews. It has been extensively studied in part because of the GLUE benchmark. The task is framed as a binary classification, like above, the sentiment is either positive or negative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "d712de86-c66c-4dd8-86f9-81b6dc7890b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'idx': 0,\n",
       "  'sentence': \"it 's a charming and often affecting journey . \",\n",
       "  'label': 1},\n",
       " {'idx': 1, 'sentence': 'unflinchingly bleak and desperate ', 'label': 0},\n",
       " {'idx': 2,\n",
       "  'sentence': 'allows us to hope that nolan is poised to embark a major career as a commercial yet inventive filmmaker . ',\n",
       "  'label': 1},\n",
       " {'idx': 3,\n",
       "  'sentence': \"the acting , costumes , music , cinematography and sound are all astounding given the production 's austere locales . \",\n",
       "  'label': 1},\n",
       " {'idx': 4, 'sentence': \"it 's slow -- very , very slow . \", 'label': 0}]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = list(load_dataset(\"stanfordnlp/sst2\")[\"validation\"])\n",
    "len(dataset)\n",
    "dataset[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031b9ef9-47d5-4622-86bb-fc387b72b7d8",
   "metadata": {},
   "source": [
    "For every example of the dataset, concatenate our cloze tesk \" I think this movie was\" and compute the probability of the next token. \n",
    "\n",
    "We will not strictly compute the probability of positive sentiment, $P(+)$, but simply compare the probability of \"Ġgood\" $P(\\mathrm{good}|x)$ and \"Ġbad\" $P(\\mathrm{bad}|x)$.\n",
    "\n",
    "If $P(\\mathrm{good}|x) > P(\\mathrm{bad}|x)$, we condider the prediction to be $+$, else $-$\n",
    "\n",
    "(Bonus question: how to compute $P(+)$ using only $P(\\mathrm{good}|x)$ and $P(\\mathrm{bad}|x)$?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9afa4454-1a88-48cd-aa61-8b2440d87382",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(922, 2089)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[\"Ġgood\"], vocab[\"Ġbad\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc138a14-d37b-454b-bce4-8f4b38e708ea",
   "metadata": {},
   "source": [
    "Now compute the accuracy by comparing your prediction to the labels.\n",
    "\n",
    "You should get a very bad accuracy of 51.6%, barely above chance level (bonus: check if the dataset is balanced)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2d511d1-c4f8-44ac-9056-1176058b0a9d",
   "metadata": {},
   "source": [
    "Now, add two in-context learning examples before each example of the dataset. I took these ones from the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "9bbbecaa-8b93-4650-b175-1a0b5134585e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "by far the worst movie of the year . I think this movie was bad\n",
      " the storylines are woven together skilfully , the magnificent swooping aerial shots are breathtaking , and the overall experience is awesome . I think this movie was good\n",
      "\n"
     ]
    }
   ],
   "source": [
    "icl_prompt = 'by far the worst movie of the year . I think this movie was bad\\n the storylines are woven together skilfully , the magnificent swooping aerial shots are breathtaking , and the overall experience is awesome . I think this movie was good\\n'\n",
    "print(icl_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d5a2a2-a0f6-4b3a-aecf-4e2b736a0619",
   "metadata": {},
   "source": [
    "You should now get a decent accuracy of 76.4%\n",
    "\n",
    "Compare this to the numbers in [Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank](https://aclanthology.org/D13-1170/) (Socher et al., EMNLP 2013).\n",
    "\n",
    "This is worse than this paper from more than a decade ago, which trained a small Recurrent Neural Network from scratch! Even worse than their baseline (a Naive Bayes classifier with bag of words features).\n",
    "\n",
    "But our model is small for today's standards and it is still quite impressive that it achieves reasonable accuracy without being trained on the task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99ca1503-4a84-45c9-ab8c-5bb93a373ec8",
   "metadata": {},
   "source": [
    "# DPO\n",
    "DPO is a simple RL-free algorithm for training language models from paired preference data. \n",
    "\n",
    "**General DPO Pipeline:**\n",
    "In the standard RLHF (Reinforcement Learning From Human Feedback) setup, the general pipeline of DPO comprises two steps: \n",
    "- Sample model generations $y_1,y_2 \\sim \\pi_{\\mathrm{ref}}(\\cdot \\; | \\; x)$ for prompts $x$, then obtain a human *preference* over the outputs denoted $y_w \\succ y_l \\ |\\ x$ where $y_w$ and $y_l$ are the preferred and dispreferred completion amongst $(y_1, y_2)$, respectively. The output is an offline preference dataset $\\mathcal{D} = \\{x^{(i)}, y_{w}^{(i)}, y_{l}^{(i)}\\}$, where $y_{w}$ represent the preferred model completions, and  $y_{l}$ represent the dispreferred model completions.\n",
    "- Optimize the language model policy $\\pi_{\\theta}$ by minimizing $\\mathcal{L}_{\\mathrm{DPO}}$ for a given $\\pi_{\\mathrm{ref}}$ (base reference model/policy), $\\mathcal{D}$ (preference dataset), and $\\beta$ (a hyperparameter controlling the deviation from $\\pi_{\\mathrm{ref}}$).\n",
    "\n",
    "\n",
    "(Note that for this exercise, we have already created the offline preference dataset for you, and you will use the provided dataset to implement the DPO training in Step 2.)\n",
    "\n",
    "In the following section, you will implement training code for DPO. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f83ddb8-3df6-40a0-a125-928349146ac9",
   "metadata": {
    "id": "Bq6IAyz7M3fH"
   },
   "source": [
    "## Data\n",
    "\n",
    "### Loading\n",
    "For DPO, we need to load in a *pairwise, preference* dataset and make a corresponding dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "df0d1f81-0557-4981-a4ff-a2306746be5e",
   "metadata": {
    "id": "h5zcRHUWNTb8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'prompt': 'Despite the fact that the plot follows the well-known recipe', 'positive': 'Despite the fact that the plot follows the well-known recipe of a romantic comedy, the film is still very charming.', 'negative': 'Despite the fact that the plot follows the well-known recipe for a romantic comedy, it still manages to be fresh and funny.'}, {'prompt': 'I had read the newspaper reviews of this film and', 'positive': 'I had read the newspaper reviews of this film and was under the impression that it was a comedy.', 'negative': 'I had read the newspaper reviews of this film and was not impressed.'}, {'prompt': 'Clearly patterned after the first gangster movies that Warner produced', 'positive': 'Clearly patterned after the first gangster movies that Warner produced in the early 1930s, this one has the same kind of gritty realism and no-nonsense attitude.', 'negative': \"Clearly patterned after the first gangster movies that Warner produced in the 1930s, this is a throwback that doesn't know when to throw in the towel.\"}, {'prompt': 'This movie is a perfect example of a film that', 'positive': 'This movie is a perfect example of a film that is better than the sum of its parts.', 'negative': 'This movie is a perfect example of a film that could have been good, but was ruined by bad writing.'}, {'prompt': \"I haven't reviewed on IMDb before but this documentary is\", 'positive': \"I haven't reviewed on IMDb before but this documentary is so important and so well done that I had to say something.\", 'negative': \"I haven't reviewed on IMDb before but this documentary is so bad that I felt compelled to do so.\"}]\n",
      "10000 1000\n"
     ]
    }
   ],
   "source": [
    "dataset_parallel = load_dataset(\"hallisky/synthetic-imdb-movie-reviews-parallel\")\n",
    "\n",
    "\n",
    "# store prompt, concatenated prompt and positive completion,\n",
    "# and concatenated prompt and negative completion\n",
    "# train_texts_parallel and test_texts_parallel are lists of triples of\n",
    "# (prompt, prompt + pos, prompt + neg).\n",
    "dataset_parallel = load_dataset(\"hallisky/synthetic-imdb-movie-reviews-parallel\")\n",
    "\n",
    "train_texts_parallel = []\n",
    "for item in dataset_parallel[\"train\"]:\n",
    "    train_texts_parallel.append(dict(prompt=item[\"prompt\"], positive=item[\"prompt\"] + item['positive_completion'], negative=item[\"prompt\"] + item['negative_completion']))\n",
    "test_texts_parallel = []\n",
    "for item in dataset_parallel[\"test\"]:\n",
    "    test_texts_parallel.append(dict(prompt=item[\"prompt\"], positive=item[\"prompt\"] + item['positive_completion'], negative=item[\"prompt\"] + item['negative_completion']))\n",
    "# Print the first 5 rows of data\n",
    "print(train_texts_parallel[:5])\n",
    "print(len(train_texts_parallel), len(test_texts_parallel))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ee8304b-8219-457d-854e-85549bddb3f3",
   "metadata": {},
   "source": [
    "### batching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d89670c4-ac60-471b-b4ab-c378df6dbbd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(train_texts_parallel, batch_size=4, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "edc80ad3-a0ce-44fc-afd7-7a32c85ebdd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in train_loader:\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d8631c3b-fe29-4e43-b4c2-ebd3d569f4ec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'prompt': ['Despite the fact that the plot follows the well-known recipe',\n",
       "  'I had read the newspaper reviews of this film and',\n",
       "  'Clearly patterned after the first gangster movies that Warner produced',\n",
       "  'This movie is a perfect example of a film that'],\n",
       " 'positive': ['Despite the fact that the plot follows the well-known recipe of a romantic comedy, the film is still very charming.',\n",
       "  'I had read the newspaper reviews of this film and was under the impression that it was a comedy.',\n",
       "  'Clearly patterned after the first gangster movies that Warner produced in the early 1930s, this one has the same kind of gritty realism and no-nonsense attitude.',\n",
       "  'This movie is a perfect example of a film that is better than the sum of its parts.'],\n",
       " 'negative': ['Despite the fact that the plot follows the well-known recipe for a romantic comedy, it still manages to be fresh and funny.',\n",
       "  'I had read the newspaper reviews of this film and was not impressed.',\n",
       "  \"Clearly patterned after the first gangster movies that Warner produced in the 1930s, this is a throwback that doesn't know when to throw in the towel.\",\n",
       "  'This movie is a perfect example of a film that could have been good, but was ruined by bad writing.']}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f833de-4ba8-4b96-8708-c2817b2edb81",
   "metadata": {},
   "source": [
    "### tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a2757553-60b9-4311-a73a-74f1ef9abff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "positive_ids = tokenizer(batch['positive'], return_tensors='pt', padding=True, truncation=True, max_length=64)['input_ids'].cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2b30a13d-70c6-48e5-afb3-e047df860ada",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 8332,   262,  1109,   326,   262,  7110,  5679,   262,   880,    12,\n",
       "          4002,  8364,   286,   257, 14348, 10997,    11,   262,  2646,   318,\n",
       "           991,   845, 23332,    13, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256],\n",
       "        [   40,   550,  1100,   262,  7533,  8088,   286,   428,  2646,   290,\n",
       "           373,   739,   262, 10647,   326,   340,   373,   257, 10997,    13,\n",
       "         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256],\n",
       "        [30638,  3912,   276,   706,   262,   717,  7706,  1706,  6918,   326,\n",
       "         14469,  4635,   287,   262,  1903, 15533,    82,    11,   428,   530,\n",
       "           468,   262,   976,  1611,   286, 39679, 30787,   290,   645,    12,\n",
       "            77, 46563,  9408,    13],\n",
       "        [ 1212,  3807,   318,   257,  2818,  1672,   286,   257,  2646,   326,\n",
       "           318,  1365,   621,   262,  2160,   286,   663,  3354,    13, 50256,\n",
       "         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256]], device='cuda:0')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "positive_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b339b121-80bf-4b3e-b0dc-d6b6dd669130",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_ids = tokenizer(batch['prompt'], return_tensors='pt', padding=True, truncation=True, max_length=64)['input_ids'].cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a9710764-a662-4f8a-b96c-87c852b6337a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 8332,   262,  1109,   326,   262,  7110,  5679,   262,   880,    12,\n",
       "          4002,  8364],\n",
       "        [   40,   550,  1100,   262,  7533,  8088,   286,   428,  2646,   290,\n",
       "         50256, 50256],\n",
       "        [30638,  3912,   276,   706,   262,   717,  7706,  1706,  6918,   326,\n",
       "         14469,  4635],\n",
       "        [ 1212,  3807,   318,   257,  2818,  1672,   286,   257,  2646,   326,\n",
       "         50256, 50256]], device='cuda:0')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14427940-69f7-4ae4-a712-80339337fc2b",
   "metadata": {},
   "source": [
    "## Computy Cross-Entropy/masking\n",
    "\n",
    "**Masking the Prompts for Computing:**\n",
    "\n",
    "As shown by the DPO loss (see below), the paradigm of DPO and other RLHF methods is to optimize the policy towards a preferred **completion** given a prompt, i.e., optimize for $\\pi_{\\theta}(y_w |  x)$, rather than optimize for the entire sequence, i.e., $\\pi_{\\theta}(x + y_w)$. As such, when computing the log probabilities of the preferred/dispreferred completions from the policy and reference policy model for your DPO loss, make sure to calculate the log probabilities **only of the completions**, not the entire sequence. In your implementation, you can do this by masking out tokens that belong to the prompts (using the information of the prompt length) to sub-select the log probabilities of only the tokens that belong to the completion part.\n",
    "\n",
    "\n",
    "In PW2, we simply set `ignore_index=tokenizer.eos_token_id` in `nn.CrossEntropyLoss`. This ensured that we did not train the model to predict padding.\n",
    "\n",
    "by default, `ignore_index=-100`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "88ee1d73-a5d2-485c-be57-a6baeae7d101",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-100"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_entropy = nn.CrossEntropyLoss(reduction='none')\n",
    "cross_entropy.ignore_index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe35766-98f5-4d5f-9c3c-a4348e2a723a",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "Our problem here is a little more complicated: we want to compute the loss only on the continuations of the prompt, excluding both padding and inputs from the prompts.\n",
    "\n",
    "For this, we will copy and the inputs to use it as labels, like in PW2, but we'll replace the indices that correspond to the prompt and padding to `-100`\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7563643b-b49a-45d9-9cda-418f45f8263f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 8332,   262,  1109,   326,   262,  7110,  5679,   262,   880,    12,\n",
       "          4002,  8364,   286,   257, 14348, 10997,    11,   262,  2646,   318,\n",
       "           991,   845, 23332,    13, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256],\n",
       "        [   40,   550,  1100,   262,  7533,  8088,   286,   428,  2646,   290,\n",
       "           373,   739,   262, 10647,   326,   340,   373,   257, 10997,    13,\n",
       "         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256],\n",
       "        [30638,  3912,   276,   706,   262,   717,  7706,  1706,  6918,   326,\n",
       "         14469,  4635,   287,   262,  1903, 15533,    82,    11,   428,   530,\n",
       "           468,   262,   976,  1611,   286, 39679, 30787,   290,   645,    12,\n",
       "            77, 46563,  9408,    13],\n",
       "        [ 1212,  3807,   318,   257,  2818,  1672,   286,   257,  2646,   326,\n",
       "           318,  1365,   621,   262,  2160,   286,   663,  3354,    13, 50256,\n",
       "         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "         50256, 50256, 50256, 50256]], device='cuda:0')"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = positive_ids.clone()\n",
    "\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6fb38dbd-e068-440f-8a4a-b4d2d10f8abe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,   286,   257, 14348, 10997,    11,   262,  2646,   318,\n",
       "           991,   845, 23332,    13,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100],\n",
       "        [ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "           373,   739,   262, 10647,   326,   340,   373,   257, 10997,    13,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100],\n",
       "        [ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,   287,   262,  1903, 15533,    82,    11,   428,   530,\n",
       "           468,   262,   976,  1611,   286, 39679, 30787,   290,   645,    12,\n",
       "            77, 46563,  9408,    13],\n",
       "        [ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "           318,  1365,   621,   262,  2160,   286,   663,  3354,    13,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100]], device='cuda:0')"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# after correctly setting inputs from the prompt and padding to -100\n",
    "labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82115f3d-1992-4f83-998b-6d4ca9938f5d",
   "metadata": {},
   "source": [
    "Now that you have masked the labels, forward the positive_ids to the model and compute the loss as in PW2/for perplexity above (don't forget to shift the labels by one).\n",
    "\n",
    "The only difference is that we set `reduction='none'` because we want to keep one score per sample in the batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3df00589-ac6b-491a-9356-74e3565eda92",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.1465e+00,\n",
       "        3.6377e+00, 7.2197e+00, 2.4396e+00, 8.0893e-01, 2.0213e+00, 3.1597e+00,\n",
       "        1.8081e+00, 3.6082e+00, 2.7459e+00, 7.9434e+00, 1.2095e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        2.7241e+00, 7.2998e+00, 5.1491e-01, 3.9269e-02, 2.2083e-01, 1.2623e+00,\n",
       "        6.6251e-01, 1.5223e+00, 4.1020e+00, 1.5778e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 2.3589e+00,\n",
       "        2.1315e-01, 2.6189e+00, 4.4915e+00, 1.0524e-01, 1.1350e+00, 2.5632e+00,\n",
       "        3.2035e+00, 3.2706e+00, 2.9058e+00, 2.7138e+00, 4.6536e+00, 5.0916e-03,\n",
       "        5.1215e+00, 3.8952e+00, 1.7343e+00, 6.4550e+00, 2.2655e+00, 8.4306e-01,\n",
       "        2.2434e-03, 3.5734e+00, 1.4887e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        2.1031e+00, 6.4677e+00, 7.9410e-01, 1.7400e+00, 6.7935e+00, 1.9802e-01,\n",
       "        3.1059e-02, 1.6730e-02, 5.1477e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "        0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],\n",
       "       device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits = model(positive_ids, return_dict=True).logits\n",
    "logits = logits[:, :-1].contiguous()\n",
    "labels = labels[:, 1:].contiguous()\n",
    "loss = cross_entropy(logits.reshape(-1, tokenizer.vocab_size), labels.reshape(-1))\n",
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc991157-dcfa-4e1c-a7da-26f4ee021f39",
   "metadata": {},
   "source": [
    "This reduces (averages) the loss for each sample in the batch (naively doing `.mean(1)` would take into account padded labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3bb6805e-40e0-4aea-b474-74a2c7281cce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([3.1457, 1.9926, 2.5281, 2.0732], device='cuda:0',\n",
       "       grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss.reshape(labels.shape).sum(1) / (labels!=cross_entropy.ignore_index).sum(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4367ed02-f477-4c3a-a2b0-181d4ef53229",
   "metadata": {
    "id": "i2MTmyllof2D"
   },
   "source": [
    "## Implementing DPO\n",
    "\n",
    "\n",
    "**Overall DPO Loss:**\n",
    "\n",
    "Recall the DPO loss,\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\mathrm{DPO}}\\left(\\pi_\\theta ; \\pi_{\\mathrm{ref}}\\right)=-\\mathbb{E}_{\\left(x, y_w, y_l\\right) \\sim \\mathcal{D}}\\left[\\log \\sigma\\left(\\beta \\log \\frac{\\pi_\\theta\\left(y_w \\mid x\\right)}{\\pi_{\\mathrm{ref}}\\left(y_w \\mid x\\right)}-\\beta \\log \\frac{\\pi_\\theta\\left(y_l \\mid x\\right)}{\\pi_{\\mathrm{ref}}\\left(y_l \\mid x\\right)}\\right)\\right]\n",
    "$$\n",
    "\n",
    "where $\\sigma$ is the logistic function, $\\mathbb{E}$ represents the expected value, $\\log \\frac{\\pi_\\theta\\left(y_w \\mid x\\right)}{\\pi_{\\mathrm{ref}}\\left(y_w \\mid x\\right)}$ represent the log of the probability ratio of the preferred completion $y_w$ from the parameterized policy vs. from the reference policy, and $\\log \\frac{\\pi_\\theta\\left(y_l \\mid x\\right)}{\\pi_{\\mathrm{ref}}\\left(y_l \\mid x\\right)}$ represents the same log ratio but for the dispreferred compeltion $y_l$. Intuitively, we want to increase the likelihood of the preferred completions $y_w$ and decrease the likelihood of the dispreferred completions $y_l$.\n",
    "\n",
    "(Note that in the starter code that we provide, we call $y_w$ and $y_l$ `chosen` and `rejected` completions, respectively.)\n",
    "\n",
    "\n",
    "\n",
    "Implement DPO by computing the DPO loss.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29771aab-9367-490e-903f-766cb323af2a",
   "metadata": {},
   "source": [
    "We first reload a copy of the pretrained model $\\pi_{\\mathrm{ref}}$, which will not be updated.\n",
    "\n",
    "When computing cross-entropy with this model, disable gradient computations using `with torch.no_grad():`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "dcf35a64-3f85-4a05-a999-2eceb4f9417e",
   "metadata": {},
   "outputs": [],
   "source": [
    "orig_model = AutoModelForCausalLM.from_pretrained(\"openai-community/gpt2\",device_map=\"auto\").eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2315248-302d-4ba4-8393-69cc12be27fc",
   "metadata": {},
   "source": [
    "\n",
    "Remember that in the end we want to compute a fraction like\n",
    "\n",
    "$$\\log \\frac{\\pi_\\theta\\left(y_w \\mid x\\right)}{\\pi_{\\mathrm{ref}}\\left(y_w \\mid x\\right)}$$\n",
    "\n",
    "And that $\\log \\frac{a}{b} = \\log a - \\log b, \\forall a,b$\n",
    "\n",
    "And that $\\log \\pi_\\theta\\left(y_w \\mid x\\right)$ is the opposite (negative) of what you have computed above using `nn.CrossEntropyLoss`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "682ddb35-0ed8-4c56-90cf-e4f758fe57e6",
   "metadata": {},
   "source": [
    "You should get identical values for `orig_model` and `model` because `model` was not updated yet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2b438a2f-c088-4447-bd3b-0dee7a73523d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([-3.1457, -1.9926, -2.5281, -2.0732], device='cuda:0',\n",
       "        grad_fn=<NegBackward0>),\n",
       " tensor([-2.6062, -2.4911, -2.4301, -2.7002], device='cuda:0',\n",
       "        grad_fn=<NegBackward0>),\n",
       " tensor([-3.1457, -1.9926, -2.5281, -2.0732], device='cuda:0'),\n",
       " tensor([-2.6062, -2.4911, -2.4301, -2.7002], device='cuda:0'))"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps = compute_cross_entropies(batch)\n",
    "policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65834a26-5740-43f9-b2cc-95b781881957",
   "metadata": {},
   "source": [
    "Once you've computed all four log-probabilities (via `nn.CrossEntropy`), combine them as in the DPO equation above .\n",
    "\n",
    "Don't forget:\n",
    "- add a final average reduce (the expectation)\n",
    "- the minus sign `-` before the expectation\n",
    "\n",
    "Note the existence of `nn.LogSigmoid`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "76abc7dd-fbe8-49da-a316-31fcf8000b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2d943f56-984b-484f-a666-74b06bf5945b",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_sigmoid = nn.LogSigmoid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3c9746c6-b784-4a89-8eb1-09ccc63886b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.6931, device='cuda:0', grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dpo_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e219537-46a6-46c8-97e0-342737fb716b",
   "metadata": {},
   "source": [
    "Because the `model` was not updated yet, the expression within $\\sigma$ should be 0, so $\\sigma(0)=0.5$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "f13c41c3-1c4f-40f2-8d56-f6f445907169",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.5000, device='cuda:0', grad_fn=<ExpBackward0>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(-dpo_loss).exp()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18adba68-9968-4102-9e0b-3671b781d24e",
   "metadata": {},
   "source": [
    "wrap all of this code in a single function so that you can easily call it in the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "13f6616f-268f-49f8-9adc-d06288dd2192",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_dpo(batch, beta = 0.1):\n",
    "    raise NotImplementedError()\n",
    "    return dpo_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d30b42f-a994-421b-82e7-85f40a9a9568",
   "metadata": {
    "id": "MzyWOXBKyAJf"
   },
   "source": [
    "## Training DPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "c336e5ba-6c8f-466a-bf55-409ebc6d4eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "a3201f57-9850-40d5-a370-b3648725e5a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter(\"logs/pw3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09d217e8-cb0e-49e5-8b30-2ead7d5bfb27",
   "metadata": {},
   "source": [
    "Run tensorboard before training. Refresh during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "9cd9a77a-c425-400a-a3e9-f04b9b825f39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-8535b73fc6c70566\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-8535b73fc6c70566\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir logs/pw3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9d15ecdd-41a9-450e-95c9-ddfaf559b0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.00001)\n",
    "\n",
    "batch_size = 4\n",
    "train_loader = torch.utils.data.DataLoader(train_texts_parallel, batch_size=batch_size, shuffle=True)\n",
    "validation_loader = torch.utils.data.DataLoader(test_texts_parallel, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "steps = 0\n",
    "for epoch in range(1):\n",
    "    for batch in train_loader:\n",
    "        dpo_loss = compute_dpo(batch)\n",
    "        writer.add_scalar(\"Loss/train\", dpo_loss.item(), steps)\n",
    "        dpo_loss.backward()\n",
    "        optimizer.step()\n",
    "        steps += 1\n",
    "\n",
    "        # validation\n",
    "        if steps % 1000 == 0:\n",
    "            with torch.no_grad():\n",
    "                model.eval()\n",
    "                valid_loss = 0\n",
    "                valid_batches = 0\n",
    "                for batch in validation_loader:\n",
    "                    dpo_loss = compute_dpo(batch)\n",
    "                    valid_loss += dpo_loss.item()\n",
    "                    valid_batches += 1\n",
    "                model.train()\n",
    "                writer.add_scalar(\"Loss/validation\", valid_loss/valid_batches, steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "c4922644-7062-4212-8f38-9a01fdaea228",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"logs/pw3/dpo\")"
   ]
  }
 ],
 "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}