def get_decode_steps(extended_indices, output_vocab_size, model_config):
  """Convert Tensor of indices in extended vocabulary to DecodeStep."""
  extended_indices = tf.to_int64(extended_indices)
  action_types = _get_action_type(extended_indices, output_vocab_size,
                                  model_config)
  action_ids = _get_action_id(extended_indices, action_types, output_vocab_size,
                              model_config)
  return DecodeSteps(action_types=action_types, action_ids=action_ids)
Convert Tensor of indices in extended vocabulary to DecodeStep.
language/xsp/model/decode_utils.py