Generate fluent English text using transformers GPT2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Example) Generate fluent English text using transformers GPT2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> From https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb#scrollTo=a8Y7cgu9ohXP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This post shows how to generate language with very little effort using \"transformers\" library"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### From http://jalammar.github.io/illustrated-gpt2/\n",
    "All of the following functionalities can be used for **auto-regressive** language generation. In short, *auto-regressive* language generation is based on the assumption that the probability distribution of a word sequence can be decomposed into the product of conditional next word distributions: \n",
    "$$ P(w_{1:T} | W_0 ) = \\prod_{t=1}^T P(w_{t} | w_{1: t-1}, W_0) \\text{ ,with }  w_{1: 0} = \\emptyset, $$\n",
    "\n",
    "and $W_0$ being the initial *context* word sequence. The length $T$ of the word sequence is usually determined *on-the-fly* and corresponds to the timestep $t=T$ the EOS token is generated from $P(w_{t} | w_{1: t-1}, W_{0})$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Auto-regressive language generation is now available for GPT2, XLNet, OpenAi-GPT, CTRL, TransfoXL, XLM, Bart, T5 in both PyTorch and Tensorflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And Let's take a looks at popular decoding method like Greedy search, Beam search, Top-K sampling and Top-p sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:12:38.375832 34540 file_utils.py:39] PyTorch version 1.5.0 available.\n",
      "I0605 23:12:38.378586 34540 file_utils.py:55] TensorFlow version 2.2.0 available.\n",
      "wandb: WARNING W&B installed but not logged in.  Run `wandb login` or set the WANDB_API_KEY env variable.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from transformers import TFGPT2LMHeadModel, GPT2Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:16.266966 34540 filelock.py:274] Lock 2935808556808 acquired on C:\\Users\\bokhy/.cache\\torch\\transformers\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71.lock\n",
      "I0605 23:13:16.268923 34540 file_utils.py:436] https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json not found in cache or force_download set to True, downloading to C:\\Users\\bokhy\\.cache\\torch\\transformers\\tmp_4wo1kkw\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "076279fbdf784b4e8095610a02119d84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1042301.0, style=ProgressStyle(descript…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:17.496487 34540 file_utils.py:440] storing https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json in cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71\n",
      "I0605 23:13:17.498482 34540 file_utils.py:443] creating metadata file for C:\\Users\\bokhy/.cache\\torch\\transformers\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71\n",
      "I0605 23:13:17.501474 34540 filelock.py:318] Lock 2935808556808 released on C:\\Users\\bokhy/.cache\\torch\\transformers\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71.lock\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:17.987948 34540 filelock.py:274] Lock 2936530419592 acquired on C:\\Users\\bokhy/.cache\\torch\\transformers\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda.lock\n",
      "I0605 23:13:17.989944 34540 file_utils.py:436] https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt not found in cache or force_download set to True, downloading to C:\\Users\\bokhy\\.cache\\torch\\transformers\\tmpehvp5b3s\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "adb16756784b4de884201e81d8bd7348",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=456318.0, style=ProgressStyle(descripti…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:19.038484 34540 file_utils.py:440] storing https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt in cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda\n",
      "I0605 23:13:19.040477 34540 file_utils.py:443] creating metadata file for C:\\Users\\bokhy/.cache\\torch\\transformers\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda\n",
      "I0605 23:13:19.042473 34540 filelock.py:318] Lock 2936530419592 released on C:\\Users\\bokhy/.cache\\torch\\transformers\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda.lock\n",
      "I0605 23:13:19.043488 34540 tokenization_utils.py:1015] loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json from cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71\n",
      "I0605 23:13:19.044468 34540 tokenization_utils.py:1015] loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt from cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "<transformers.tokenization_gpt2.GPT2Tokenizer object at 0x000002AB8BD37F08>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:19.556293 34540 filelock.py:274] Lock 2936530781384 acquired on C:\\Users\\bokhy/.cache\\torch\\transformers\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.db13c9bc9c7bdd738ec89e069621d88e05dc670366092d809a9cbcac6798e24e.lock\n",
      "I0605 23:13:19.558256 34540 file_utils.py:436] https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json not found in cache or force_download set to True, downloading to C:\\Users\\bokhy\\.cache\\torch\\transformers\\tmp4kmv67nt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41bfe48ebbbe4b95a40f347ba8defcd7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=665.0, style=ProgressStyle(description_…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:19.971308 34540 file_utils.py:440] storing https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json in cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.db13c9bc9c7bdd738ec89e069621d88e05dc670366092d809a9cbcac6798e24e\n",
      "I0605 23:13:19.973305 34540 file_utils.py:443] creating metadata file for C:\\Users\\bokhy/.cache\\torch\\transformers\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.db13c9bc9c7bdd738ec89e069621d88e05dc670366092d809a9cbcac6798e24e\n",
      "I0605 23:13:19.975299 34540 filelock.py:318] Lock 2936530781384 released on C:\\Users\\bokhy/.cache\\torch\\transformers\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.db13c9bc9c7bdd738ec89e069621d88e05dc670366092d809a9cbcac6798e24e.lock\n",
      "I0605 23:13:19.983277 34540 configuration_utils.py:285] loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json from cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.db13c9bc9c7bdd738ec89e069621d88e05dc670366092d809a9cbcac6798e24e\n",
      "I0605 23:13:19.985306 34540 configuration_utils.py:321] Model config GPT2Config {\n",
      "  \"activation_function\": \"gelu_new\",\n",
      "  \"architectures\": [\n",
      "    \"GPT2LMHeadModel\"\n",
      "  ],\n",
      "  \"attn_pdrop\": 0.1,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"embd_pdrop\": 0.1,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"layer_norm_epsilon\": 1e-05,\n",
      "  \"model_type\": \"gpt2\",\n",
      "  \"n_ctx\": 1024,\n",
      "  \"n_embd\": 768,\n",
      "  \"n_head\": 12,\n",
      "  \"n_layer\": 12,\n",
      "  \"n_positions\": 1024,\n",
      "  \"pad_token_id\": 50256,\n",
      "  \"resid_pdrop\": 0.1,\n",
      "  \"summary_activation\": null,\n",
      "  \"summary_first_dropout\": 0.1,\n",
      "  \"summary_proj_to_labels\": true,\n",
      "  \"summary_type\": \"cls_index\",\n",
      "  \"summary_use_proj\": true,\n",
      "  \"task_specific_params\": {\n",
      "    \"text-generation\": {\n",
      "      \"do_sample\": true,\n",
      "      \"max_length\": 50\n",
      "    }\n",
      "  },\n",
      "  \"vocab_size\": 50257\n",
      "}\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:13:20.513489 34540 filelock.py:274] Lock 2936530781384 acquired on C:\\Users\\bokhy/.cache\\torch\\transformers\\132dec44f9ced4b20f1b1c88a426b1d3dab5ba9e5f24a82541833dae44d5b8db.afd2261c07481427cd087f622388c2c086be9c62875f5945922c7adb2239b63a.h5.lock\n",
      "I0605 23:13:20.514487 34540 file_utils.py:436] https://cdn.huggingface.co/gpt2-tf_model.h5 not found in cache or force_download set to True, downloading to C:\\Users\\bokhy\\.cache\\torch\\transformers\\tmpo6pl4cj8\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1f550151e5e4999bd541c1d04e0c665",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=497933648.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0605 23:15:38.321964 34540 file_utils.py:440] storing https://cdn.huggingface.co/gpt2-tf_model.h5 in cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\132dec44f9ced4b20f1b1c88a426b1d3dab5ba9e5f24a82541833dae44d5b8db.afd2261c07481427cd087f622388c2c086be9c62875f5945922c7adb2239b63a.h5\n",
      "I0605 23:15:38.324956 34540 file_utils.py:443] creating metadata file for C:\\Users\\bokhy/.cache\\torch\\transformers\\132dec44f9ced4b20f1b1c88a426b1d3dab5ba9e5f24a82541833dae44d5b8db.afd2261c07481427cd087f622388c2c086be9c62875f5945922c7adb2239b63a.h5\n",
      "I0605 23:15:38.326950 34540 filelock.py:318] Lock 2936530781384 released on C:\\Users\\bokhy/.cache\\torch\\transformers\\132dec44f9ced4b20f1b1c88a426b1d3dab5ba9e5f24a82541833dae44d5b8db.afd2261c07481427cd087f622388c2c086be9c62875f5945922c7adb2239b63a.h5.lock\n",
      "I0605 23:15:38.327948 34540 modeling_tf_utils.py:393] loading weights file https://cdn.huggingface.co/gpt2-tf_model.h5 from cache at C:\\Users\\bokhy/.cache\\torch\\transformers\\132dec44f9ced4b20f1b1c88a426b1d3dab5ba9e5f24a82541833dae44d5b8db.afd2261c07481427cd087f622388c2c086be9c62875f5945922c7adb2239b63a.h5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<transformers.modeling_tf_gpt2.TFGPT2LMHeadModel at 0x2abb7164e88>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "print(tokenizer)\n",
    "\n",
    "# add the EOS token as PAD token to avoid warnings\n",
    "model = TFGPT2LMHeadModel.from_pretrained(\"gpt2\", pad_token_id=tokenizer.eos_token_id)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. **Greedy Search**\n",
    "\n",
    "Greedy search simply selects the word with the highest probability as its next word\n",
    "\n",
    "![Greedy Search](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/greedy_search.png)\n",
    "\n",
    "Starting from the word $\\text{\"The\"}$, the algorithm \n",
    "greedily chooses the next word of highest probability $\\text{\"nice\"}$ and so on, so that the final generated word sequence is $\\text{\"The\", \"nice\", \"woman\"}$ having an overall probability of $0.5 \\times 0.4 = 0.2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, Let's generate word sequences using GPT2 on the context $(\\text{\"I\", \"enjoy\", \"walking\", \"with\", \"my\", \"cute\", \"dog\"})$. Let's see how greedy search can be used in `transformers` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with my dog. I'm not sure if I'll ever be able to walk with my dog.\n",
      "\n",
      "I'm not sure if I'll\n"
     ]
    }
   ],
   "source": [
    "# encode context the generation is conditioned on\n",
    "input_ids = tokenizer.encode('I enjoy walking with my cute dog', return_tensors='tf')\n",
    "\n",
    "# generate text until the output length (which includes the context length) reaches 50\n",
    "greedy_output = model.generate(input_ids, max_length=50)\n",
    "\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "print(tokenizer.decode(greedy_output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright! The generated words following the context looks good, but the model quickly starts repeating itself"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The major drawback of greedy search though is that it misses high probability words hidden behind a low probability word as can be seen in our sketch above:\n",
    "\n",
    "The word \\text{\"has\"}\"has\" with its high conditional probability of 0.90.9 is hidden behind the word \\text{\"dog\"}\"dog\", which has only the second-highest conditional probability, so that greedy search misses the word sequence \\text{\"The\"}, \\text{\"dog\"}, \\text{\"has\"}\"The\",\"dog\",\"has\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. **Beam search**\n",
    "\n",
    "Beam search reduces the risk of missing hidden high probability word sequences by keeping the most likely `num_beams` of hypotheses at each time step and eventually choosing the hypothesis that has the overall highest probability. Let's illustrate with `num_beams=2`:\n",
    "\n",
    "![Beam search](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/beam_search.png)\n",
    "\n",
    "At time step $1$, besides the most likely hypothesis $\\text{\"The\", \"woman\"}$, beam search also keeps track of the second most likely one $\\text{\"The\", \"dog\"}$. At time step $2$, beam search finds that the word sequence $\\text{\"The\", \"dog\", \"has\"}$ has with $0.36$ a higher probability than $\\text{\"The\", \"nice\", \"woman\"}$, which has $0.2$. Great, it has found the most likely word sequence in our toy example! \n",
    "\n",
    "Beam search will always find an output sequence with higher probability than greedy search, but is not guaranteed to find the most likely output. \n",
    "\n",
    "Let's see how beam search can be used in `transformers`. We set `num_beams > 1` and `early_stopping=True` so that generation is finished when all beam hypotheses reached the EOS token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to take a break\n"
     ]
    }
   ],
   "source": [
    "# activate beam search and early_stopping\n",
    "beam_output = model.generate(\n",
    "    input_ids, \n",
    "    max_length=50, \n",
    "    num_beams=5, \n",
    "    no_repeat_ngram_size=2, # set no_repeat_ngram_size to 2, in order to remove repetitions.\n",
    "    early_stopping=True\n",
    ")\n",
    "\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "print(tokenizer.decode(beam_output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For \"no_repeat_ngram_size\", it have to be used with care. An article generated about the city \"New York\" should not use a 2-gram penalty or otherwise, the name of the city would only appear once in the whole text!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another important feature about beam search is that we can compare the top beams after generation and choose the generated beam that fits our purpose best.\n",
    "\n",
    "In transformers, we simply set the parameter num_return_sequences to the number of highest scoring beams that should be returned."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make sure though that num_return_sequences <= num_beams!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "0: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to take a break\n",
      "1: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to get back to\n",
      "2: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with her again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to take a break\n",
      "3: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with her again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to get back to\n",
      "4: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.\n",
      "\n",
      "I've been thinking about this for a while now, and I think it's time for me to take a step\n"
     ]
    }
   ],
   "source": [
    "# set return_num_sequences > 1\n",
    "beam_outputs = model.generate(\n",
    "    input_ids, \n",
    "    max_length=50, \n",
    "    num_beams=5, \n",
    "    no_repeat_ngram_size=2, \n",
    "    num_return_sequences=5, \n",
    "    early_stopping=True\n",
    ")\n",
    "\n",
    "# now we have 3 output sequences\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "for i, beam_output in enumerate(beam_outputs):\n",
    "    print(\"{}: {}\".format(i, tokenizer.decode(beam_output, skip_special_tokens=True)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Findings:\n",
    "1. Beam search can work very well in tasks where the length of the desired generation is more or less predictable\n",
    "2. High quality human language does not follow a distribution of high probability next words, but Beam search still has repetitive generation of texts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. **Top-K Sampling**\n",
    "\n",
    "[Fan et. al (2018)](https://arxiv.org/pdf/1805.04833.pdf) introduced a simple, but very powerful sampling scheme, called ***Top-K*** sampling. In *Top-K* sampling, the *K* most likely next words are filtered and the probability mass is redistributed among only those *K* next words. \n",
    "GPT2 adopted this sampling scheme, which was one of the reasons for its success in story generation. \n",
    "\n",
    "We extend the range of words used for both sampling steps in the example above from 3 words to 10 words to better illustrate *Top-K* sampling.\n",
    "\n",
    "![top_k_sampling](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/top_k_sampling.png)\n",
    "\n",
    "Having set $K = 6$, in both sampling steps we limit our sampling pool to 6 words. While the 6 most likely words, defined as $V_{\\text{top-K}}$ encompass only *ca.* two-thirds of the whole probability mass in the first step, it includes almost all of the probability mass in the second step. Nevertheless, we see that it successfully eliminates the rather weird candidates $\\text{\"not\", \"the\", \"small\", \"told\"}$ \n",
    "in the second sampling step.\n",
    "\n",
    "\n",
    "Let's see how *Top-K* can be used in the library by setting `top_k=50`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "I enjoy walking with my cute dog so much,\" says my cat, Lulu.\n",
      "\n",
      "Lulu, now with her five-year-old son and a long-distance friend, is the new mom to Pemberton. She and her\n"
     ]
    }
   ],
   "source": [
    "tf.random.set_seed(623)\n",
    "\n",
    "# set top_k to 50\n",
    "sample_output = model.generate(\n",
    "    input_ids, \n",
    "    do_sample=True, \n",
    "    max_length=50, \n",
    "    top_k=50\n",
    ")\n",
    "\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "print(tokenizer.decode(sample_output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "it does not dynamically adapt the number of words that are filtered from the next word probability distribution  P(w|w1:t−1) . This can be problematic as some words might be sampled from a very sharp distribution (distribution on the right in the graph above), whereas others from a much more flat distribution (distribution on the left in the graph above).\n",
    "\n",
    "In step  t=1 , Top-K eliminates the possibility to sample  \"people\", \"big\", \"house\", \"cat\" , which seem like reasonable candidates. On the other hand, in step  t=2  the method includes the arguably ill-fitted words  \"down\", \"a\"  in the sample pool of words. Thus, limiting the sample pool to a fixed size K could endanger the model to produce gibberish for sharp distributions and limit the model's creativity for flat distribution. This intuition led Ari Holtzman et al. (2019) to create Top-p- or nucleus-sampling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. **Top-p (nucleus) sampling**\n",
    "\n",
    "Instead of sampling only from the most likely *K* words, in *Top-p* sampling chooses from the smallest possible set of words whose cumulative probability exceeds the probability *p*. The probability mass is then redistributed among this set of words. This way, the size of the set of words (*a.k.a* the number of words in the set) can dynamically increase and decrease according to the next word's probability distribution. Ok, that was very wordy, let's visualize.\n",
    "\n",
    "![top_p_sampling](https://github.com/patrickvonplaten/scientific_images/blob/master/top_p_sampling.png?raw=true)\n",
    "\n",
    "Having set $p=0.92$, *Top-p* sampling picks the *minimum* number of words to exceed together $p=92\\%$ of the probability mass, defined as $V_{\\text{top-p}}$. In the first example, this included the 9 most likely words, whereas it only has to pick the top 3 words in the second example to exceed 92%. Quite simple actually! It can be seen that it keeps a wide range of words where the next word is arguably less predictable, *e.g.* $P(w | \\text{\"The\"})$, and only a few words when the next word seems more predictable, *e.g.* $P(w | \\text{\"The\", \"car\"})$.\n",
    "\n",
    "Alright, time to check it out in `transformers`!\n",
    "We activate *Top-p* sampling by setting `0 < top_p < 1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "0: I enjoy walking with my cute dog but this dog gets into really bad situations as much. So I can't walk and get into really good positions for this dog. I am very happy and will be more than happy to take my dog out to see\n",
      "1: I enjoy walking with my cute dog, but I'm really into my cat, too,\" he said.\n",
      "\n",
      "As for whether he'd consider an end to his job, he said that, \"no, I think I'd have to work again\n",
      "2: I enjoy walking with my cute dog and eating lunch on a Tuesday. I love having a good time with my dogs and having a good time with my family. I'm really lucky to have my dogs with me when I have two dogs at home,\n"
     ]
    }
   ],
   "source": [
    "tf.random.set_seed(623)\n",
    "\n",
    "# deactivate top_k sampling and sample only from 92% most likely words\n",
    "# set top_k = 50 and set top_p = 0.95 and num_return_sequences = 3\n",
    "sample_outputs = model.generate(\n",
    "    input_ids,\n",
    "    do_sample=True, \n",
    "    max_length=50, \n",
    "    top_k=50, \n",
    "    top_p=0.95, \n",
    "    num_return_sequences=3 # to get multiple independently sampled outputs, set num_return_sequences > 1:\n",
    ")\n",
    "\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "for i, sample_output in enumerate(sample_outputs):\n",
    "    print(\"{}: {}\".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This sounds like it could have been written by a human!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. top-p and top-K sampling seem to produce more fluent text than traditional greedy - and beam search on open-ended language generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. According to human evaluations, beam search can generate more fluent text than Top-p sampling"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}