<commit_before># Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""A decoder that performs beam search."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import numpy as np

import tensorflow as tf

from tensorflow_addons.seq2seq import attention_wrapper
from tensorflow_addons.seq2seq import decoder
from tensorflow_addons.utils import keras_utils
from tensorflow_addons.utils.resource_loader import get_path_to_datafile

_beam_search_ops_so = tf.load_op_library(
    get_path_to_datafile("custom_ops/seq2seq/_beam_search_ops.so"))
gather_tree = _beam_search_ops_so.addons_gather_tree


class BeamSearchDecoderState(
        collections.namedtuple("BeamSearchDecoderState",
                               ("cell_state", "log_probs", "finished",
                                "lengths", "accumulated_attention_probs"))):
    pass


class BeamSearchDecoderOutput(
        collections.namedtuple("BeamSearchDecoderOutput",
                               ("scores", "predicted_ids", "parent_ids"))):
    pass


class FinalBeamSearchDecoderOutput(
        collections.namedtuple(
            "FinalBeamDecoderOutput",
            ["predicted_ids", "beam_search_decoder_output"])):
    """Final outputs returned by the beam search after all decoding is
    finished.

    Args:
      predicted_ids: The final prediction. A tensor of shape
        `[batch_size, T, beam_width]` (or `[T, batch_size, beam_width]` if
        `output_time_major` is True). Beams are ordered from best to worst.
      beam_search_decoder_output: An instance of `BeamSearchDecoderOutput` that
        describes the state of the beam search.
    """
    pass


def _tile_batch(t, multiplier):
    """Core single-tensor implementation of tile_batch."""
    t = tf.convert_to_tensor(t, name="t")
    shape_t = tf.shape(t)
    if t.shape.ndims is None or t.shape.ndims < 1:
        raise ValueError("t must have statically known rank")
    tiling = [1] * (t.shape.ndims + 1)
    tiling[1] = multiplier
    tiled_static_batch_size = (t.shape.dims[0].value * multiplier
                               if t.shape.dims[0].value is not None else None)
    tiled = tf.tile(tf.expand_dims(t, 1), tiling)
    tiled = tf.reshape(tiled,
                       tf.concat(([shape_t[0] * multiplier], shape_t[1:]), 0))
    tiled.set_shape(
        tf.TensorShape([tiled_static_batch_size]).concatenate(t.shape[1:]))
    return tiled


def tile_batch(t, multiplier, name=None):
    """Tile the batch dimension of a (possibly nested structure of) tensor(s)
    t.

    For each tensor t in a (possibly nested structure) of tensors,
    this function takes a tensor t shaped `[batch_size, s0, s1, ...]` composed
    of minibatch entries `t[0], ..., t[batch_size - 1]` and tiles it to have a
    shape `[batch_size * multiplier, s0, s1, ...]` composed of minibatch
    entries `t[0], t[0], ..., t[1], t[1], ...` where each minibatch entry is
    repeated `multiplier` times.

    Args:
      t: `Tensor` shaped `[batch_size, ...]`.
      multiplier: Python int.
      name: Name scope for any created operations.

    Returns:
      A (possibly nested structure of) `Tensor` shaped
      `[batch_size * multiplier, ...]`.

    Raises:
      ValueError: if tensor(s) `t` do not have a statically known rank or
      the rank is < 1.
    """
    flat_t = tf.nest.flatten(t)
    with tf.name_scope(name or "tile_batch"):
        return tf.nest.map_structure(lambda t_: _tile_batch(t_, multiplier), t)


def gather_tree_from_array(t, parent_ids, sequence_length):
    """Calculates the full beams for `TensorArray`s.

    Args:
      t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
        shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
        where `s` is the depth shape.
      parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
      sequence_length: The sequence length of shape `[batch_size, beam_width]`.

    Returns:
      A `Tensor` which is a stacked `TensorArray` of the same size and type as
      `t` and where beams are sorted in each `Tensor` according to
      `parent_ids`.
    """
    max_time = parent_ids.shape.dims[0].value or tf.shape(parent_ids)[0]
    batch_size = parent_ids.shape.dims[1].value or tf.shape(parent_ids)[1]
    beam_width = parent_ids.shape.dims[2].value or tf.shape(parent_ids)[2]

    # Generate beam ids that will be reordered by gather_tree.
    beam_ids = tf.expand_dims(tf.expand_dims(tf.range(beam_width), 0), 0)
    beam_ids = tf.tile(beam_ids, [max_time, batch_size, 1])

    max_sequence_lengths = tf.cast(
        tf.reduce_max(sequence_length, axis=1), tf.int32)
    sorted_beam_ids = gather_tree(
        step_ids=beam_ids,
        parent_ids=parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=beam_width + 1)

    # For out of range steps, simply copy the same beam.
    in_bound_steps = tf.transpose(
        tf.sequence_mask(sequence_length, maxlen=max_time), perm=[2, 0, 1])
    sorted_beam_ids = tf.where(in_bound_steps, x=sorted_beam_ids, y=beam_ids)

    # Generate indices for gather_nd.
    time_ind = tf.tile(
        tf.reshape(tf.range(max_time), [-1, 1, 1]),
        [1, batch_size, beam_width])
    batch_ind = tf.tile(
        tf.reshape(tf.range(batch_size), [-1, 1, 1]),
        [1, max_time, beam_width])
    batch_ind = tf.transpose(batch_ind, perm=[1, 0, 2])
    indices = tf.stack([time_ind, batch_ind, sorted_beam_ids], -1)

    # Gather from a tensor with collapsed additional dimensions.
    gather_from = t
    final_shape = tf.shape(gather_from)
    gather_from = tf.reshape(gather_from,
                             [max_time, batch_size, beam_width, -1])
    ordered = tf.gather_nd(gather_from, indices)
    ordered = tf.reshape(ordered, final_shape)

    return ordered


def _check_ndims(t):
    if t.shape.ndims is None:
        raise ValueError(
            "Expected tensor (%s) to have known rank, but ndims == None." % t)


def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
    """Raises an exception if dimensions are known statically and can not be
    reshaped to [batch_size, beam_size, -1]."""
    reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
    assert len(shape.dims) > 0
    if batch_size is None or shape.dims[0].value is None:<commit_msg>Simplify gather_tree_from_array implementation (#576)<commit_after># Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""A decoder that performs beam search."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import numpy as np

import tensorflow as tf

from tensorflow_addons.seq2seq import attention_wrapper
from tensorflow_addons.seq2seq import decoder
from tensorflow_addons.utils import keras_utils
from tensorflow_addons.utils.resource_loader import get_path_to_datafile

_beam_search_ops_so = tf.load_op_library(
    get_path_to_datafile("custom_ops/seq2seq/_beam_search_ops.so"))
gather_tree = _beam_search_ops_so.addons_gather_tree


class BeamSearchDecoderState(
        collections.namedtuple("BeamSearchDecoderState",
                               ("cell_state", "log_probs", "finished",
                                "lengths", "accumulated_attention_probs"))):
    pass


class BeamSearchDecoderOutput(
        collections.namedtuple("BeamSearchDecoderOutput",
                               ("scores", "predicted_ids", "parent_ids"))):
    pass


class FinalBeamSearchDecoderOutput(
        collections.namedtuple(
            "FinalBeamDecoderOutput",
            ["predicted_ids", "beam_search_decoder_output"])):
    """Final outputs returned by the beam search after all decoding is
    finished.

    Args:
      predicted_ids: The final prediction. A tensor of shape
        `[batch_size, T, beam_width]` (or `[T, batch_size, beam_width]` if
        `output_time_major` is True). Beams are ordered from best to worst.
      beam_search_decoder_output: An instance of `BeamSearchDecoderOutput` that
        describes the state of the beam search.
    """
    pass


def _tile_batch(t, multiplier):
    """Core single-tensor implementation of tile_batch."""
    t = tf.convert_to_tensor(t, name="t")
    shape_t = tf.shape(t)
    if t.shape.ndims is None or t.shape.ndims < 1:
        raise ValueError("t must have statically known rank")
    tiling = [1] * (t.shape.ndims + 1)
    tiling[1] = multiplier
    tiled_static_batch_size = (t.shape.dims[0].value * multiplier
                               if t.shape.dims[0].value is not None else None)
    tiled = tf.tile(tf.expand_dims(t, 1), tiling)
    tiled = tf.reshape(tiled,
                       tf.concat(([shape_t[0] * multiplier], shape_t[1:]), 0))
    tiled.set_shape(
        tf.TensorShape([tiled_static_batch_size]).concatenate(t.shape[1:]))
    return tiled


def tile_batch(t, multiplier, name=None):
    """Tile the batch dimension of a (possibly nested structure of) tensor(s)
    t.

    For each tensor t in a (possibly nested structure) of tensors,
    this function takes a tensor t shaped `[batch_size, s0, s1, ...]` composed
    of minibatch entries `t[0], ..., t[batch_size - 1]` and tiles it to have a
    shape `[batch_size * multiplier, s0, s1, ...]` composed of minibatch
    entries `t[0], t[0], ..., t[1], t[1], ...` where each minibatch entry is
    repeated `multiplier` times.

    Args:
      t: `Tensor` shaped `[batch_size, ...]`.
      multiplier: Python int.
      name: Name scope for any created operations.

    Returns:
      A (possibly nested structure of) `Tensor` shaped
      `[batch_size * multiplier, ...]`.

    Raises:
      ValueError: if tensor(s) `t` do not have a statically known rank or
      the rank is < 1.
    """
    flat_t = tf.nest.flatten(t)
    with tf.name_scope(name or "tile_batch"):
        return tf.nest.map_structure(lambda t_: _tile_batch(t_, multiplier), t)


def gather_tree_from_array(t, parent_ids, sequence_length):
    """Calculates the full beams for `TensorArray`s.

    Args:
      t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
        shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
        where `s` is the depth shape.
      parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
      sequence_length: The sequence length of shape `[batch_size, beam_width]`.

    Returns:
      A `Tensor` which is a stacked `TensorArray` of the same size and type as
      `t` and where beams are sorted in each `Tensor` according to
      `parent_ids`.
    """
    max_time = parent_ids.shape.dims[0].value or tf.shape(parent_ids)[0]
    batch_size = parent_ids.shape.dims[1].value or tf.shape(parent_ids)[1]
    beam_width = parent_ids.shape.dims[2].value or tf.shape(parent_ids)[2]

    # Generate beam ids that will be reordered by gather_tree.
    beam_ids = tf.reshape(tf.range(beam_width), [1, 1, -1])
    beam_ids = tf.tile(beam_ids, [max_time, batch_size, 1])

    max_sequence_lengths = tf.cast(
        tf.reduce_max(sequence_length, axis=1), tf.int32)
    sorted_beam_ids = gather_tree(
        step_ids=beam_ids,
        parent_ids=parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=beam_width + 1)

    # For out of range steps, simply copy the same beam.
    in_bound_steps = tf.transpose(
        tf.sequence_mask(sequence_length, maxlen=max_time), perm=[2, 0, 1])
    sorted_beam_ids = tf.where(in_bound_steps, x=sorted_beam_ids, y=beam_ids)

    # Gather from a tensor with collapsed additional dimensions.
    final_shape = tf.shape(t)
    gather_from = tf.reshape(t, [max_time, batch_size, beam_width, -1])
    ordered = tf.gather(gather_from, sorted_beam_ids, axis=2, batch_dims=2)
    ordered = tf.reshape(ordered, final_shape)

    return ordered


def _check_ndims(t):
    if t.shape.ndims is None:
        raise ValueError(
            "Expected tensor (%s) to have known rank, but ndims == None." % t)


def _check_static_batch_beam_maybe(shape, batch_size, beam_width):
    """Raises an exception if dimensions are known statically and can not be
    reshaped to [batch_size, beam_size, -1]."""
    reshaped_shape = tf.TensorShape([batch_size, beam_width, None])
    assert len(shape.dims) > 0
    if batch_size is None or shape.dims[0].value is None:
{'fraction_non_alphanumeric': 0.11141959523635567, 'fraction_numerical': 0.00905969167823482, 'mean_word_length': 3.3275371482769525, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5561', 'n_tokens_mistral': 4366, 'n_tokens_neox': 4235, 'n_words': 1269}