12.4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/udlbook/udlbook/blob/main/Notebooks/Chap12/12_4_Decoding_Strategies.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RnIUiieJWu6e"
   },
   "source": [
    "# **Notebook 12.4: Decoding strategies**\n",
    "\n",
    "This practical investigates neural decoding from transformer models.  \n",
    "\n",
    "Work through the cells below, running each cell in turn. In various places you will see the words \"TO DO\". Follow the instructions at these places and make predictions about what is going to happen or write code to complete the functions.\n",
    "\n",
    "Contact me at udlbookmail@gmail.com if you find any mistakes or have any suggestions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "7abjZ9pMVj3k"
   },
   "outputs": [],
   "source": [
    "#!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "sMOyD0zem2Ef"
   },
   "outputs": [],
   "source": [
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer, set_seed\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "pZgfxbzKWNSR"
   },
   "outputs": [],
   "source": [
    "# Load model and tokenizer\n",
    "model = GPT2LMHeadModel.from_pretrained('gpt2')\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TfhAGy0TXEvV"
   },
   "source": [
    "# Decoding from GPT2\n",
    "\n",
    "This tutorial investigates how to use GPT2 (the forerunner of GPT3) to generate text.  There are a number of ways to do this that trade-off the realism of the text against the amount of variation.\n",
    "\n",
    "At every stage, GPT2 takes an input string and returns a probability for each of the possible subsequent tokens.  We can choose what to do with these probability.  We could always *greedily choose* the most likely next token, or we could draw a *sample* randomly according to the probabilities.  There are also intermediate strategies such as *top-k sampling* and *nucleus sampling*, that have some controlled randomness.\n",
    "\n",
    "We'll also investigate *beam search* -- the idea is that rather than greedily take the next best token at each stage, we maintain a set of hypotheses  (beams)as we add each subsequent token and return the most likely overall hypothesis.  This is not necessarily the same result we get from greedily choosing the next token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vsmO9ptzau3_"
   },
   "source": [
    "First, let's investigate the token themselves.  The code below prints out the vocabulary size and shows 20 random tokens.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "dmmBNS5GY_yk"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens in dictionary = 50257\n",
      "Token: 33003  Mormons\n",
      "Token: 12172  cam\n",
      "Token: 5192  trig\n",
      "Token: 32511 ojure\n",
      "Token: 50057  gist\n",
      "Token: 43723  Petition\n",
      "Token: 7813  sin\n",
      "Token: 21440  Witness\n",
      "Token: 32912  Remy\n",
      "Token: 20609 isure\n",
      "Token: 49100  creeps\n",
      "Token: 7751  fasc\n",
      "Token: 43757  Alc\n",
      "Token: 31228  messenger\n",
      "Token: 36230  SYSTEM\n",
      "Token: 32025  precipitation\n",
      "Token: 21758  cores\n",
      "Token: 45413  Forestry\n",
      "Token: 35730  guru\n",
      "Token: 8444  Disc\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(1)\n",
    "print(\"Number of tokens in dictionary = %d\"%(tokenizer.vocab_size))\n",
    "for i in range(20):\n",
    "  index = np.random.randint(tokenizer.vocab_size)\n",
    "  print(\"Token: %d \"%(index)+tokenizer.decode(torch.tensor(index), skip_special_tokens=True))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MUM3kLEjbTso"
   },
   "source": [
    "# Sampling\n",
    "\n",
    "Each time we run GPT2 it will take in a set of tokens, and return a probability over each of the possible next tokens.  The simplest thing we could do is to just draw a sample from this probability distribution each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "TIyNgg0FkJKO"
   },
   "outputs": [],
   "source": [
    "def sample_next_token(input_tokens, model, tokenizer):\n",
    "  # Run model to get prediction over next output\n",
    "  outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])\n",
    "  # Find prediction\n",
    "  prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]\n",
    "  # TODO Draw a random token according to the probabilities\n",
    "  # next_token should be an array with an sole integer in it (as below)\n",
    "  # Use:  https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html\n",
    "  # Replace this line\n",
    "  next_token = [np.random.choice(len(prob_over_tokens), p=prob_over_tokens)]\n",
    "\n",
    "\n",
    "  # Append token to sentence\n",
    "  output_tokens = input_tokens\n",
    "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n",
    "  output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)\n",
    "  output_tokens['last_token_prob'] = prob_over_tokens[next_token]\n",
    "\n",
    "  return output_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "BHs-IWaz9MNY"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best thing about Bath is that\n",
      "The best thing about Bath is that they\n",
      "The best thing about Bath is that they don\n",
      "The best thing about Bath is that they don't\n",
      "The best thing about Bath is that they don't even\n",
      "The best thing about Bath is that they don't even change\n",
      "The best thing about Bath is that they don't even change or\n",
      "The best thing about Bath is that they don't even change or shrink\n",
      "The best thing about Bath is that they don't even change or shrink anymore\n",
      "The best thing about Bath is that they don't even change or shrink anymore.\n"
     ]
    }
   ],
   "source": [
    "# Expected output:\n",
    "# \"The best thing about Bath is that they don't even change or shrink anymore.\"\n",
    "\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = sample_next_token(input_tokens, model, tokenizer)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "yN98_7WqbvIe"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Those who can't see these\n",
      "Those who can't see these are\n",
      "Those who can't see these are being\n",
      "Those who can't see these are being dem\n",
      "Those who can't see these are being demoted\n",
      "Those who can't see these are being demoted.\n",
      "Those who can't see these are being demoted. \"\n",
      "Those who can't see these are being demoted. \"This\n"
     ]
    }
   ],
   "source": [
    "# TODO Modify the code below by changing the number of tokens generated and the initial sentence\n",
    "# to get a feel for how well this works.  Since I didn't reset the seed, it will give a different\n",
    "# answer every time that you run it.\n",
    "\n",
    "# TODO Experiment with changing this line:\n",
    "input_txt = \"Those who can't see\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "# TODO Experiment with changing this line:\n",
    "for i in range(8):\n",
    "    input_tokens = sample_next_token(input_tokens, model, tokenizer)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7eHFLCeZcmmg"
   },
   "source": [
    "# Greedy token selection\n",
    "\n",
    "You probably (correctly) got the impression that the text from pure sampling of the probability model can be kind of random.  How about if we choose most likely token at each step?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "OhRzynEjxpZF"
   },
   "outputs": [],
   "source": [
    "def get_best_next_token(input_tokens, model, tokenizer):\n",
    "  # Run model to get prediction over next output\n",
    "  outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])\n",
    "  # Find prediction\n",
    "  prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]\n",
    "\n",
    "  # TODO -- find the token index with the maximum probability\n",
    "  # It should be returns as a list (i.e., put squared brackets around it)\n",
    "  # Use https://numpy.org/doc/stable/reference/generated/numpy.argmax.html\n",
    "  # Replace this line\n",
    "  next_token = [np.argmax(prob_over_tokens)]\n",
    "\n",
    "  # Append token to sentence\n",
    "  output_tokens = input_tokens\n",
    "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n",
    "  output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)\n",
    "  output_tokens['last_token_prob'] = prob_over_tokens[next_token]\n",
    "  return output_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "gKB1Mgndj-Hm"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best thing about Bath is that\n",
      "The best thing about Bath is that it\n",
      "The best thing about Bath is that it's\n",
      "The best thing about Bath is that it's a\n",
      "The best thing about Bath is that it's a place\n",
      "The best thing about Bath is that it's a place where\n",
      "The best thing about Bath is that it's a place where you\n",
      "The best thing about Bath is that it's a place where you can\n",
      "The best thing about Bath is that it's a place where you can go\n",
      "The best thing about Bath is that it's a place where you can go to\n"
     ]
    }
   ],
   "source": [
    "# Expected output:\n",
    "# The best thing about Bath is that it's a place where you can go to\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = get_best_next_token(input_tokens, model, tokenizer)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "L1YHKaYFfC0M"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Those who can't see the\n",
      "Those who can't see the difference\n",
      "Those who can't see the difference between\n",
      "Those who can't see the difference between the\n",
      "Those who can't see the difference between the two\n",
      "Those who can't see the difference between the two are\n",
      "Those who can't see the difference between the two are not\n",
      "Those who can't see the difference between the two are not going\n",
      "Those who can't see the difference between the two are not going to\n",
      "Those who can't see the difference between the two are not going to be\n",
      "Those who can't see the difference between the two are not going to be able\n",
      "Those who can't see the difference between the two are not going to be able to\n",
      "Those who can't see the difference between the two are not going to be able to see\n",
      "Those who can't see the difference between the two are not going to be able to see it\n",
      "Those who can't see the difference between the two are not going to be able to see it.\n"
     ]
    }
   ],
   "source": [
    "# TODO Modify the code below by changing the number of tokens generated and the initial sentence\n",
    "# to get a feel for how well this works.\n",
    "\n",
    "# TODO Experiment with changing this line:\n",
    "input_txt = \"Those who can't see\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "# TODO Experiment with changing this line:\n",
    "for i in range(15):\n",
    "    input_tokens = get_best_next_token(input_tokens, model, tokenizer)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1ORFXYX_gBDT"
   },
   "source": [
    "# Top-K sampling\n",
    "\n",
    "You probably noticed that the greedy strategy produces quite realistic text, but it's kind of boring.  It produces generic answers.  Also, if this was a chatbot, then we wouldn't necessarily want it to produce the same answer to a question each time.  \n",
    "\n",
    "Top-K sampling is a compromise strategy that samples randomly from the top K most probable tokens.  We could just choose them with a uniform distribution, or (as here) we could sample them according to their original probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "7RFbn6c-0Z4v"
   },
   "outputs": [],
   "source": [
    "def get_top_k_token(input_tokens, model, tokenizer, k=20):\n",
    "  # Run model to get prediction over next output\n",
    "  outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])\n",
    "  # Find prediction\n",
    "  prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]\n",
    "\n",
    "  # Draw a sample from the top K most likely tokens.\n",
    "  # Take copy of the probabilities and sort from largest to smallest (use np.sort)\n",
    "  # TODO -- replace this line\n",
    "  sorted_prob_over_tokens = prob_over_tokens[np.argsort(prob_over_tokens)[-k:]]\n",
    "\n",
    "  # Find the probability at the k'th position\n",
    "  # TODO -- replace this line\n",
    "  kth_prob_value = sorted_prob_over_tokens[0] # the way i did, should be the first\n",
    "\n",
    "  # Set all probabilities below this value to zero\n",
    "  prob_over_tokens[prob_over_tokens<kth_prob_value] = 0\n",
    "\n",
    "  # Renormalize the probabilities so that they sum to one\n",
    "  # TODO -- replace this line\n",
    "  prob_over_tokens /= prob_over_tokens.sum()\n",
    "\n",
    "  # Draw random token\n",
    "  next_token = np.random.choice(len(prob_over_tokens), 1, replace=False, p=prob_over_tokens)\n",
    "\n",
    "  # Append token to sentence\n",
    "  output_tokens = input_tokens\n",
    "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n",
    "  output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)\n",
    "  output_tokens['last_token_prob'] = prob_over_tokens[next_token]\n",
    "  return output_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "G3w1GVED4HYv"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Poggers\\AppData\\Local\\Temp\\ipykernel_9716\\1512861585.py:28: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at C:\\actions-runner\\_work\\pytorch\\pytorch\\builder\\windows\\pytorch\\torch\\csrc\\utils\\tensor_new.cpp:278.)\n",
      "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best thing about Bath is that\n",
      "The best thing about Bath is that you\n",
      "The best thing about Bath is that you get\n",
      "The best thing about Bath is that you get to\n",
      "The best thing about Bath is that you get to see\n",
      "The best thing about Bath is that you get to see all\n",
      "The best thing about Bath is that you get to see all the\n",
      "The best thing about Bath is that you get to see all the beautiful\n",
      "The best thing about Bath is that you get to see all the beautiful faces\n",
      "The best thing about Bath is that you get to see all the beautiful faces of\n"
     ]
    }
   ],
   "source": [
    "# Expected output:\n",
    "# The best thing about Bath is that you get to see all the beautiful faces of\n",
    "\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = get_top_k_token(input_tokens, model, tokenizer, k=10)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "RySu2bzqpW9E"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best thing about Bath is that\n",
      "The best thing about Bath is that all\n",
      "The best thing about Bath is that all you\n",
      "The best thing about Bath is that all you have\n",
      "The best thing about Bath is that all you have to\n",
      "The best thing about Bath is that all you have to do\n",
      "The best thing about Bath is that all you have to do is\n",
      "The best thing about Bath is that all you have to do is submit\n",
      "The best thing about Bath is that all you have to do is submit applications\n",
      "The best thing about Bath is that all you have to do is submit applications in\n"
     ]
    }
   ],
   "source": [
    "# TODO\n",
    "# Experiment with different values of k\n",
    "# If you set it to a lower number (say 3) the text will be less random\n",
    "# If you set it to a higher number (say 5000) the text will be more random\n",
    "\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = get_top_k_token(input_tokens, model, tokenizer, k=5000)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fOHak_QJfU-2"
   },
   "source": [
    "# Nucleus sampling\n",
    "\n",
    "Top-K sampling has the disadvantage that sometimes there are only a few plausible next tokens, and sometimes there are a lot.  How do we adapt to this situation?  One way is to sample from a fixed proportion of the probability mass.  That is we order the tokens in terms of probability and cut off the possibility of sampling when the cumulative sum is greater than a threshold.\n",
    "\n",
    "This way, we adapt the number of possible tokens that we can choose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "PtxS4kNDyUcm"
   },
   "outputs": [],
   "source": [
    "def get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh=0.25):\n",
    "  # Run model to get prediction over next output\n",
    "  outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])\n",
    "  # Find prediction\n",
    "  prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]\n",
    "\n",
    "  # Find the most likely tokens that make up the first (thresh) of the probability\n",
    "  # TODO -- sort the probabilities in decreasing order\n",
    "  # Replace this line\n",
    "  sorted_probs_decreasing = prob_over_tokens[np.argsort(prob_over_tokens)[::-1]]\n",
    "  # TODO -- compute the cumulative sum of these probabilities\n",
    "  # Replace this line\n",
    "  cum_sum_probs = np.cumsum(sorted_probs_decreasing)\n",
    "\n",
    "\n",
    "  # Find index where that the cumulative sum is greater than the threshold\n",
    "  thresh_index = np.argmax(cum_sum_probs>thresh)\n",
    "  print(\"Choosing from %d tokens\"%(thresh_index))\n",
    "  # TODO:  Find the probability value to threshold\n",
    "  # Replace this line:\n",
    "  thresh_prob = sorted_probs_decreasing[thresh_index]\n",
    "\n",
    "\n",
    "  # Set any probabilities less than this to zero\n",
    "  prob_over_tokens[prob_over_tokens<thresh_prob] = 0\n",
    "  # Renormalize\n",
    "  prob_over_tokens = prob_over_tokens / np.sum(prob_over_tokens)\n",
    "  # Draw random token\n",
    "  next_token = np.random.choice(len(prob_over_tokens), 1, replace=False, p=prob_over_tokens)\n",
    "\n",
    "  # Append token to sentence\n",
    "  output_tokens = input_tokens\n",
    "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n",
    "  output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)\n",
    "  output_tokens['last_token_prob'] = prob_over_tokens[next_token]\n",
    "  return output_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "K2Vk1Ly40S6c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing from 0 tokens\n",
      "The best thing about Bath is that\n",
      "Choosing from 0 tokens\n",
      "The best thing about Bath is that it\n",
      "Choosing from 0 tokens\n",
      "The best thing about Bath is that it's\n",
      "Choosing from 2 tokens\n",
      "The best thing about Bath is that it's not\n",
      "Choosing from 1 tokens\n",
      "The best thing about Bath is that it's not a\n",
      "Choosing from 25 tokens\n",
      "The best thing about Bath is that it's not a city\n",
      "Choosing from 2 tokens\n",
      "The best thing about Bath is that it's not a city that\n",
      "Choosing from 1 tokens\n",
      "The best thing about Bath is that it's not a city that has\n",
      "Choosing from 1 tokens\n",
      "The best thing about Bath is that it's not a city that has been\n",
      "Choosing from 11 tokens\n",
      "The best thing about Bath is that it's not a city that has been around\n"
     ]
    }
   ],
   "source": [
    "# Expected output:\n",
    "# The best thing about Bath is that it's not a city that has been around\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh = 0.2)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "id": "eQNNHe14wDvC"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Choosing from 40 tokens\n",
      "The best thing about Bath is actually\n",
      "Choosing from 494 tokens\n",
      "The best thing about Bath is actually seeing\n",
      "Choosing from 164 tokens\n",
      "The best thing about Bath is actually seeing it\n",
      "Choosing from 342 tokens\n",
      "The best thing about Bath is actually seeing it without\n",
      "Choosing from 1579 tokens\n",
      "The best thing about Bath is actually seeing it without the\n",
      "Choosing from 5428 tokens\n",
      "The best thing about Bath is actually seeing it without the dual\n",
      "Choosing from 2556 tokens\n",
      "The best thing about Bath is actually seeing it without the dual use\n",
      "Choosing from 2086 tokens\n",
      "The best thing about Bath is actually seeing it without the dual use dome\n",
      "Choosing from 442 tokens\n",
      "The best thing about Bath is actually seeing it without the dual use dome;\n",
      "Choosing from 218 tokens\n",
      "The best thing about Bath is actually seeing it without the dual use dome; it\n"
     ]
    }
   ],
   "source": [
    "# TODO -- experiment with setting the threshold probability to larger or smaller values\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "for i in range(10):\n",
    "    input_tokens = get_nucleus_sampling_token(input_tokens, model, tokenizer, thresh = 0.9)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WMMNeLixwlgM"
   },
   "source": [
    "# Beam search\n",
    "\n",
    "All of the methods we've seen so far choose the tokens one by one.  But this isn't necessarily sensible.  Even greedily choosing the best token doesn't necessarily retrieve the sequence with the highest probability.  It might be that the most likely token only has very unlikely tokens following it.\n",
    "\n",
    "Beam search maintains $K$ hypotheses about the best possible continuation.  It starts with the top $K$ continuations.  Then for each of those, it finds the top K continuations, giving $K^2$ hypotheses.  Then it retains just the top $K$ of these so that the number of hypotheses stays the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "id": "sAI2bClXCe2F"
   },
   "outputs": [],
   "source": [
    "# This routine returns the k'th most likely next token.\n",
    "# If k =0 then it returns the most likely token, if k=1 it returns the next most likely and so on\n",
    "# We will need this for beam search\n",
    "def get_kth_most_likely_token(input_tokens, model, tokenizer, k):\n",
    "  # Run model to get prediction over next output\n",
    "  outputs = model(input_ids = input_tokens['input_ids'], attention_mask = input_tokens['attention_mask'])\n",
    "  # Find prediction\n",
    "  prob_over_tokens = F.softmax(outputs.logits, dim=-1).detach().numpy()[0,-1]\n",
    "\n",
    "  # Find the k'th most likely token\n",
    "  # TODO Sort the probabilities from largest to smallest\n",
    "  # Replace this line:\n",
    "  sorted_prob_over_tokens = prob_over_tokens[np.argsort(prob_over_tokens)[::-1]]\n",
    "  # TODO Find the k'th sorted probability\n",
    "  # Replace this line\n",
    "  kth_prob_value = sorted_prob_over_tokens[k]\n",
    "\n",
    "\n",
    "  # Find position of this token.\n",
    "  next_token = np.where(prob_over_tokens == kth_prob_value)[0]\n",
    "\n",
    "  # Append token to sentence\n",
    "  output_tokens = input_tokens\n",
    "  output_tokens[\"input_ids\"] = torch.cat((output_tokens['input_ids'],torch.tensor([next_token])),dim=1)\n",
    "  output_tokens['attention_mask'] = torch.cat((output_tokens['attention_mask'],torch.tensor([[1]])),dim=1)\n",
    "  output_tokens['last_token_prob'] = prob_over_tokens[next_token]\n",
    "  output_tokens['log_prob'] = output_tokens['log_prob'] + np.log(prob_over_tokens[next_token])\n",
    "  return output_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "6kSc0WrTELMd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The best thing about Bath is the\n",
      "The best thing about Bath is the way\n",
      "The best thing about Bath is the way you\n",
      "The best thing about Bath is the way you get\n",
      "The best thing about Bath is the way you get the\n",
      "The best thing about Bath is the way you get the most\n",
      "The best thing about Bath is the way you get the most bang\n",
      "The best thing about Bath is the way you get the most bang out\n",
      "The best thing about Bath is the way you get the most bang outta\n",
      "The best thing about Bath is the way you get the most bang outta the\n",
      "The best thing about Bath is mixed\n",
      "The best thing about Bath is mixed profits\n",
      "The best thing about Bath is mixed profits partnerships\n",
      "The best thing about Bath is mixed profits partnerships»\n",
      "The best thing about Bath is mixed profits partnerships» buy\n",
      "The best thing about Bath is mixed profits partnerships» buy generic\n",
      "The best thing about Bath is mixed profits partnerships» buy generic+\n",
      "The best thing about Bath is mixed profits partnerships» buy generic+ Honda\n",
      "The best thing about Bath is mixed profits partnerships» buy generic+ Honda throttle\n",
      "The best thing about Bath is mixed profits partnerships» buy generic+ Honda throttlecont\n"
     ]
    }
   ],
   "source": [
    "# We can test this code and see that if we choose the 2nd most likely (K=1) token each time\n",
    "# then we get much better generation results than if we choose the 2001st most likely token\n",
    "\n",
    "# Expected output:\n",
    "# The best thing about Bath is the way you get the most bang outta the\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "input_tokens['log_prob'] = 0.0\n",
    "for i in range(10):\n",
    "    input_tokens = get_kth_most_likely_token(input_tokens, model, tokenizer, k=1)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))\n",
    "\n",
    "# Expected output:\n",
    "# The best thing about Bath is mixed profits partnerships» buy generic+ Honda throttlecont\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "input_tokens['log_prob'] = 0.0\n",
    "for i in range(10):\n",
    "    input_tokens = get_kth_most_likely_token(input_tokens, model, tokenizer, k=2000)\n",
    "    print(tokenizer.decode(input_tokens[\"input_ids\"][0], skip_special_tokens=True))\n",
    "\n",
    "# TODO -- play around with different values of K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "Y4hFfwPFFxka"
   },
   "outputs": [],
   "source": [
    "# Print out each beam plus the log probability\n",
    "def print_beams(beams):\n",
    "  for index,beam in enumerate(beams):\n",
    "    print(\"Beam %d, Prob %3.3f: \"%(index,beam['log_prob'])+tokenizer.decode(beam[\"input_ids\"][0], skip_special_tokens=True))\n",
    "  print('---')\n",
    "\n",
    "\n",
    "# TODO:  Read this code carefully!\n",
    "def do_beam_search(input_tokens_in, model, tokenizer, n_beam=5, beam_length=10):\n",
    "  # Store beams in a list\n",
    "  input_tokens['log_prob'] = 0.0\n",
    "\n",
    "  # Initialize with n_beam most likely continuations\n",
    "  beams = [None] * n_beam\n",
    "  for c_k in range(n_beam):\n",
    "    beams[c_k] = dict(input_tokens_in)\n",
    "    beams[c_k] = get_kth_most_likely_token(beams[c_k], model, tokenizer, c_k)\n",
    "\n",
    "  print_beams(beams)\n",
    "\n",
    "  # For each token in the sequence we will add\n",
    "  for c_pos in range(beam_length-1):\n",
    "    # Now for each beam, we continue it in the most likely ways, making n_beam*n_beam type hypotheses\n",
    "    beams_all = [None] * (n_beam*n_beam)\n",
    "    log_probs_all = np.zeros(n_beam*n_beam)\n",
    "    # For each current hypothesis\n",
    "    for c_beam in range(n_beam):\n",
    "      # For each continuation\n",
    "      for c_k in range(n_beam):\n",
    "        # Store the continuation and the probability\n",
    "        beams_all[c_beam * n_beam + c_k] = dict(get_kth_most_likely_token(beams[c_beam], model, tokenizer, c_k))\n",
    "        log_probs_all[c_beam * n_beam + c_k] = beams_all[c_beam * n_beam + c_k]['log_prob']\n",
    "\n",
    "    # Keep the best n_beams sequences with the highest probabilities\n",
    "    sorted_index = np.argsort(np.array(log_probs_all)*-1)\n",
    "    for c_k in range(n_beam):\n",
    "      beams[c_k] = dict(beams_all[sorted_index[c_k]])\n",
    "\n",
    "    # Print the beams\n",
    "    print_beams(beams)\n",
    "\n",
    "  return beams[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "id": "0YWKwZmz4NXb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beam 0, Prob -0.727: The best thing about Bath is that\n",
      "Beam 1, Prob -2.161: The best thing about Bath is the\n",
      "Beam 2, Prob -3.177: The best thing about Bath is it\n",
      "Beam 3, Prob -3.468: The best thing about Bath is how\n",
      "Beam 4, Prob -3.536: The best thing about Bath is you\n",
      "---\n",
      "Beam 0, Prob -1.899: The best thing about Bath is that it\n",
      "Beam 1, Prob -3.931: The best thing about Bath is it's\n",
      "Beam 2, Prob -4.079: The best thing about Bath is that it is\n",
      "Beam 3, Prob -4.433: The best thing about Bath is the fact\n",
      "Beam 4, Prob -4.553: The best thing about Bath is you can\n",
      "---\n",
      "Beam 0, Prob -2.741: The best thing about Bath is that it's\n",
      "Beam 1, Prob -4.657: The best thing about Bath is the fact that\n",
      "Beam 2, Prob -5.331: The best thing about Bath is that it's not\n",
      "Beam 3, Prob -6.227: The best thing about Bath is it's a\n",
      "Beam 4, Prob -6.264: The best thing about Bath is that it is a\n",
      "---\n",
      "Beam 0, Prob -4.938: The best thing about Bath is that it's a\n",
      "Beam 1, Prob -6.012: The best thing about Bath is the fact that it\n",
      "Beam 2, Prob -7.313: The best thing about Bath is that it's not a\n",
      "Beam 3, Prob -7.907: The best thing about Bath is the fact that it is\n",
      "Beam 4, Prob -8.111: The best thing about Bath is that it's a great\n",
      "---\n",
      "Beam 0, Prob -6.919: The best thing about Bath is the fact that it's\n",
      "Beam 1, Prob -7.764: The best thing about Bath is that it's a place\n",
      "Beam 2, Prob -8.930: The best thing about Bath is that it's a great place\n",
      "Beam 3, Prob -9.479: The best thing about Bath is that it's a place that\n",
      "Beam 4, Prob -9.543: The best thing about Bath is the fact that it's not\n",
      "---\n",
      "Beam 0, Prob -8.451: The best thing about Bath is that it's a place where\n",
      "Beam 1, Prob -9.125: The best thing about Bath is the fact that it's a\n",
      "Beam 2, Prob -9.285: The best thing about Bath is that it's a great place to\n",
      "Beam 3, Prob -10.536: The best thing about Bath is that it's a place where people\n",
      "Beam 4, Prob -11.255: The best thing about Bath is that it's a place that's\n",
      "---\n",
      "Beam 0, Prob -9.306: The best thing about Bath is that it's a place where you\n",
      "Beam 1, Prob -11.327: The best thing about Bath is that it's a place where people can\n",
      "Beam 2, Prob -11.839: The best thing about Bath is that it's a great place to live\n",
      "Beam 3, Prob -12.031: The best thing about Bath is that it's a place where you don\n",
      "Beam 4, Prob -12.366: The best thing about Bath is the fact that it's a place\n",
      "---\n",
      "Beam 0, Prob -9.753: The best thing about Bath is that it's a place where you can\n",
      "Beam 1, Prob -12.031: The best thing about Bath is that it's a place where you don't\n",
      "Beam 2, Prob -12.629: The best thing about Bath is that it's a place where you can get\n",
      "Beam 3, Prob -13.006: The best thing about Bath is that it's a great place to live.\n",
      "Beam 4, Prob -13.108: The best thing about Bath is the fact that it's a place where\n",
      "---\n",
      "Beam 0, Prob -12.274: The best thing about Bath is that it's a place where you can go\n",
      "Beam 1, Prob -12.603: The best thing about Bath is that it's a place where you don't have\n",
      "Beam 2, Prob -14.056: The best thing about Bath is the fact that it's a place where you\n",
      "Beam 3, Prob -14.324: The best thing about Bath is that it's a place where you can go and\n",
      "Beam 4, Prob -14.492: The best thing about Bath is that it's a great place to live. It\n",
      "---\n",
      "Beam 0, Prob -12.662: The best thing about Bath is that it's a place where you don't have to\n",
      "Beam 1, Prob -13.658: The best thing about Bath is that it's a place where you can go to\n",
      "Beam 2, Prob -14.469: The best thing about Bath is the fact that it's a place where you can\n",
      "Beam 3, Prob -14.785: The best thing about Bath is that it's a great place to live. It's\n",
      "Beam 4, Prob -14.896: The best thing about Bath is that it's a place where you don't have to be\n",
      "---\n",
      "Beam search result:\n",
      "The best thing about Bath is that it's a place where you don't have to\n"
     ]
    }
   ],
   "source": [
    "# Expected output:\n",
    "# The best thing about Bath is that it's a place where you don't have to\n",
    "\n",
    "set_seed(0)\n",
    "input_txt = \"The best thing about Bath is\"\n",
    "input_tokens = tokenizer(input_txt, return_tensors='pt')\n",
    "\n",
    "# Now let's call the beam search\n",
    "# It takes a while as it has to run the model multiple times to add a token\n",
    "n_beams = 5\n",
    "best_beam = do_beam_search(input_tokens,model,tokenizer)\n",
    "print(\"Beam search result:\")\n",
    "print(tokenizer.decode(best_beam[\"input_ids\"][0], skip_special_tokens=True))\n",
    "\n",
    "# You should see that the best answer is not the same as the greedy solution we found above\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-SXpjZPYsMhv"
   },
   "source": [
    "You can read about more decoding strategies in this blog (which uses a recursive neural network, not a transformer, but the principles are the same).\n",
    "\n",
    "https://www.borealisai.com/research-blogs/tutorial-6-neural-natural-language-generation-decoding-algorithms/\n",
    "\n",
    "You can also look at other possible language models via hugging face:\n",
    "\n",
    "https://huggingface.co/docs/transformers/v4.25.1/en/model_summary#decoders-or-autoregressive-models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyPsZjfqVeHYh95Hzt+hCIO7",
   "include_colab_link": true,
   "provenance": []
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}