decoding_api.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXLA5InzXydn"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RuRlpLL-X0R_"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2X-XaMSVcLua"
      },
      "source": [
        "# Decoding API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYEwGTeCXnnX"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/tfmodels/nlp/decoding_api\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/models/blob/master/docs/nlp/decoding_api.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/models/blob/master/docs/nlp/decoding_api.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/models/docs/nlp/decoding_api.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fsACVQpVSifi"
      },
      "source": [
        "### Install the TensorFlow Model Garden pip package\n",
        "\n",
        "*  `tf-models-official` is the stable Model Garden package. Note that it may not include the latest changes in the `tensorflow_models` github repo. To include latest changes, you may install `tf-models-nightly`,\n",
        "which is the nightly Model Garden package created daily automatically.\n",
        "*  pip will install all models and dependencies automatically."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G4BhAu01HZcM"
      },
      "outputs": [],
      "source": [
        "!pip uninstall -y opencv-python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2j-xhrsVQOQT"
      },
      "outputs": [],
      "source": [
        "!pip install tf-models-official"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BjP7zwxmskpY"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow_models import nlp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T92ccAzlnGqh"
      },
      "outputs": [],
      "source": [
        "def length_norm(length, dtype):\n",
        "  \"\"\"Return length normalization factor.\"\"\"\n",
        "  return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0AWgyo-IQ5sP"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This API provides an interface to experiment with different decoding strategies used for auto-regressive models.\n",
        "\n",
        "1. The following sampling strategies are provided in sampling_module.py, which inherits from the base Decoding class:\n",
        "  *   [top_p](https://arxiv.org/abs/1904.09751) : [github](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/ops/sampling_module.py#L65) \n",
        "\n",
        "      This implementation chooses the most probable logits with cumulative probabilities up to top_p.\n",
        "\n",
        "  *   [top_k](https://arxiv.org/pdf/1805.04833.pdf) : [github](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/ops/sampling_module.py#L48)\n",
        "\n",
        "      At each timestep, this implementation samples from top-k logits based on their probability distribution\n",
        "\n",
        "  *   Greedy : [github](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/ops/sampling_module.py#L26)\n",
        "\n",
        "      This implementation returns the top logits based on probabilities.\n",
        "\n",
        "2. Beam search is provided in beam_search.py. [github](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/ops/beam_search.py)\n",
        "\n",
        "      This implementation reduces the risk of missing hidden high probability logits by keeping the most likely num_beams of logits at each time step and eventually choosing the logits that has the overall highest probability."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfOj7oaBRQnS"
      },
      "source": [
        "## Initialize Sampling Module in TF-NLP.\n",
        "\n",
        "\n",
        "\u003e **symbols_to_logits_fn** : This is a closure implemented by the users of the API. The input to this closure will be  \n",
        "```\n",
        "Args:\n",
        "  1] ids [batch_size, .. (index + 1 or 1 if padded_decode is True)],\n",
        "  2] index [scalar] : current decoded step,\n",
        "  3] cache [nested dictionary of tensors].\n",
        "Returns:\n",
        "  1] tensor for next-step logits [batch_size, vocab]\n",
        "  2] the updated_cache [nested dictionary of tensors].\n",
        "```\n",
        "This closure calls the model to predict the logits for the 'index+1' step. The cache is used for faster decoding.\n",
        "Here is a [reference](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/ops/beam_search_test.py#L88) implementation for the above closure.\n",
        "\n",
        "\n",
        "\u003e **length_normalization_fn** : Closure for returning length normalization parameter.\n",
        "```\n",
        "Args: \n",
        "  1] length : scalar for decoded step index.\n",
        "  2] dtype : data-type of output tensor\n",
        "Returns:\n",
        "  1] value of length normalization factor.\n",
        "Example :\n",
        "  def _length_norm(length, dtype):\n",
        "    return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)\n",
        "```\n",
        "\n",
        "\u003e **vocab_size** : Output vocabulary size.\n",
        "\n",
        "\u003e **max_decode_length** : Scalar for total number of decoding steps.\n",
        "\n",
        "\u003e **eos_id** : Decoding will stop if all output decoded ids in the batch have this ID.\n",
        "\n",
        "\u003e **padded_decode** : Set this to True if running on TPU. Tensors are padded to max_decoding_length if this is True.\n",
        "\n",
        "\u003e **top_k** : top_k is enabled if this value is \u003e 1.\n",
        "\n",
        "\u003e **top_p** : top_p is enabled if this value is \u003e 0 and \u003c 1.0\n",
        "\n",
        "\u003e **sampling_temperature** : This is used to re-estimate the softmax output. Temperature skews the distribution towards high-probability tokens and lowers the mass in the tail distribution. Value has to be positive. Low temperature is equivalent to greedy and makes the distribution sharper, while high temperature makes it flatter.\n",
        "\n",
        "\u003e **enable_greedy** : By default, this is true and greedy decoding is enabled.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lV1RRp6ihnGX"
      },
      "source": [
        "## Initialize the Model Hyper-parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eTsGp2gaKLdE"
      },
      "outputs": [],
      "source": [
        "params = {\n",
        "    'num_heads': 2,\n",
        "    'num_layers': 2,\n",
        "    'batch_size': 2,\n",
        "    'n_dims': 256,\n",
        "    'max_decode_length': 4}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CYXkoplAij01"
      },
      "source": [
        "## Initialize cache. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGvmd0_dRFYI"
      },
      "source": [
        "In auto-regressive architectures like Transformer based [Encoder-Decoder](https://arxiv.org/abs/1706.03762) models, \n",
        "Cache is used for fast sequential decoding.\n",
        "It is a nested dictionary storing pre-computed hidden-states (key and values in the self-attention blocks and  the cross-attention blocks) for every layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D6kfZOOKgkm1"
      },
      "outputs": [],
      "source": [
        "cache = {\n",
        "    'layer_%d' % layer: {\n",
        "        'k': tf.zeros(\n",
        "            shape=[params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims'] // params['num_heads']],\n",
        "            dtype=tf.float32),\n",
        "        'v': tf.zeros(\n",
        "            shape=[params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims'] // params['num_heads']],\n",
        "            dtype=tf.float32)\n",
        "        } for layer in range(params['num_layers'])\n",
        "    }\n",
        "print(\"cache value shape for layer 1 :\", cache['layer_1']['k'].shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "syl7I5nURPgW"
      },
      "source": [
        "### Create model_fn\n",
        "  In practice, this will be replaced by an actual model implementation such as [here](https://github.com/tensorflow/models/blob/master/official/nlp/transformer/transformer.py#L236)\n",
        "```\n",
        "Args:\n",
        "i : Step that is being decoded.\n",
        "Returns:\n",
        "  logit probabilities of size [batch_size, 1, vocab_size]\n",
        "```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AhzSkRisRdB6"
      },
      "outputs": [],
      "source": [
        "probabilities = tf.constant([[[0.3, 0.4, 0.3], [0.3, 0.3, 0.4],\n",
        "                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],\n",
        "                            [[0.2, 0.5, 0.3], [0.2, 0.7, 0.1],\n",
        "                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]])\n",
        "def model_fn(i):\n",
        "  return probabilities[:, i, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FAJ4CpbfVdjr"
      },
      "outputs": [],
      "source": [
        "def _symbols_to_logits_fn():\n",
        "  \"\"\"Calculates logits of the next tokens.\"\"\"\n",
        "  def symbols_to_logits_fn(ids, i, temp_cache):\n",
        "    del ids\n",
        "    logits = tf.cast(tf.math.log(model_fn(i)), tf.float32)\n",
        "    return logits, temp_cache\n",
        "  return symbols_to_logits_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R_tV3jyWVL47"
      },
      "source": [
        "## Greedy \n",
        "Greedy decoding selects the token id with the highest probability as its next id: $id_t = argmax_{w}P(id | id_{1:t-1})$ at each timestep $t$. The following sketch shows greedy decoding. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aGt9idSkVQEJ"
      },
      "outputs": [],
      "source": [
        "greedy_obj = sampling_module.SamplingModule(\n",
        "    length_normalization_fn=None,\n",
        "    dtype=tf.float32,\n",
        "    symbols_to_logits_fn=_symbols_to_logits_fn(),\n",
        "    vocab_size=3,\n",
        "    max_decode_length=params['max_decode_length'],\n",
        "    eos_id=10,\n",
        "    padded_decode=False)\n",
        "ids, _ = greedy_obj.generate(\n",
        "    initial_ids=tf.constant([9, 1]), initial_cache=cache)\n",
        "print(\"Greedy Decoded Ids:\", ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s4pTTsQXVz5O"
      },
      "source": [
        "## top_k sampling\n",
        "In *Top-K* sampling, the *K* most likely next token ids are filtered and the probability mass is redistributed among only those *K* ids. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pCLWIn6GV5_G"
      },
      "outputs": [],
      "source": [
        "top_k_obj = sampling_module.SamplingModule(\n",
        "    length_normalization_fn=length_norm,\n",
        "    dtype=tf.float32,\n",
        "    symbols_to_logits_fn=_symbols_to_logits_fn(),\n",
        "    vocab_size=3,\n",
        "    max_decode_length=params['max_decode_length'],\n",
        "    eos_id=10,\n",
        "    sample_temperature=tf.constant(1.0),\n",
        "    top_k=tf.constant(3),\n",
        "    padded_decode=False,\n",
        "    enable_greedy=False)\n",
        "ids, _ = top_k_obj.generate(\n",
        "    initial_ids=tf.constant([9, 1]), initial_cache=cache)\n",
        "print(\"top-k sampled Ids:\", ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jp3G-eE_WI4Y"
      },
      "source": [
        "## top_p sampling\n",
        "Instead of sampling only from the most likely *K* token ids, in *Top-p* sampling chooses from the smallest possible set of ids whose cumulative probability exceeds the probability *p*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rEGdIWcuWILO"
      },
      "outputs": [],
      "source": [
        "top_p_obj = sampling_module.SamplingModule(\n",
        "    length_normalization_fn=length_norm,\n",
        "    dtype=tf.float32,\n",
        "    symbols_to_logits_fn=_symbols_to_logits_fn(),\n",
        "    vocab_size=3,\n",
        "    max_decode_length=params['max_decode_length'],\n",
        "    eos_id=10,\n",
        "    sample_temperature=tf.constant(1.0),\n",
        "    top_p=tf.constant(0.9),\n",
        "    padded_decode=False,\n",
        "    enable_greedy=False)\n",
        "ids, _ = top_p_obj.generate(\n",
        "    initial_ids=tf.constant([9, 1]), initial_cache=cache)\n",
        "print(\"top-p sampled Ids:\", ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2hcuyJ2VWjDz"
      },
      "source": [
        "## Beam search decoding\n",
        "Beam search reduces the risk of missing hidden high probability token ids by keeping the most likely num_beams of hypotheses at each time step and eventually choosing the hypothesis that has the overall highest probability. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cJ3WzvSrWmSA"
      },
      "outputs": [],
      "source": [
        "beam_size = 2\n",
        "params['batch_size'] = 1\n",
        "beam_cache = {\n",
        "    'layer_%d' % layer: {\n",
        "        'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32),\n",
        "        'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32)\n",
        "        } for layer in range(params['num_layers'])\n",
        "    }\n",
        "print(\"cache key shape for layer 1 :\", beam_cache['layer_1']['k'].shape)\n",
        "ids, _ = beam_search.sequence_beam_search(\n",
        "    symbols_to_logits_fn=_symbols_to_logits_fn(),\n",
        "    initial_ids=tf.constant([9], tf.int32),\n",
        "    initial_cache=beam_cache,\n",
        "    vocab_size=3,\n",
        "    beam_size=beam_size,\n",
        "    alpha=0.6,\n",
        "    max_decode_length=params['max_decode_length'],\n",
        "    eos_id=10,\n",
        "    padded_decode=False,\n",
        "    dtype=tf.float32)\n",
        "print(\"Beam search ids:\", ids)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "decoding_api_in_tf_nlp.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}