<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2021 The TensorFlow Authors.
Step1:  Decoding API
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  Initialize Sampling Module in TF-NLP.
symbols_to_logits_fn 
Step3:  In auto-regressive architectures like Transformer based Encoder-Decoder models, 
Cache is used for fast sequential decoding.
It is a nested dictionary storing pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) for every layer.
Initialize cache.
Step5:  Define closure for length normalization if needed.
This is used for normalizing the final scores of generated sequences and is optional
Step6:  Create model_fn
In practice, this will be replaced by an actual model implementation such as here
Args
Step8:  Initialize symbols_to_logits_fn
Step9:  Greedy
Greedy decoding selects the token id with the highest probability as its next id
Step10:  top_k sampling
In Top-K sampling, the K most likely next token ids are filtered and the probability mass is redistributed among only those K ids.
Step11:  top_p sampling
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.
Step12:  Beam search decoding
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.
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2021 The TensorFlow Authors.
End of explanation
!pip install -q -U "tensorflow-text==2.8.*"
!pip install -q tf-models-official==2.7.0
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from official import nlp
from official.nlp.modeling.ops import sampling_module
from official.nlp.modeling.ops import beam_search
Explanation: Decoding API
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/text/guide/decoding_api"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/text/blob/master/docs/guide/decoding_api.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/text/blob/master/docs/guide/decoding_api.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/text/docs/guide/decoding_api.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>
Overview
In the recent past, there has been a lot of research in language generation with auto-regressive models. In auto-regressive language generation, the probability distribution of token at time step K is dependent on the model's token-predictions till step K-1. For these models, decoding strategies such as Beam search, Greedy, Top-p, and Top-k are critical components of the model and largely influence the style/nature of the generated output token at a given time step K. 
For example, Beam search reduces the risk of missing hidden high probability tokens by
keeping the most likely num_beams of hypotheses at each time step and eventually
choosing the hypothesis that has the overall highest probability.
Murray et al. (2018) and
Yang et al. (2018) show that beam search
works well in Machine Translation tasks.
Both Beam search & Greedy strategies have a possibility of generating
repeating tokens.
Fan et. al (2018) introduced Top-K
sampling, in which the K most likely tokens are filtered and the probability mass
is redistributed among only those K tokens.
Ari Holtzman et. al (2019) introduced
Top-p sampling, which chooses from the smallest possible set of tokens with
cumulative probability that adds upto the probability p. The probability mass is then
redistributed among this set. This way, the size of the set of tokens can
dynamically increase and decrease.
Top-p, Top-k are generally used in tasks such as story-generation.
The Decoding API provides an interface to experiment with different decoding strategies on auto-regressive models.
The following sampling strategies are provided in sampling_module.py, which inherits from the base Decoding class:
top_p : github 
top_k : github
Greedy : github
Beam search is provided in beam_search.py. github
Setup
End of explanation
params = {}
params['num_heads'] = 2
params['num_layers'] = 2
params['batch_size'] = 2
params['n_dims'] = 256
params['max_decode_length'] = 4
Explanation: Initialize Sampling Module in TF-NLP.
symbols_to_logits_fn : Use this closure to call the model to predict the logits for the index+1 step. Inputs and outputs for this closure are as follows:
Args:
  1] ids : Current decoded sequences. int tensor with shape (batch_size, index + 1 or 1 if padded_decode is True)],
  2] index [scalar] : current decoded step,
  3] cache [nested dictionary of tensors] : Only used for faster decoding to store pre-computed attention hidden states for keys and values. More explanation in the cell below.
Returns:
  1] tensor for next-step logits [batch_size, vocab]
  2] the updated_cache [nested dictionary of tensors].
The cache is used for faster decoding.
Here is a reference implementation for the above closure.
length_normalization_fn : Use this closure for returning length normalization parameter.
Args: 
  1] length : scalar for decoded step index.
  2] dtype : data-type of output tensor
Returns:
  1] value of length normalization factor.
vocab_size : Output vocabulary size.
max_decode_length : Scalar for total number of decoding steps.
eos_id : Decoding will stop if all output decoded ids in the batch have this eos_id.
padded_decode : Set this to True if running on TPU. Tensors are padded to max_decoding_length if this is True.
top_k : top_k is enabled if this value is > 1.
top_p : top_p is enabled if this value is > 0 and < 1.0
sampling_temperature : This is used to re-estimate the softmax output. Temperature skews the distribution towards high probability tokens and lowers the mass in tail distribution. Value has to be positive. Low temperature is equivalent to greedy and makes the distribution sharper, while high temperature makes it more flat.
enable_greedy : By default, this is True and greedy decoding is enabled. To experiment with other strategies, please set this to False.
Initialize the Model hyperparameters
End of explanation
cache = {
    'layer_%d' % layer: {
        'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32),
        'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], int(params['n_dims']/params['num_heads'])], dtype=tf.float32)
        } for layer in range(params['num_layers'])
    }
print("cache key shape for layer 1 :", cache['layer_1']['k'].shape)
Explanation: In auto-regressive architectures like Transformer based Encoder-Decoder models, 
Cache is used for fast sequential decoding.
It is a nested dictionary storing pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) for every layer.
Initialize cache.
End of explanation
def length_norm(length, dtype):
  Return length normalization factor.
  return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)
Explanation: Define closure for length normalization if needed.
This is used for normalizing the final scores of generated sequences and is optional
End of explanation
probabilities = tf.constant([[[0.3, 0.4, 0.3], [0.3, 0.3, 0.4],
                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                            [[0.2, 0.5, 0.3], [0.2, 0.7, 0.1],
                              [0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]])
def model_fn(i):
  return probabilities[:, i, :]
Explanation: Create model_fn
In practice, this will be replaced by an actual model implementation such as here
Args:
i : Step that is being decoded.
Returns:
  logit probabilities of size [batch_size, 1, vocab_size]
End of explanation
def _symbols_to_logits_fn():
  Calculates logits of the next tokens.
  def symbols_to_logits_fn(ids, i, temp_cache):
    del ids
    logits = tf.cast(tf.math.log(model_fn(i)), tf.float32)
    return logits, temp_cache
  return symbols_to_logits_fn
Explanation: Initialize symbols_to_logits_fn
End of explanation
greedy_obj = sampling_module.SamplingModule(
    length_normalization_fn=None,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    padded_decode=False)
ids, _ = greedy_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("Greedy Decoded Ids:", ids)
Explanation: Greedy
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.
End of explanation
top_k_obj = sampling_module.SamplingModule(
    length_normalization_fn=length_norm,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    sample_temperature=tf.constant(1.0),
    top_k=tf.constant(3),
    padded_decode=False,
    enable_greedy=False)
ids, _ = top_k_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("top-k sampled Ids:", ids)
Explanation: top_k sampling
In Top-K sampling, the K most likely next token ids are filtered and the probability mass is redistributed among only those K ids.
End of explanation
top_p_obj = sampling_module.SamplingModule(
    length_normalization_fn=length_norm,
    dtype=tf.float32,
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    vocab_size=3,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    sample_temperature=tf.constant(1.0),
    top_p=tf.constant(0.9),
    padded_decode=False,
    enable_greedy=False)
ids, _ = top_p_obj.generate(
    initial_ids=tf.constant([9, 1]), initial_cache=cache)
print("top-p sampled Ids:", ids)
Explanation: top_p sampling
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.
End of explanation
beam_size = 2
params['batch_size'] = 1
beam_cache = {
    'layer_%d' % layer: {
        'k': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32),
        'v': tf.zeros([params['batch_size'], params['max_decode_length'], params['num_heads'], params['n_dims']], dtype=tf.float32)
        } for layer in range(params['num_layers'])
    }
print("cache key shape for layer 1 :", beam_cache['layer_1']['k'].shape)
ids, _ = beam_search.sequence_beam_search(
    symbols_to_logits_fn=_symbols_to_logits_fn(),
    initial_ids=tf.constant([9], tf.int32),
    initial_cache=beam_cache,
    vocab_size=3,
    beam_size=beam_size,
    alpha=0.6,
    max_decode_length=params['max_decode_length'],
    eos_id=10,
    padded_decode=False,
    dtype=tf.float32)
print("Beam search ids:", ids)
Explanation: Beam search decoding
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.
End of explanation
<END_TASK>