notebooks/02_how_to_generate.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "02_how_to_generate.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/huggingface/blog/blob/main/notebooks/02_how_to_generate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vp3XPuaTu9jl",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "# How to generate text: using different decoding methods for language generation with Transformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KxLvv6UaPa33",
        "colab_type": "text"
      },
      "source": [
        "### **Introduction**\n",
        "\n",
        "In recent years, there has been an increasing interest in open-ended language generation thanks to the rise of large transformer-based language models trained on millions of webpages, such as OpenAI's famous [GPT2 model](https://openai.com/blog/better-language-models/). The results on conditioned open-ended language generation are impressive, e.g. [GPT2 on unicorns](https://openai.com/blog/better-language-models/#samples), [XLNet](https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e), [Controlled language with CTRL](https://blog.einstein.ai/introducing-a-conditional-transformer-language-model-for-controllable-generation/). Besides the improved transformer architecture and massive unsupervised training data, **better decoding methods** have also played an important role. \n",
        "\n",
        "This blog post gives a brief overview of different decoding strategies and more importantly shows how *you* can implement them with very little effort using the popular `transformers` library! \n",
        "\n",
        "All of the following functionalities can be used for **auto-regressive** language generation ([here](http://jalammar.github.io/illustrated-gpt2/) a refresher). 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})$.\n",
        "\n",
        "\n",
        "Auto-regressive language generation is now available for `GPT2`, `XLNet`, `OpenAi-GPT`, `CTRL`, `TransfoXL`, `XLM`, `Bart`, `T5` in both PyTorch and Tensorflow >= 2.0!\n",
        "\n",
        "We will give a tour of the currently most prominent decoding methods, mainly *Greedy search*, *Beam search*, *Top-K sampling* and *Top-p sampling*.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Si4GyYhOQMzi",
        "colab_type": "text"
      },
      "source": [
        "Let's quickly install transformers and load the model. We will use GPT2 in Tensorflow 2.1 for demonstration, but the API is 1-to-1 the same for PyTorch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XbzZ_IVTtoQe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install -q git+https://github.com/huggingface/transformers.git\n",
        "!pip install -q tensorflow==2.1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ue2kOQhXTAMU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "from transformers import TFGPT2LMHeadModel, GPT2Tokenizer\n",
        "\n",
        "\n",
        "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\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)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a8Y7cgu9ohXP",
        "colab_type": "text"
      },
      "source": [
        "### **Greedy Search**\n",
        "\n",
        "Greedy search simply selects the word with the highest probability as its next word: $w_t = argmax_{w}P(w | w_{1:t-1})$ at each timestep $t$. The following sketch shows greedy search. \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$.\n",
        "\n",
        "In the following we will 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",
      "metadata": {
        "id": "OWLd_J6lXz_t",
        "colab_type": "code",
        "outputId": "3b9dfd1e-21e6-44f4-f27f-8e975010f9af",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "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))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BBn1ePmJvhrl",
        "colab_type": "text"
      },
      "source": [
        "Alright! We have generated our first short text with GPT2 😊. The generated words following the context are reasonable, but the model quickly starts repeating itself! This is a very common problem in language generation in general and seems to be even more so in greedy and beam search - check out [Vijayakumar et al., 2016](https://arxiv.org/abs/1610.02424) and [Shao et al., 2017](https://arxiv.org/abs/1701.03185).\n",
        "\n",
        "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\"}$ with its high conditional probability of $0.9$ is hidden behind the word $\\text{\"dog\"}$, which has only the second-highest conditional probability, so that greedy search misses the word sequence $\\text{\"The\"}, \\text{\"dog\"}, \\text{\"has\"}$.\n",
        "\n",
        "Thankfully, we have beam search to alleviate this problem!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g8DnXZ1WiuNd",
        "colab_type": "text"
      },
      "source": [
        "### **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\", \"nice\"}$, 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",
      "metadata": {
        "id": "R1R5kx30Ynej",
        "colab_type": "code",
        "outputId": "574f068b-f418-48b5-8334-8451d2221032",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "# activate beam search and early_stopping\n",
        "beam_output = model.generate(\n",
        "    input_ids,  \n",
        "    max_length=50, \n",
        "    num_beams=5, \n",
        "    early_stopping=True\n",
        ")\n",
        "\n",
        "print(\"Output:\\n\" + 100 * '-')\n",
        "print(tokenizer.decode(beam_output[0], skip_special_tokens=True))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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'm not sure if I'll ever be able to walk with him again. I'm not sure if I'll\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZ6xs-KLi9jT",
        "colab_type": "text"
      },
      "source": [
        "While the result is arguably more fluent, the output still includes repetitions of the same word sequences.  \n",
        "A simple remedy is to introduce *n-grams* (*a.k.a* word sequences of $n$ words) penalties as introduced by [Paulus et al. (2017)](https://arxiv.org/abs/1705.04304) and [Klein et al. (2017)](https://arxiv.org/abs/1701.02810). The most common *n-grams* penalty makes sure that no *n-gram* appears twice by manually setting the probability of next words that could create an already seen *n-gram* to $0$.\n",
        "\n",
        "Let's try it out by setting `no_repeat_ngram_size=2` so that no *2-gram* appears twice:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jy3iVJgfnkMi",
        "colab_type": "code",
        "outputId": "4d3e6511-711a-4594-a715-aaeb6e48e1a9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "# set no_repeat_ngram_size to 2\n",
        "beam_output = model.generate(\n",
        "    input_ids, \n",
        "    max_length=50, \n",
        "    num_beams=5, \n",
        "    no_repeat_ngram_size=2, \n",
        "    early_stopping=True\n",
        ")\n",
        "\n",
        "print(\"Output:\\n\" + 100 * '-')\n",
        "print(tokenizer.decode(beam_output[0], skip_special_tokens=True))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nxsksOGDpmA0",
        "colab_type": "text"
      },
      "source": [
        "Nice, that looks much better! We can see that the repetition does not appear anymore. Nevertheless, *n-gram* penalties 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!\n",
        "\n",
        "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. Make sure though that `num_return_sequences <= num_beams`!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5ClO3VphqGp6",
        "colab_type": "code",
        "outputId": "2296891c-024f-4fd2-9071-bff7c11a3e04",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "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)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HhLKyfdbsjXc",
        "colab_type": "text"
      },
      "source": [
        "As can be seen, the five beam hypotheses are only marginally different to each other - which should not be too surprising when using only 5 beams.\n",
        "\n",
        "In open-ended generation, a couple of reasons have recently been brought forward why beam search might not be the best possible option:\n",
        "\n",
        "- Beam search can work very well in tasks where the length of the desired generation is more or less predictable as in machine translation or summarization - see [Murray et al. (2018)](https://arxiv.org/abs/1808.10006) and [Yang et al. (2018)](https://arxiv.org/abs/1808.09582). But this is not the case for open-ended generation where the desired output length can vary greatly, e.g. dialog and story generation.\n",
        "\n",
        "- We have seen that beam search heavily suffers from repetitive generation. This is especially hard to control with *n-gram*- or other penalties in story generation since finding a good trade-off between forced \"no-repetition\" and repeating cycles of identical *n-grams* requires a lot of finetuning.\n",
        "\n",
        "- As argued in [Ari Holtzman et al. (2019)](https://arxiv.org/abs/1904.09751), high quality human language does not follow a distribution of high probability next words. In other words, as humans, we want generated text to surprise us and not to be boring/predictable. The authors show this nicely by plotting the probability, a model would give to human text vs. what beam search does.\n",
        "\n",
        "![alt text](https://blog.fastforwardlabs.com/images/2019/05/Screen_Shot_2019_05_08_at_3_06_36_PM-1557342561886.png)\n",
        "\n",
        "\n",
        "So let's stop being boring and introduce some randomness 🤪."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XbbIyK84wHq6",
        "colab_type": "text"
      },
      "source": [
        "### **Sampling**\n",
        "\n",
        "In its most basic form, sampling means randomly picking the next word $w_t$ according to its conditional probability distribution:\n",
        "\n",
        "$$w_t \\sim P(w|w_{1:t-1})$$\n",
        "\n",
        "Taking the example from above, the following graphic visualizes language generation when sampling.\n",
        "\n",
        "![vanilla_sampling](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/sampling_search.png)\n",
        "\n",
        "It becomes obvious that language generation using sampling is not *deterministic* anymore. The word \n",
        "$\\text{\"car\"}$ is sampled from the conditioned probability distribution $P(w | \\text{\"The\"})$, followed by sampling $\\text{\"drives\"}$ from $P(w | \\text{\"The\"}, \\text{\"car\"})$.\n",
        "\n",
        "In `transformers`, we set `do_sample=True` and deactivate *Top-K* sampling (more on this later) via `top_k=0`. In the following, we will fix `random_seed=0` for illustration purposes. Feel free to change the `random_seed` to play around with the model.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "1b78d191-15f6-4cbe-e2b1-23c77366fc21",
        "id": "aRAz4D-Ks0_4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        }
      },
      "source": [
        "# set seed to reproduce results. Feel free to change the seed though to get different results\n",
        "tf.random.set_seed(0)\n",
        "\n",
        "# activate sampling and deactivate top_k by setting top_k sampling to 0\n",
        "sample_output = model.generate(\n",
        "    input_ids, \n",
        "    do_sample=True, \n",
        "    max_length=50, \n",
        "    top_k=0\n",
        ")\n",
        "\n",
        "print(\"Output:\\n\" + 100 * '-')\n",
        "print(tokenizer.decode(sample_output[0], skip_special_tokens=True))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Output:\n",
            "----------------------------------------------------------------------------------------------------\n",
            "I enjoy walking with my cute dog. He just gave me a whole new hand sense.\"\n",
            "\n",
            "But it seems that the dogs have learned a lot from teasing at the local batte harness once they take on the outside.\n",
            "\n",
            "\"I take\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mQHuo911wfT-",
        "colab_type": "text"
      },
      "source": [
        "Interesting! The text seems alright - but when taking a closer look, it is not very coherent. the *3-grams* *new hand sense* and *local batte harness* are very weird and don't sound like they were written by a human. That is the big problem when sampling word sequences: The models often generate incoherent gibberish, *cf.* [Ari Holtzman et al. (2019)](https://arxiv.org/abs/1904.09751).\n",
        "\n",
        "A trick is to make the distribution $P(w|w_{1:t-1})$ sharper (increasing the likelihood of high probability words and decreasing the likelihood of low probability words) by lowering the so-called `temperature` of the [softmax](https://en.wikipedia.org/wiki/Softmax_function#Smooth_arg_max). \n",
        "\n",
        "An illustration of applying temperature to our example from above could look as follows.\n",
        "\n",
        "![top_p_sampling](https://github.com/patrickvonplaten/scientific_images/blob/master/sampling_search_with_temp.png?raw=true)\n",
        "\n",
        "The conditional next word distribution of step $t=1$ becomes much sharper leaving almost no chance for word $\\text{\"car\"}$ to be selected.\n",
        "\n",
        "\n",
        "Let's see how we can cool down the distribution in the library by setting `temperature=0.7`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WgJredc-0j0Z",
        "colab_type": "code",
        "outputId": "a4e79355-8e3c-4788-fa21-c4e28bf61c5b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "# set seed to reproduce results. Feel free to change the seed though to get different results\n",
        "tf.random.set_seed(0)\n",
        "\n",
        "# use temperature to decrease the sensitivity to low probability candidates\n",
        "sample_output = model.generate(\n",
        "    input_ids, \n",
        "    do_sample=True, \n",
        "    max_length=50, \n",
        "    top_k=0, \n",
        "    temperature=0.7\n",
        ")\n",
        "\n",
        "print(\"Output:\\n\" + 100 * '-')\n",
        "print(tokenizer.decode(sample_output[0], skip_special_tokens=True))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Output:\n",
            "----------------------------------------------------------------------------------------------------\n",
            "I enjoy walking with my cute dog, but I don't like to be at home too much. I also find it a bit weird when I'm out shopping. I am always away from my house a lot, but I do have a few friends\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kzGuu24hZZnq",
        "colab_type": "text"
      },
      "source": [
        "OK. There are less weird n-grams and the output is a bit more coherent now! While applying temperature can make a distribution less random, in its limit, when setting `temperature` $ \\to 0$, temperature scaled sampling becomes equal to greedy decoding and will suffer from the same problems as before. \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "binNTroyzQBu",
        "colab_type": "text"
      },
      "source": [
        "### **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",
      "metadata": {
        "id": "HBtDOdD0wx3l",
        "colab_type": "code",
        "outputId": "cfc97fac-0956-42ee-a6e5-cad14fc942d3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 156
        }
      },
      "source": [
        "# set seed to reproduce results. Feel free to change the seed though to get different results\n",
        "tf.random.set_seed(0)\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))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Output:\n",
            "----------------------------------------------------------------------------------------------------\n",
            "I enjoy walking with my cute dog. It's so good to have an environment where your dog is available to share with you and we'll be taking care of you.\n",
            "\n",
            "We hope you'll find this story interesting!\n",
            "\n",
            "I am from\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y77H5m4ZmhEX",
        "colab_type": "text"
      },
      "source": [
        "Not bad at all! The text is arguably the most *human-sounding* text so far. \n",
        "One concern though with *Top-K* sampling is that it does not dynamically adapt the number of words that are filtered from the next word probability distribution $P(w|w_{1:t-1})$.\n",
        "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 \n",
        "sample $\\text{\"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 $\\text{\"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.\n",
        "This intuition led [Ari Holtzman et al. (2019)](https://arxiv.org/abs/1904.09751) to create ***Top-p***- or ***nucleus***-sampling. \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ki9LAaexzV3H",
        "colab_type": "text"
      },
      "source": [
        "### **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",
      "metadata": {
        "id": "EvwIc7YAx77F",
        "colab_type": "code",
        "outputId": "57e2b785-5dcb-4e06-9869-078b758b6a82",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        }
      },
      "source": [
        "# set seed to reproduce results. Feel free to change the seed though to get different results\n",
        "tf.random.set_seed(0)\n",
        "\n",
        "# deactivate top_k sampling and sample only from 92% most likely words\n",
        "sample_output = model.generate(\n",
        "    input_ids, \n",
        "    do_sample=True, \n",
        "    max_length=50, \n",
        "    top_p=0.92, \n",
        "    top_k=0\n",
        ")\n",
        "\n",
        "print(\"Output:\\n\" + 100 * '-')\n",
        "print(tokenizer.decode(sample_output[0], skip_special_tokens=True))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Output:\n",
            "----------------------------------------------------------------------------------------------------\n",
            "I enjoy walking with my cute dog. He will never be the same. I watch him play.\n",
            "\n",
            "\n",
            "Guys, my dog needs a name. Especially if he is found with wings.\n",
            "\n",
            "\n",
            "What was that? I had a lot of\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tn-8gLaR4lat",
        "colab_type": "text"
      },
      "source": [
        "Great, that sounds like it could have been written by a human. Well, maybe not quite yet. \n",
        "\n",
        "While in theory, *Top-p* seems more elegant than *Top-K*, both methods work  well in practice. *Top-p* can also be used in combination with *Top-K*, which can avoid very low ranked words while allowing for some dynamic selection.\n",
        "\n",
        "Finally, to get multiple independently sampled outputs, we can *again* set the parameter `num_return_sequences > 1`: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3kY8P9VG8Gi9",
        "colab_type": "code",
        "outputId": "6103051e-1681-4ab9-a9c1-1fad437c299d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 190
        }
      },
      "source": [
        "# set seed to reproduce results. Feel free to change the seed though to get different results\n",
        "tf.random.set_seed(0)\n",
        "\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\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)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Output:\n",
            "----------------------------------------------------------------------------------------------------\n",
            "0: I enjoy walking with my cute dog. It's so good to have the chance to walk with a dog. But I have this problem with the dog and how he's always looking at us and always trying to make me see that I can do something\n",
            "1: I enjoy walking with my cute dog, she loves taking trips to different places on the planet, even in the desert! The world isn't big enough for us to travel by the bus with our beloved pup, but that's where I find my love\n",
            "2: I enjoy walking with my cute dog and playing with our kids,\" said David J. Smith, director of the Humane Society of the US.\n",
            "\n",
            "\"So as a result, I've got more work in my time,\" he said.\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-vRPfMl88rk0",
        "colab_type": "text"
      },
      "source": [
        "Cool, now you should have all the tools to let your model write your stories with `transformers`!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NsWd7e98Vcs3",
        "colab_type": "text"
      },
      "source": [
        "### **Conclusion**\n",
        "\n",
        "As *ad-hoc* decoding methods, *top-p* and *top-K* sampling seem to produce more fluent text than traditional *greedy* - and *beam* search on open-ended language generation. \n",
        "Recently, there has been more evidence though that the apparent flaws of *greedy* and *beam* search - mainly generating repetitive word sequences - are  caused by the model (especially the way the model is trained), rather than the decoding method, *cf.* [Welleck et al. (2019)](https://arxiv.org/pdf/1908.04319.pdf). Also, as demonstrated in [Welleck et al. (2020)](https://arxiv.org/abs/2002.02492), it looks as *top-K* and *top-p* sampling also suffer from generating repetitive word sequences.\n",
        "\n",
        "In [Welleck et al. (2019)](https://arxiv.org/pdf/1908.04319.pdf), the authors show that according to human evaluations, *beam* search can generate more fluent text than *Top-p* sampling, when adapting the model's training objective.\n",
        "\n",
        "Open-ended language generation is a rapidly evolving field of research and as it is often the case there is no one-size-fits-all method here, so one has to see what works best in one's specific use case.\n",
        "\n",
        "Good thing, that *you* can try out all the different decoding methods in `transfomers` 🤗. \n",
        "\n",
        "That was a short introduction on how to use different decoding methods in `transformers` and recent trends in open-ended language generation. \n",
        "\n",
        "Feedback and questions are very welcome on the [Github repository](https://github.com/huggingface/transformers).\n",
        "\n",
        "For more fun generating stories, please take a look at [Writing with Transformers](https://transformer.huggingface.co).\n",
        "\n",
        "Thanks to everybody, who has contributed to the blog post: Alexander Rush, Julien Chaumand, Thomas Wolf, Victor Sanh, Sam Shleifer, Clément Delangue, Yacine Jernite, Oliver Åstrand and John de Wasseige.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4CYi91h11yd",
        "colab_type": "text"
      },
      "source": [
        "### **Appendix**\n",
        "\n",
        "There are a couple of additional parameters for the `generate` method that were not mentioned above. We will explain them here briefly!\n",
        "\n",
        "- `min_length` can be used to force the model to not produce an EOS token (= not finish the sentence) before `min_length` is reached. This is used quite frequently in summarization, but can be useful in general if the user wants to have longer outputs.\n",
        "- `repetition_penalty` can be used to penalize words that were already generated or belong to the context. It was first introduced by [Kesker et al. (2019)](https://arxiv.org/abs/1909.05858) and is also used in the training objective in [Welleck et al. (2019)](https://arxiv.org/pdf/1908.04319.pdf). It can be quite effective at preventing repetitions, but seems to be very sensitive to different models and use cases, *e.g.* see this [discussion](https://github.com/huggingface/transformers/pull/2303) on Github.\n",
        "\n",
        "- `attention_mask` can be used to mask padded tokens\n",
        "- `pad_token_id`, `bos_token_id`, `eos_token_id`: If the model does not have those tokens by default, the user can manually choose other token ids to represent them.\n",
        "\n",
        "For more information please also look into the `generate` function [docstring](https://huggingface.co/transformers/main_classes/model.html?highlight=generate#transformers.TFPreTrainedModel.generate)."
      ]
    }
  ]
}