def _DerivePaddingsAndIds(src_ids, tgt_labels):
"""tgt_ids is tgt_labels shifted right by one, with a SOS ID prepended."""
tgt_ids = tf.concat([[p.sos_id], tgt_labels[:-1]], axis=0)
src_paddings = tf.zeros(tf.shape(src_ids), dtype=tf.float32)
tgt_paddings = tf.zeros(tf.shape(tgt_ids), dtype=tf.float32)
tgt_weights = tf.ones(tf.shape(tgt_ids), dtype=tf.float32)
bucket_key = tf.cast(
tf.maximum(
tf.reduce_sum(1.0 - src_paddings),
tf.reduce_sum(1.0 - tgt_paddings)), tf.int32)
return src_paddings, tgt_ids, tgt_paddings, tgt_weights, bucket_key
def _MergeCandidates(tokens, candidates):
"""Merge in the reverse binary tree."""
best_id = tf.argmin(candidates, output_type=tf.int32)
# Perform the merge at position best_id.
tokens = tf.concat([
tokens[:best_id], [candidates[best_id]], tokens[best_id + 2:]
],
axis=0)
# Recompute the merge candidates.
# Only the neighbors of best_id need to be recomputed.
empty = tf.zeros([0], dtype=candidates.dtype)
def _MergeLeft():
return tf.concat([
candidates[:best_id - 1],
_MergeOneToken(tokens, best_id - 1)
],
axis=0)
left_candidates = tf.cond(tf.equal(best_id, 0), lambda: empty,
_MergeLeft)
def _MergeRight():
return tf.concat([
_MergeOneToken(tokens, best_id), candidates[best_id + 2:]
],
axis=0)
right_candidates = tf.cond(
tf.greater_equal(best_id,
tf.size(tokens) - 1), lambda: empty,
_MergeRight)
candidates = tf.concat([left_candidates, right_candidates], axis=0)
return tokens, candidates
def _CreateVariableStub(name,
params,
reuse=None,
trainable=True,
collections=None,
default_seed=None,
synchronization=None,
aggregation=None):
"""Return a zero tensor of the right shape instead of creating variable."""
del reuse
del default_seed
del synchronization
del aggregation
dtype = params.dtype
shape = py_utils.ToStaticShape(params.shape)
# For total samples counters we have to actually create variables so that
# we can access the 'value' attribute during construction.
if 'total_samples' in name:
var = tf.get_variable(name,
shape,
dtype,
tf.constant_initializer(0),
collections=collections,
trainable=trainable,
validate_shape=True)
else:
key = (tf.get_default_graph(), tuple(shape))
if key in variable_cache:
var = variable_cache[key]
else:
var = tf.zeros(shape, dtype)
variable_cache[key] = var
return var, var
def ReOrderHyps(x_in):
"""Reorders x_in based on prev hyp ids."""
if isinstance(x_in, tf.Tensor) and x_in.shape.ndims > 0:
# For rank > 1 tensors we make use of an efficient matmul based gather
# on tpu that takes in account the range of the values. For R1, we
# rely on the tf.gather and xla to optimize it efficiently for R1
# layout.
if x_in.shape.ndims > 1:
if p.batch_major_state:
num_hyps = tf.shape(old_hyp_ids)[0]
x_out = beam_search_tpu_ops.fast_gather(
x_in,
old_hyp_ids,
num_hyps,
max_value=None,
batch_major_state=p.batch_major_state)
else:
# Use corrected indices only here for batch major compute as
# key/value caches are the states being affected.
correct_old_hyp_ids = (old_hyp_ids_in_cache_order
if p.batch_major_compute else
old_hyp_ids)
def _GatherStep(x_in, t):
"""Gather for one time step.
Args:
x_in: in the shape of [T, B, ...] we first get slice(t) from the
tensors, then gather old_hyp_ids from the slice and write the
interpolated slice inplace to update the original x_in.
t: current time step
Returns:
Updated x_in and time step
"""
x = tf.gather(tf.gather(x_in, t),
correct_old_hyp_ids)
return inplace_ops.alias_inplace_update(
x_in, t, x), t + 1
x_out, _ = tf.while_loop(
lambda _, t: t <= cur_step, _GatherStep,
(x_in, tf.zeros([], tf.int32)))
else:
x_out = tf.gather(x_in, old_hyp_ids)
x_out.set_shape(x_in.get_shape())
return x_out
else:
return x_in
def GetState(self, theta):
"""Gets the state from theta."""
p = self.params
if p.is_inference:
# State is not used for inference. Just return dummy.
return tf.zeros([1], tf.float32)
else:
# Calculations/vars need to be float but these can be ints in the params.
clip_end_step = tf.cast(p.clip_end_step, tf.float32)
clip_start_step = tf.cast(p.clip_start_step, tf.float32)
quant_start_step = tf.cast(p.quant_start_step, tf.float32)
global_step = tf.cast(theta.global_step, tf.float32)
# Will be negative if before clipping starts.
clip_ratio = (tf.minimum(clip_end_step - clip_start_step,
global_step - clip_start_step) /
tf.maximum(1.0, clip_end_step - clip_start_step))
# Currently fq is either on (1.0) or off (-1.0). Progressive quantization
# may later occupy 0..1.0.
fq_ratio = tf.where(global_step < quant_start_step, -1.0, 1.0)
return tf.stack([clip_ratio, fq_ratio])
def ZeroState(self, theta, prepared_inputs, batch_size):
"""Produce a zero state for this step.
Args:
theta: A `.NestedMap` object containing weights' values of this layer and
its children layers.
prepared_inputs: A set of inputs pre-processed by using
PrepareExternalInputs.
batch_size: Number of elements in the batched input.
Returns:
state0, a state parameter to pass to FProp on its first invocation.
"""
max_seq_length = py_utils.GetShape(prepared_inputs.src, 3)[0]
atten_state = self.atten.ZeroAttentionState(max_seq_length, batch_size)
(new_atten_context, _,
new_atten_states) = self.atten.ComputeContextVectorWithSource(
theta.atten,
prepared_inputs.packed_src,
tf.zeros([batch_size, self.params.atten.query_dim],
dtype=py_utils.FPropDtype(self.params)),
attention_state=atten_state)
return py_utils.NestedMap(atten_context=new_atten_context,
atten_state=new_atten_states)
def GetState(self, theta):
return tf.zeros([1], tf.float32)
def FProp(self,
theta,
x,
x_paddings=None,
eos_id=1,
force_sample_last_token=True):
"""Applies SymbolInsertionLayer.
We take in a `x`, which represents the groundtruth sequence (i.e., English
sequence). We return a sampled rollin (observed) canvas (i.e., random subset
of the English sequence), as well as the target (indices) for an
insertion-based model (i.e., the targets given the random observed subset).
Args:
theta: Ignored, this can be None.
x: The symbol ids of shape `[batch_size, time_dim]`.
x_paddings: The paddings (1 or 0) of shape `[batch_size, time_dim]` where
0 is valid and 1 is invalid.
eos_id: The <eos> token id to represent end-of-slot.
force_sample_last_token: Set True to force sample the last token of `x`.
Returns:
A `NestedMap`.
- canvas: The canvas (based off of the `rollin_policy`) of shape
[batch_size, c_dim]. Note that, `c_dim` <= `time_dim` but need not be
equal.
- canvas_indices: The canvas indices (into `x`).
- canvas_paddings: The paddings of `canvas_indices`.
- target_indices: The target indices of shape [num_targets, 3].
`num_targets` is the number of total targets in the entire batch.
[:, 0] captures the batch, [:, 1] captures the slot, and [:, 2]
captures the token. Each row [batch, slot, vocab] represents the
indices of the target -- i.e., the batch, slot and vocab combination
of the target. Typical usage of these indices is to tf.gather_nd
the log-probs (from the softmax layer).
- target_weights: The target weights.
Raises:
ValueError: If invalid params.
"""
p = self.params
batch_size = py_utils.GetShape(x)[0]
time_dim = py_utils.GetShape(x)[1]
if x_paddings is None:
x_paddings = tf.zeros([batch_size, time_dim], tf.float32)
oracle_policy = p.oracle_policy
rollin_policy = (oracle_policy
if p.rollin_policy == 'oracle' else p.rollin_policy)
if rollin_policy != 'uniform':
raise ValueError('Unknown or unsupported rollin policy: %s' %
rollin_policy)
if oracle_policy != 'uniform':
raise ValueError('Unknown or unsupported oracle policy: %s' %
oracle_policy)
x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
# Compute the desired length per example in the batch.
ratio = tf.random.uniform([batch_size], 0.0, 1.0, seed=p.random_seed)
if force_sample_last_token:
c_len = tf.minimum(
tf.cast(ratio * tf.cast(x_len, tf.float32), tf.int32),
x_len - 1) + 1
else:
c_len = tf.minimum(
tf.cast(ratio * tf.cast(x_len + 1, tf.float32), tf.int32),
x_len)
# Compute the maximum length across the batch.
c_len_max = tf.reduce_max(c_len)
# Grab subset of random valid indices per example.
z_logits = tf.cast(
tf.expand_dims(tf.range(time_dim), 0) >= tf.expand_dims(x_len, 1),
tf.float32) * -1e9
if force_sample_last_token:
# Force sample the last token -- i.e., as indexed by `x_len - 1`. We can
# accomplish this by add +LARGE_NUMBER to the logits.
z_logits += tf.cast(
tf.equal(tf.expand_dims(tf.range(time_dim), 0),
tf.expand_dims(x_len - 1, 1)), tf.float32) * 1e9
# Gumbel-max trick to sample (we only sample valid positions per sample in
# the batch).
z = -tf.math.log(-tf.math.log(
tf.random.uniform([batch_size, time_dim], seed=p.random_seed)))
unused_c_values, c_indices = tf.nn.top_k(z_logits + z, time_dim)
# Trim everything > c_len_max.
c_indices = c_indices[:, :c_len_max]
# Invalidate any indices >= c_len, we use the last index as the default
# invalid index.
c_indices = tf.where(
tf.expand_dims(tf.range(c_len_max), 0) < tf.expand_dims(c_len, 1),
c_indices, tf.fill(py_utils.GetShape(c_indices), time_dim - 1))
# Materialize the canvas.
c_indices = tf.sort(c_indices)
c = tf.gather_nd(
x,
tf.stack([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, c_len_max]), [-1]),
tf.reshape(c_indices, [-1])
], 1))
c = tf.reshape(c, [batch_size, c_len_max])
# Compute the paddings.
c_paddings = 1 - tf.sequence_mask(
c_len, c_len_max, dtype=x_paddings.dtype)
c *= tf.cast(1 - c_paddings, tf.int32)
indices = tf.concat([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, c_len_max]), [batch_size * c_len_max, 1]),
tf.reshape(c_indices, [batch_size * c_len_max, 1])
], 1)
x_token_is_observed = tf.scatter_nd(
indices, tf.ones([batch_size * c_len_max], tf.int32),
py_utils.GetShape(x))
# `x_segments` captures which slot each `x` belongs to (both observed and
# tokens that need to be observed).
x_segments = tf.cumsum(x_token_is_observed, 1, exclusive=True)
x_token_is_observed = tf.cast(x_token_is_observed, tf.bool)
prev_x_token_is_observed = tf.pad(x_token_is_observed[:, :-1],
[[0, 0], [1, 0]],
constant_values=True)
x_token_is_observed = tf.reshape(x_token_is_observed, [-1])
prev_x_token_is_observed = tf.reshape(prev_x_token_is_observed, [-1])
x_is_valid = tf.cast(1 - x_paddings, tf.bool)
x_is_valid = tf.reshape(x_is_valid, [-1])
# Remap all the observed to <eos>, note some of these need a zero weight
# (or else there would be <eos> and valid token in the same slot).
target_indices = tf.cast(tf.reshape(x, [-1, 1]), tf.int32)
target_indices = tf.where(
x_token_is_observed,
tf.fill(py_utils.GetShape(target_indices), eos_id), target_indices)
# TODO(williamchan): We give uniform 1.0 weight, however, math suggests
# we may want to weigh this term by the original sequence length.
target_weights = tf.ones_like(target_indices, tf.float32)
# We need to set all the weights for <eos> which actually have valid tokens
# in the slot to zero.
target_weights = tf.where(
x_token_is_observed & ~prev_x_token_is_observed,
tf.zeros_like(target_weights), target_weights)
# TODO(williamchan): Consider dropping the entries w/ weight zero.
# Add the batch and slot indices.
target_indices = tf.concat([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, time_dim]), [batch_size * time_dim, 1]),
tf.reshape(x_segments, [-1, 1]), target_indices
], 1)
# Select only the valid indices. The selected valid ones include slots w/
# <eos>.
target_indices = target_indices[x_is_valid]
target_weights = target_weights[x_is_valid]
return py_utils.NestedMap(canvas=c,
canvas_indices=c_indices,
canvas_paddings=c_paddings,
target_indices=target_indices,
target_weights=target_weights)
def _BeamSearchDecodeIds(self,
theta,
encoder_outputs,
num_hyps_per_beam,
init_beam_search_state=None,
pre_beam_search_step_callback=None,
post_beam_search_step_callback=None,
max_steps=None):
"""Performs beam-search based decoding.
Args:
theta: A NestedMap object containing weights' values of the decoder layer
and its children layers.
encoder_outputs: A NestedMap computed by encoder.
num_hyps_per_beam: Number of hyps per beam.
init_beam_search_state: The InitBeamSearchState callback. Please refer to
the class header comments for more details.
pre_beam_search_step_callback: The PreBeamSearchStepCallback callback.
Please refer to the class header comments for more details.
post_beam_search_step_callback: The PostBeamSearchStepCallback callback.
Please refer to the class header comments for more details.
max_steps: maximum beam search steps. If None, use
self.params.target_seq_len.
Returns:
hyps: A tensor of shape [time, b * k] with ids of the token selected.
prev_hyps: A tensor of shape [time, b * k] with index to the previous hyps
which was selected.
done_hyps: A boolean tensor of shape [time, b * k] where value
indicates if hyps was terminated.
scores: A tensor of shape [time, b * k] with scores of the token
selected.
atten_probs: A tensor of shape [time, b * k, seq_len] which contain the
attention probabilities over the source words against word in the
previous hyps.
eos_scores: A tensor of shape [time, b * k] with scores of the eos token
selected.
eos_atten_probs: A tensor of shape [time, b * k, seq_len] which contain
the attention probabilities over the source words against word in the
previous hyps.
source_seq_lengths: A tensor of shape [time] containing the source
seq_lengths.
flat_final_other_states: A array of tensors that are part of other states.
"""
p = self.params
source_paddings = encoder_outputs.padding
initial_results, other_states = init_beam_search_state(
theta, encoder_outputs, num_hyps_per_beam)
num_hyps = tf.shape(initial_results.log_probs)[0]
num_beams = num_hyps // num_hyps_per_beam
# We cache the NestedMap as member variable so that we can use it to
# pack the final outputs. Tpu rewrite methods forces us to strictly pass
# in Tensors, and output Tensors
self._other_states = other_states
step_ids = tf.fill([num_hyps, 1],
tf.constant(p.target_sos_id, dtype=tf.int32))
min_score = -1e36
fprop_dtype = py_utils.FPropDtype(p)
best_scores = (tf.zeros(shape=[num_beams], dtype=fprop_dtype) +
min_score)
cumulative_scores = tf.zeros(shape=[num_hyps], dtype=fprop_dtype)
histories = tf.zeros(shape=[num_hyps], dtype=tf.int32)
in_scores = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_prev_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_done_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_atten_probs = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_eos_scores = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_eos_atten_probs = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
cur_step = tf.constant(0, dtype=tf.int32)
all_done = tf.constant(False, dtype=tf.bool)
# States for beam search that are inputs into Beam search step.
accum_bs_states = [best_scores, cumulative_scores, histories]
# States that are not accumulators.
non_accum_bs_states = [
in_scores,
in_hyps,
in_prev_hyps,
in_done_hyps,
in_atten_probs,
in_eos_scores,
in_eos_atten_probs,
]
core_bs_states = tuple(accum_bs_states + non_accum_bs_states)
flat_other_states = other_states.Flatten()
# If there is an optimized implementation for short sequence, LoopBodyShort
# will run first for short_seq_limit steps (after which the
# LoopBodyShort does not have performance benefit). Then LoopBodyLong (the
# default implementation) is used to continue the rest of the steps. For
# decoders which do not have the short sequence specific implementation,
# only the LoopBodyLong (the default implementation) will run.
if p.short_seq_limit > 0:
def LoopContinueShort(cur_step, all_done, unused_step_ids,
unused_core_bs_states,
unused_other_states_list):
"""Use short_seq optimization when cur_step is smaller than limit."""
return tf.math.logical_and(cur_step < p.short_seq_limit,
tf.math.logical_not(all_done))
def LoopBodyShort(cur_step, unused_all_done, step_ids,
core_bs_states, other_states_list):
"""Loop body of short_seq optimization.
Instead of doing computation for the entire padded sequence, while loop
with early exit is used within each _BeamSearchStep to do computation
for only the actual sequence (seq_length <= cur_step).
use_short_seq_opt is used as the flag to pass this information down to
the decoder implementation.
Args:
cur_step: A scalar int tensor, the current time step, 0-based.
unused_all_done: A tf.bool, indicating whether the decoding finishes.
step_ids: An int32 tensor of shape [num_hyps, 1]. The input ids to the
current search step.
core_bs_states: A tuple of core beam search states.
other_states_list: A flattened NestedMap of other beam search states.
Returns:
The updated input tuple, with the same shape.
"""
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta,
encoder_outputs,
cur_step,
step_ids,
core_bs_states,
other_states.Pack(other_states_list),
num_hyps_per_beam,
pre_beam_search_step_callback,
post_beam_search_step_callback,
use_short_seq_opt=True)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
(cur_step, all_done, step_ids, core_bs_states,
flat_other_states) = tf.while_loop(
LoopContinueShort,
LoopBodyShort,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(
tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
tuple(
list(_GetShapes(accum_bs_states)) +
list(_GetShapes(non_accum_bs_states,
none_shapes=True))),
_GetShapes(flat_other_states, none_shapes=True)),
maximum_iterations=max_steps)
def LoopContinueLong(cur_step, all_done, unused_step_ids,
unused_core_bs_states, unused_other_states_list):
"""Continue default implementation until decoding finishes."""
return tf.math.logical_and(cur_step < max_steps,
tf.math.logical_not(all_done))
def LoopBodyLong(cur_step, unused_all_done, step_ids, core_bs_states,
other_states_list):
"""Loop body of default long_seq implementation."""
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta,
encoder_outputs,
cur_step,
step_ids,
core_bs_states,
other_states.Pack(other_states_list),
num_hyps_per_beam,
pre_beam_search_step_callback,
post_beam_search_step_callback,
use_short_seq_opt=False)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
_, _, _, final_bs_states, flat_final_other_states = tf.while_loop(
LoopContinueLong,
LoopBodyLong,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(
tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
tuple(
list(_GetShapes(accum_bs_states)) +
list(_GetShapes(non_accum_bs_states, none_shapes=True))),
_GetShapes(flat_other_states, none_shapes=False)),
maximum_iterations=max_steps)
if isinstance(source_paddings, py_utils.NestedMap):
source_seq_lengths = tf.cast(tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings.Flatten()[0]),
1)),
dtype=tf.int32)
else:
source_seq_lengths = tf.cast(tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings), 1)),
dtype=tf.int32)
# Concatenate all outputs on axis=0.
scores = final_bs_states[3].stack()
hyps = final_bs_states[4].stack()
prev_hyps = final_bs_states[5].stack()
done_hyps = tf.cast(final_bs_states[6].stack(), tf.bool)
atten_probs = final_bs_states[7].stack()
eos_scores = final_bs_states[8].stack()
eos_atten_probs = final_bs_states[9].stack()
rets = (hyps, prev_hyps, done_hyps, scores, atten_probs, eos_scores,
eos_atten_probs, source_seq_lengths)
# TODO(rohananil): Only send a single R1 tensor to host instead of 3 after
# b/111131551 is resolved.
# Canonical shapes for tensors of various. ranks
r_shapes = [
py_utils.GetShape(source_seq_lengths),
py_utils.GetShape(hyps),
py_utils.GetShape(atten_probs)
]
# Reshape all tensors to [-1] to avoid cost of copy due to padding.
rets_r1 = [tf.reshape(r, [-1]) for r in rets]
return tuple(r_shapes) + tuple(rets_r1) + tuple(
flat_final_other_states)
def DefaultValue(self): """Returns the default value of the accumulator.""" return tf.zeros(self.shape, dtype=self.dtype)
def BeamSearchDecode(self,
theta,
encoder_outputs,
num_hyps_per_beam_override=0,
init_beam_search_state=None,
pre_beam_search_step_callback=None,
post_beam_search_step_callback=None,
max_steps=None):
"""Performs beam-search based decoding.
Args:
theta: A NestedMap object containing weights' values of the decoder layer
and its children layers.
encoder_outputs: A NestedMap containing encoder outputs to be passed to
the callbacks. Mostly opaque to BeamSearchHelper, except that it should
contain either a 'seq_lengths' field of shape [source_batch_size] or
a 'paddings' field of shape [source_max_lengths, source_batch_size].
num_hyps_per_beam_override: If set to a value <= 0, this parameter is
ignored. If set to a value > 0, then this value will be used to override
`p.num_hyps_per_beam`.
init_beam_search_state: The `InitBeamSearchState` callback. Please refer
to the class header comments for more details.
pre_beam_search_step_callback: The `PreBeamSearchStepCallback` callback.
Please refer to the class header comments for more details.
post_beam_search_step_callback: The `PostBeamSearchStepCallback` callback.
Please refer to the class header comments for more details.
max_steps: maximum beam search steps. If None, use
self.params.target_seq_len.
Returns:
A `BeamSearchDecodeOutput`.
"""
p = self.params
num_hyps_per_beam = p.num_hyps_per_beam
if num_hyps_per_beam_override > 0:
num_hyps_per_beam = num_hyps_per_beam_override
if max_steps is None:
max_steps = p.target_seq_len
initial_results, other_states = init_beam_search_state(
theta, encoder_outputs, num_hyps_per_beam)
num_hyps = tf.shape(initial_results.log_probs)[0]
num_beams = num_hyps // num_hyps_per_beam
if 'step_ids' in initial_results:
# [num_hyps, 1]
step_ids = tf.ensure_shape(initial_results.step_ids, [None, 1])
else:
step_ids = tf.fill([num_hyps, 1],
tf.constant(p.target_sos_id, dtype=tf.int32))
min_score = -1e36
best_scores = (tf.zeros(shape=[num_beams], dtype=p.dtype) + min_score)
cumulative_scores = tf.zeros(shape=[num_hyps], dtype=p.dtype)
in_scores = tf.zeros([max_steps, num_hyps], dtype=p.dtype)
in_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32)
in_prev_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32)
in_done_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.string)
bs_atten_probs = tf.zeros(
[max_steps, num_hyps,
tf.shape(initial_results.atten_probs)[1]],
dtype=p.dtype)
cur_step = tf.constant(0, dtype=tf.int32)
all_done = tf.constant(False, dtype=tf.bool)
core_bs_states = (best_scores, cumulative_scores, in_scores, in_hyps,
in_prev_hyps, in_done_hyps, bs_atten_probs)
def LoopContinue(cur_step, all_done, unused_step_ids,
unused_core_bs_states, unused_other_states_list):
return tf.math.logical_and(cur_step < max_steps,
tf.math.logical_not(all_done))
def LoopBody(cur_step, unused_all_done, step_ids, core_bs_states,
other_states_list):
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta, encoder_outputs, cur_step, step_ids, core_bs_states,
other_states.Pack(other_states_list), num_hyps_per_beam,
pre_beam_search_step_callback, post_beam_search_step_callback)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
flat_other_states = other_states.Flatten()
_, _, _, final_bs_states, flat_final_other_states = tf.while_loop(
LoopContinue,
LoopBody,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
_GetShapes(core_bs_states),
_GetShapes(flat_other_states, none_shapes=True)))
# [target_seq_len, num_beams * num_hyps_per_beam].
final_done_hyps = final_bs_states[5]
final_other_states = other_states.Pack(flat_final_other_states)
# Assume that `paddings` has shape [source_max_lengths, source_batch_size]
# by default, and compute `encoded_seq_lengths` accordingly. This can be
# overridden by directly passing `seq_lengths` in the `encoder_outputs`
# NestedMap.
encoded_seq_lengths = getattr(encoder_outputs, 'seq_lengths', None)
if encoded_seq_lengths is None:
source_paddings = encoder_outputs.padding
if isinstance(source_paddings, py_utils.NestedMap):
encoded_seq_lengths = tf.cast(
tf.round(
tf.reduce_sum(
1.0 - tf.transpose(source_paddings.Flatten()[0]),
1)), tf.int32)
else:
encoded_seq_lengths = tf.cast(
tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings), 1)),
tf.int32)
# [num_beams, num_hyps_per_beam].
topk_hyps = ops.top_k_terminated_hyps(
final_done_hyps,
encoded_seq_lengths,
k=num_hyps_per_beam,
num_hyps_per_beam=num_hyps_per_beam,
length_normalization=p.length_normalization,
coverage_penalty=p.coverage_penalty,
target_seq_length_ratio=p.target_seq_length_ratio,
eoc_id=p.target_eoc_id,
merge_paths=p.merge_paths)
# [num_beams * num_hyps_per_beam, ...].
max_seq_length = 0 if isinstance(max_steps, tf.Tensor) else max_steps
topk_ids, topk_lens, topk_scores = ops.unpack_hyp(
tf.reshape(topk_hyps, [-1]), max_seq_length=max_seq_length)
# [num_beams, num_hyps_per_beam].
topk_scores = tf.reshape(topk_scores, tf.shape(topk_hyps))
return BeamSearchDecodeOutput(final_done_hyps, topk_hyps, topk_ids,
topk_lens, topk_scores, None,
final_other_states)
def merge_hyps(global_score_values, histories_in, mask, num_beams,
num_hyps_per_beam):
"""Merges candidate hypotheses with identical histories.
This function takes a set of candidate hypotheses, represented as Tensors of
scores and histories, and merges all pairs of hypotheses that have identical
history hashes. When two hypotheses are merged, the hyp with lower global
score gets "deleted" and has its probability mass added to the higher scoring
one. Hypotheses are "deleted" by giving them empty history and a large
negative global score. The function output is a tuple of new
(global_score_values, histories) Tensors.
All input Tensors are assumed to be in "div" hypothesis ordering. That is,
element [i, ...] corresponds to the j-th hyp of the n-th beam, where j = i % k
and n = i / k.
Example:
Suppose num_beams = 1, num_hyps_per_beam = 2, candidates_per_hyp = 5,
global_score_values is
[[11 12 13 14 15],
[17 16 10 19 20]]
and histories_in is
[[1 2 3 4 5],
[5 6 3 7 8]].
There are two pairs of hypotheses with identical histories that should
be merged -- two with hash value 3 and two with hash 5. In each pair, the
one with lower score will be deleted and merged into the one with higher
score.
The output is a new set of global_score_values,
[[ 11 12 13.04 14 -1e34 ],
17.13 16 -1e34 19 20 ]]
and new histories
[[1 2 3 4 0],
[5 6 0 7 8]].
Hypotheses deleted in the merge now have zero history and a large negative
score. The destination of each merge now has additional probability mass.
(Note _log_sum_exp(13, 10) ~= 13.04 and _log_sum_exp(15, 17) ~= 17.13.)
Args:
global_score_values: Tensor of shape [b * k, candidates_per_hyp], the global
scores of each candidate hypothesis.
histories_in: int32 Tensor of shape [b * k, candidates_per_hyp], the
histories of each candidate hypothesis.
mask: Tensor of shape [b * k, 1] indicating which entries in
global_score_values and histories_in are valid.
num_beams: int, the number of beams (b above).
num_hyps_per_beam: int, the number of hypotheses per beam (k above).
Returns:
A tuple of new (global_score_values, histories) updated so that input
hypotheses with identical histories are now merged. Hypotheses deleted in
the merge have a new global score of BEST_SCORES_INIT and a history of 0.
"""
values_dtype = global_score_values.dtype
candidates_per_hyp = histories_in.get_shape()[1]
k = num_hyps_per_beam
# High-level strategy: To detect hyps to merge, we'll permute the hypotheses
# within each beam so that their histories are in sorted order. We can then
# in parallel check whether each history is equal to its left or right
# neighbor (i.e. whether the hyps should be merged), and if so, which of them
# has the higher global score (the direction of the merge). When two hyps need
# to be merged, we'll "delete" the one with lower score (by giving it a large
# negative score and empty history) and add its probability mass to the other.
#
# Note we only have to do pair-wise merging once per beam search step, because
# (ignoring hash collisions) there are at most two candidate hypotheses with
# any particular history. This follows from the fact that hypotheses are
# unique at the start of the beam search step, as are the top K non-epsilon
# extensions of those hypotheses. Thus, if there are two paths with
# identical histories, they must have the form
# h_i <eps> == h_j s (for some i != j, s != eps),
# where h_i and h_j are distinct input hypotheses, and s is some non-epsilon
# symbol.
# Reshape inputs to [b, num_hyps_per_beam * candidates_per_hyp] so they're
# grouped by beam.
histories = histories_in
orig_scores_shape = tf.shape(global_score_values)
histories = tf.reshape(histories, [num_beams, -1])
histories_valid = tf.cast(
tf.reshape(tf.tile(mask, [1, candidates_per_hyp]), [num_beams, -1]),
values_dtype)
# Compute the permutation of hyps within each beam that put the histories in
# sorted order, and the one that permutates the sorted hyps back to the
# original order.
sorted_history_indices = tf.argsort(histories, axis=1)
inverse_indices = tf.argsort(sorted_history_indices, axis=1)
def to_flat_indices(column_indices_per_row):
column_indices_per_row.shape.assert_has_rank(2)
flat_indices = (column_indices_per_row +
num_hyps_per_beam * candidates_per_hyp *
tf.reshape(tf.range(num_beams), [num_beams, 1]))
return tf.reshape(flat_indices, [-1])
# Convert to linear indices so we can use fast_gather.
sorted_history_indices_flat = to_flat_indices(sorted_history_indices)
inverse_indices_flat = to_flat_indices(inverse_indices)
def history_sort(values):
return tf.reshape(
fast_gather(tf.reshape(values,
[-1, 1]), sorted_history_indices_flat,
num_beams * k * candidates_per_hyp),
[num_beams, k * candidates_per_hyp])
def history_unsort(values):
return tf.reshape(
fast_gather(tf.reshape(values, [-1, 1]), inverse_indices_flat,
num_beams * k * candidates_per_hyp), orig_scores_shape)
sorted_histories = history_sort(histories)
sorted_histories_valid = history_sort(histories_valid)
# Indicators of whether each hypothesis is a duplicate of its left/right
# neighbors.
# [num_batches, k * candidates_per_hyp - 1]
dup_mask = tf.cast(
tf.equal(sorted_histories[:, 1:], sorted_histories[:, :-1]),
values_dtype) * (sorted_histories_valid[:, 1:] *
sorted_histories_valid[:, :-1])
padding = tf.zeros([num_beams, 1], dtype=values_dtype)
is_dup_of_left = tf.concat([padding, dup_mask], axis=1)
is_dup_of_right = tf.concat([dup_mask, padding], axis=1)
# Examine global scores to see which hyps should be merged, and within those
# cases, which hyps get deleted/retained in the merge.
sorted_global_scores = history_sort(global_score_values)
# Global scores of each hyp's left and right neighbors.
right_global_scores = tf.concat([sorted_global_scores[:, 1:], padding],
axis=1)
left_global_scores = tf.concat([padding, sorted_global_scores[:, :-1]],
axis=1)
# Masks indicating whether each candidate hyp is better or worse than its
# left or right neighbor.
is_better_than_right = tf.cast(
tf.greater_equal(sorted_global_scores, right_global_scores),
values_dtype)
is_worse_than_right = 1.0 - is_better_than_right
is_better_than_left = tf.cast(
tf.greater(sorted_global_scores, left_global_scores), values_dtype)
is_worse_than_left = 1.0 - is_better_than_left
# Determine which hypotheses need to be merged.
is_merge_source = tf.minimum(
is_dup_of_left * is_worse_than_left +
is_dup_of_right * is_worse_than_right, 1.0)
is_left_merge_dest = is_dup_of_left * is_better_than_left
is_right_merge_dest = is_dup_of_right * is_better_than_right
is_merge_dest = tf.minimum(is_left_merge_dest + is_right_merge_dest, 1.0)
# Mask of hyps unaffected by merging.
is_unchanged = tf.maximum(1.0 - is_merge_source - is_merge_dest, 0.0)
sorted_global_scores = (
is_unchanged * sorted_global_scores +
is_merge_source * BEST_SCORES_INIT + is_left_merge_dest *
_log_sum_exp(left_global_scores, sorted_global_scores) +
is_right_merge_dest *
_log_sum_exp(right_global_scores, sorted_global_scores))
# Set histories of deleted (merge source) hyps to zero.
sorted_histories *= tf.cast(1.0 - is_merge_source, sorted_histories.dtype)
# Put everything back in its original order and rank.
global_score_values_out = history_unsort(sorted_global_scores)
histories_out = history_unsort(sorted_histories)
return global_score_values_out, histories_out
def Sample(self, decoder_theta, encoder_outputs, random_seed,
init_state_callback, pre_step_callback, post_step_callback):
"""Samples target sequences, one target sequence per source sequence.
(Please see beam_search_helper.py for description of decoder callbacks.)
Args:
decoder_theta: A NestedMap object containing weights' values of the
decoder layer and its children layers, to be passed to decoder
callbacks.
encoder_outputs: the outputs of the encoder, to be passed to callbacks.
random_seed: a scalar int32 tensor representing the random seed.
init_state_callback: decoder._InitBeamSearchStateCallback.
pre_step_callback: decoder._PreBeamSearchStepCallback.
post_step_callback: decoder._PostBeamSearchStepCallback.
Returns:
A NestedMap containing the following tensors
- 'logits': [batch, max_target_length, vocab_size], representing the
distribution from which target sequences are sampled.
- 'ids': [batch, max_target_length] of int32, representing the target
sequence ids, not including target_sos_id, but maybe ending with
target_eos_id if end-of-sequence is reached before target_seq_len.
- 'paddings': [batch, max_target_length] of 0/1, where 1 represents
a padded timestep.
"""
p = self.params
assert p.temperature > 0
if getattr(encoder_outputs, 'segment_id', 1) is None:
# Remove None values, which are not supported by recurrent.
del encoder_outputs['segment_id']
# init_state_callback may modify 'encoder_outputs', e.g., by inserting
# 'packed_src'.
bs_result, bs_state = init_state_callback(decoder_theta,
encoder_outputs,
num_hyps_per_beam=1)
# 'recurrent_theta' represents all cross-timestep information used by the
# recurrent loop below, including layer theta and encoder outputs.
recurrent_theta = py_utils.NestedMap(theta=decoder_theta,
random_seed=random_seed,
encoder_outputs=encoder_outputs)
batch = tf.shape(bs_result.log_probs)[0]
recurrent_state0 = py_utils.NestedMap(
timestep=tf.zeros(shape=[], dtype=tf.int32),
logits=bs_result.log_probs,
# Start with target_sos_id.
ids=tf.fill([batch], tf.cast(p.target_sos_id, tf.int32)),
bs_state=bs_state)
inputs = py_utils.NestedMap(dummy=tf.zeros([p.target_seq_len, batch]))
def Step(recurrent_theta, state0, inputs):
"""Computes one decoder step."""
del inputs
with tf.name_scope('single_sampler_step'):
# Compute logits and states.
bs_result, bs_state1 = pre_step_callback(
recurrent_theta.theta,
recurrent_theta.encoder_outputs,
tf.expand_dims(state0.ids, 1), # [batch, 1].
state0.bs_state,
num_hyps_per_beam=1)
batch = tf.shape(bs_result.log_probs)[0]
state1 = py_utils.NestedMap(timestep=state0.timestep + 1)
state1.logits = bs_result.log_probs
# Sample ids from logits. [batch].
state1.ids = tf.reshape(
tf.random.stateless_categorical(
state1.logits / p.temperature,
num_samples=1,
seed=tf.stack(
[recurrent_theta.random_seed, state0.timestep]),
dtype=state0.ids.dtype,
name='sample_next_id'), [batch])
if 'is_last_chunk' in bs_result and p.target_eoc_id >= 0:
state1.ids = tf.where(
tf.math.logical_and(
bs_result.is_last_chunk,
tf.equal(state1.ids, p.target_eoc_id)),
tf.fill(tf.shape(state1.ids), p.target_eos_id),
state1.ids)
state1.bs_state = post_step_callback(
recurrent_theta.theta, recurrent_theta.encoder_outputs,
state1.ids, bs_state1)
return state1, py_utils.NestedMap()
accumulated_states, _ = recurrent.Recurrent(
recurrent_theta,
recurrent_state0,
inputs,
Step,
allow_implicit_capture=True)
result = py_utils.NestedMap(logits=tf.transpose(
accumulated_states.logits, [1, 0, 2]),
ids=tf.transpose(accumulated_states.ids))
result.paddings = tf.cast(
_ComputePaddings(result.ids, p.target_eos_id), result.logits.dtype)
# Force ids to be eos_id if the timestep is padded.
result.ids = tf.where(tf.equal(result.paddings, 0), result.ids,
tf.fill(tf.shape(result.ids), p.target_eos_id))
static_batch_size = bs_result.log_probs.shape[0]
result.ids.set_shape([static_batch_size, p.target_seq_len])
result.paddings.set_shape([static_batch_size, p.target_seq_len])
return result
def FProp(self, theta, *args):
"""Runs p.repeat copies of self.body.FProp independently.
Args:
theta: Layer model parameters. The shape of each variable in theta is
always [p.repeat, ...]. And the i-th slice theta[i] becomes theta of the
i-th copy of self.body.
*args: Input arguments. The shape of each tensor in args is always
[p.repeat, ....]. And the list [arg[i] for arg in args] becomes inputs
to the i-th copy of self.body.FProp.
Returns:
The accumulated output_tensors. Each tensor t in the return has the shape
[p.repeat, ....] and the tuple (t[i] for i in output_tensors) is the
return tuple of the i-th self.body.FProp.
"""
p = self.params
for arg in args:
if arg is not None:
arg = py_utils.HasShape(arg, [p.repeat], ndims=1)
theta_stack = _MaybeStackExtraTheta(theta.body, self.body.vars,
p.repeat)
inputs = py_utils.NestedMap(theta=theta_stack, args=list(args))
# Infer out_shapes from FPropMeta.
out_shapes = self._InferOutShapes(args)
def _CellFn(unused_theta, unused_state0, inputs):
"""Recurrent cell function wrapper of body.FProp."""
# Sets shapes for both theta and inputs to self.body.FProp.
for dst, src in zip(inputs.args + inputs.theta.Flatten(),
list(args) + theta_stack.Flatten()):
if src is not None:
dst.set_shape(tf.TensorShape(src.shape.as_list()[1:]))
# Runs the actual body.FProp
fprop_outputs = self.body.FProp(inputs.theta, *inputs.args)
fprop_outputs = _ToTuple(fprop_outputs)
assert len(fprop_outputs) == len(out_shapes)
# Passes fprop outputs to the next layer through state.
state1 = py_utils.NestedMap(outputs=list(fprop_outputs))
return state1, py_utils.NestedMap()
with tf.name_scope(p.name):
# Initiate state0 with inferred output shapes.
state0 = py_utils.NestedMap(outputs=[
tf.zeros(shape, args[0].dtype) for shape in out_shapes
])
# Runs body.FProp p.repeat times using Recurrent.
acc_states, _ = recurrent.Recurrent(theta=py_utils.NestedMap(),
state0=state0,
inputs=inputs,
cell_fn=_CellFn)
# Retrieves fprop outputs from state1 and sets shapes.
output_tensors = tuple(acc_states.outputs)
for out_idx in range(len(output_tensors)):
output_tensors[out_idx].set_shape(
tf.TensorShape([p.repeat] + out_shapes[out_idx].as_list()))
return output_tensors[0] if len(args) == 1 else tuple(
output_tensors)
def _TfZeros(t_shape):
if t_shape is None:
return None
return tf.zeros(t_shape.ToTensorShape().as_list(), dtype=state_dtype)
def DefaultValue(self):
return tf.zeros([3], dtype=self.dtype, name='qstate_zero')
def beam_search_step(in_scores,
in_atten_probs,
in_best_scores,
in_cumulative_scores,
in_histories,
cur_step,
eos_id,
num_beams,
beam_size,
num_hyps_per_beam,
valid_eos_max_logit_delta=5.0,
local_eos_threshold=-100.0,
merge_paths=False,
is_last_chunk=None,
eoc_id=0):
"""A single step of beam search.
Let "b" be the number of beams, "k" be the number hyps in each beam. This
function supports values with dtypes tf.float32 or tf.bfloat16.
The following data structures are allocated before the first decoding step and
are passed along from cur step to the next step:
Args:
in_scores: A tensor of shape [b * k, vocab_size], where [i, ...] is the
token score of the j-th hyps of the n-th beam. j = (i / k), and n = i % k
in_atten_probs: A tensor of shape [b*k, s_len], where in_atten_probs[i, ...]
is the attention probabilities over the source words of the j-th hyps of
n-th beam (where j, and n are derived as above).
in_best_scores: A vector of size [b], best scores of terminated hyps so far
in each of the beams.
in_cumulative_scores: A vector of size [b * k]. The cumulative score of each
active hyp before the current step.
in_histories: An int32 vector of size [b * k] containing hashes of the
histories of each active hyp. If 'merge_paths' is enabled, the histories
are used to identify hypotheses that are identical modulo epsilons (e.g.
"a <eps> b" and "a b <eps>") and merge them. See 'update_histories'
docstring for details.
cur_step: Current step id.
eos_id: Token id of the special end of sequence token.
num_beams: Number of beams.
beam_size: Search terminates if the delta between the scores of the active
hyps.
num_hyps_per_beam: Number of hyps in a beam.
valid_eos_max_logit_delta: We allow </s> to terminate a hyp only if its
logit is no more than 'valid_eos_max_logit_delta' away from the logit of
the best candidate.
local_eos_threshold: We allow </s> to terminate a hyp if the local score for
</s> is greater than local_eos_threshold.
merge_paths: If true, hyps which are identical when epsilons are removed
will be combined into a single hyp. The probability for that combined hyp
will be the sum of the probabilities of the component hyps. This can only
be applied for epsilon-emitting models (RNN-T and NT).
is_last_chunk: A tensor of shape [b * k, 1]. Used by neural transducer,
determines whether the current hypothesis reaches the last chunk and
should treat the next end-of-chunk symbol as end-of-sentence.
eoc_id: int, the id of the end of chunk (a.k.a epsilon) token used by neural
transducer models. Only relevant if 'merge_paths' is True or
'is_last_chunk' is provided.
Returns:
out_best_scores: A tensor of shape [b] of updated best scores for each of
the beams.
out_cumulative_scores: A tensor of shape [b * k]. The cumulative score of
the new hyps after the current decoding step.
out_scores: A tensor of shape [b * k] with scores of the token selected.
out_eos_scores: A tensor of shape [b * k] with token scores for the EOS, in
case the hyp was terminated, otherwise 0.0.
out_hyps: A tensor of shape [b * k] with ids of the token selected.
out_prev_hyps: A tensor of shape [b * k] with index to the previous hyps
which was selected.
out_done_hyps: A boolean tensor of shape [b * k] where value indicates
if hyps was terminated.
out_atten_probs: A tensor of shape [b * k, seq_len] which contain the
attention probabilities over the source words against word in the previous
hyps.
out_eos_atten_probs: A tensor of shape [b * k, seq_len] which contains the
attention probabilities over the source against word in the current hyp
which was terminated.
out_all_done: A scalar, whether decoding should terminate for all beams.
out_histories: A tensor of shape [b * k] containing new history hashes for
the active hypotheses. See 'update_histories' docstring for details.
Raises:
ValueError: if inputs are invalid.
"""
num_hyps_per_beam = int(num_hyps_per_beam)
if num_hyps_per_beam <= 0:
raise ValueError("num_hyps_per_beam = {} and must be > 0.".format(
num_hyps_per_beam))
in_scores = tf.convert_to_tensor(in_scores)
in_scores.shape.assert_has_rank(2)
num_classes = in_scores.get_shape()[1]
in_atten_probs = tf.convert_to_tensor(in_atten_probs)
in_atten_probs.shape.assert_has_rank(2)
in_best_scores = tf.convert_to_tensor(in_best_scores)
in_best_scores.shape.assert_has_rank(1)
in_cumulative_scores = tf.convert_to_tensor(in_cumulative_scores)
in_cumulative_scores.shape.assert_has_rank(1)
in_histories = tf.convert_to_tensor(in_histories)
in_histories.shape.assert_has_rank(1)
with tf.name_scope("beam_search_step"):
# For k = num_hyps_per_beam
# First step of beam search is to find the top tokens based on its score.
# Normally we select k+1, where the extra +1 is to make sure we have k
# non-eos tokens to select if EOS token is in the top-k. If path merging is
# on, we actually need to select k+2; this ensures there are k+1 tokens left
# after the merge, at least k of which are not EOS.
# TODO(b/118644069): Avoid casts when there is a XLA op available that takes
# in bfloat16.
num_candidates_per_input_hyp = (num_hyps_per_beam + 2 if merge_paths
else num_hyps_per_beam + 1)
# [b * k, num_candidates_per_input_hyp]
local_score_values, local_indices = xla_ops.top_k_with_unique(
tf.cast(in_scores, tf.float32), k=num_candidates_per_input_hyp)
local_score_values = tf.cast(local_score_values, in_scores.dtype)
# Compute the global score which is sum of the local score, and the
# cumulative scores for each of the hyps.
# [b * k, num_candidates_per_input_hyp]
global_score_values = local_score_values + tf.expand_dims(
in_cumulative_scores, 1)
values_dtype = local_score_values.dtype
is_first_step = tf.cast(tf.equal(cur_step, 0), values_dtype)
# Preprocessing to reorder the tensor from `mod` sharding to `div` so that
# we can use matrix/vector operations to complete the beam search.
# [b * k, num_candidates_per_input_hyp]
global_score_values = reorder_tensor("mod_to_div", global_score_values,
num_beams, num_hyps_per_beam)
local_score_values = reorder_tensor("mod_to_div", local_score_values,
num_beams, num_hyps_per_beam)
local_indices = reorder_tensor("mod_to_div",
local_indices,
num_beams,
num_hyps_per_beam,
max_value=num_classes - 1)
# [b * k, 1]
histories = reorder_tensor("mod_to_div",
tf.expand_dims(in_histories, 1), num_beams,
num_hyps_per_beam)
if is_last_chunk is None:
is_last_chunk = tf.zeros([num_beams * num_hyps_per_beam, 1],
tf.bool)
else:
is_last_chunk = tf.cast(
reorder_tensor(
"mod_to_div",
tf.reshape(is_last_chunk,
[num_beams * num_hyps_per_beam, 1]), num_beams,
num_hyps_per_beam), tf.bool)
# For the first step mask everything but the first row.
# [num_hyps_per_beam]
per_example_mask = tf.concat([
tf.constant([1.0], dtype=values_dtype),
tf.zeros([num_hyps_per_beam - 1], dtype=values_dtype)
], 0)
# [num_hyps_per_beam, num_beams] => [b*k, 1]
mask = tf.reshape(
tf.tile(per_example_mask, tf.expand_dims(num_beams, 0)),
[-1, 1]) * is_first_step + (1.0 - is_first_step)
local_score_values *= mask
global_score_values *= mask
# We add a large negative value for the unmasked values.
per_example_additive_mask = tf.concat([
tf.constant([0.0], dtype=values_dtype),
tf.constant(BEST_SCORES_INIT,
shape=[num_hyps_per_beam - 1],
dtype=values_dtype)
], 0)
additive_mask = tf.reshape(
tf.tile(per_example_additive_mask, tf.expand_dims(num_beams, 0)),
[-1, 1]) * is_first_step
local_score_values += additive_mask
global_score_values += additive_mask
if merge_paths:
with tf.name_scope("merge_paths"):
# Compute new history hashes for each hypothesis + new token.
# [b * k, num_candidates_per_input_hyp]
histories = update_histories(histories,
local_indices,
mask,
epsilon_id=eoc_id)
global_score_values, histories = merge_hyps(
global_score_values, histories, mask, num_beams,
num_hyps_per_beam)
# As we keep num_candidates_per_input_hyp, we have a total of
# num_candidates_per_input_hyp * k hyps active per example.
num_candidate_hyps = num_candidates_per_input_hyp * num_hyps_per_beam
batch_shape = [-1, num_candidate_hyps]
# Reshape score values so that each row corresponds to a particular example.
# [num_beams, num_candidate_hyps]
global_score_values_batch = tf.reshape(global_score_values,
batch_shape)
# First for each beam: Find the top 2 * num_hyps_per_beam candidates.
# The factor of 2 is to be able to process non EOS token ids in the case
# where top scoring token for each hyps is EOS token.
# [k * b, 2 * num_hyps_per_beam]
_, candidates_indices_in_top_k = xla_ops.top_k_with_unique(
tf.cast(global_score_values_batch, tf.float32),
k=2 * num_hyps_per_beam)
# Find the previous hyps of the candidate. We divide here by (k+1) to
# identify which hyps this token came from.
hyps_id = candidates_indices_in_top_k // num_candidates_per_input_hyp
# Add in offset so that we can get the candidate index in the [b * k] space.
offset = tf.expand_dims(tf.range(num_beams) * num_candidate_hyps, 1)
flat_candidates_indices_in_top_k = tf.reshape(
candidates_indices_in_top_k + offset, [-1])
flat_local_indices = tf.reshape(local_indices, [1, -1])
flat_token_scores = tf.reshape(local_score_values, [-1, 1])
flat_global_scores = tf.reshape(global_score_values, [-1, 1])
# Gather the token scores for each of 2*k candidates. We use tf.one_hot()
# followed by a tf.matmul() to speedup gather on TPUs.
total_num_candidates = num_beams * num_candidate_hyps
token_scores_for_beam = tf.reshape(
fast_gather(flat_token_scores, flat_candidates_indices_in_top_k,
total_num_candidates),
[num_beams, 2 * num_hyps_per_beam])
token_scores_for_beam_shape = tf.shape(token_scores_for_beam)
global_scores_for_beam = tf.reshape(
fast_gather(flat_global_scores, flat_candidates_indices_in_top_k,
total_num_candidates), token_scores_for_beam_shape)
# Local indices value's are between [0, vocab_size-1], hence we use the
# slower version of gather.
token_ids_for_beam = tf.reshape(
fast_gather(flat_local_indices,
flat_candidates_indices_in_top_k,
total_num_candidates,
max_value=num_classes - 1,
axis=1), token_scores_for_beam_shape)
# We have access to 2*num_hyps_per_beam hyps per beam.
# We shrink back to num_hyps_per_beam that does not include EOS, and move
# EOS that occurs in top-num_hyps_per_beam to the EOS done matrix.
# To determine the threshold at which eos is allowed to terminate a hyp,
# we need to know the maximum global score for that hyp with any additional
# token. If path merging is *not* enabled, the global_score_values are
# by construction in sorted order, so we can just look at its 0th column. If
# path merging is enabled, the global scores of deleted (merged) hyps break
# the sorted order, which means we have to do a full reduce_max.
if merge_paths:
max_global_score_per_input_hyp = tf.reduce_max(global_score_values,
axis=1,
keepdims=True)
else:
max_global_score_per_input_hyp = global_score_values[:, 0:1]
# [num_beams * num_hyps_per_beam, 1]
global_eos_threshold = (max_global_score_per_input_hyp -
valid_eos_max_logit_delta)
local_eos_threshold_tensor = local_eos_threshold * tf.ones_like(
global_eos_threshold)
# Find EOS in top num_hyps_per_beam token ids. We also treat EOC as EOS if
# the model has indicated this is the last chunk.
local_index_is_eos = tf.equal(local_indices, eos_id)
local_index_is_last_chunk_eoc = tf.math.logical_and(
tf.equal(local_indices, eoc_id), is_last_chunk)
eos_mask = tf.math.logical_and(
tf.math.logical_and(
tf.math.logical_and(
tf.greater(
local_score_values,
tf.tile(local_eos_threshold_tensor,
[1, num_candidates_per_input_hyp])),
tf.greater(
global_score_values,
tf.tile(global_eos_threshold,
[1, num_candidates_per_input_hyp]))),
tf.math.logical_or(local_index_is_eos,
local_index_is_last_chunk_eoc)),
tf.cast(mask, tf.bool))
end_hyps_bool_mask = tf.reshape(tf.reduce_any(eos_mask, 1), [-1, 1])
end_hyps_bool_mask = reorder_tensor("div_to_mod", end_hyps_bool_mask,
num_beams, num_hyps_per_beam)
eos_atten_probs = in_atten_probs * tf.cast(end_hyps_bool_mask,
in_atten_probs.dtype)
eos_atten_probs = tf.reshape(eos_atten_probs,
[num_beams * num_hyps_per_beam, -1])
# A boolean tensor of shape [b * k] where value indicates if hyps was
# terminated.
out_done_hyps = tf.reshape(end_hyps_bool_mask, [-1])
# Scores for EOS token.
eos_float_mask = tf.cast(eos_mask, values_dtype)
eos_local_scores = eos_float_mask * local_score_values
eos_additive_float_mask = (1.0 - eos_float_mask) * BEST_SCORES_INIT
eos_local_scores += eos_additive_float_mask
out_eos_scores = tf.reshape(tf.reduce_max(eos_local_scores, 1),
[-1, 1])
out_eos_scores = tf.reshape(
reorder_tensor("div_to_mod", out_eos_scores, num_beams,
num_hyps_per_beam), [-1])
# A tensor of shape [b] of updated best scores for each of the beams.
eos_global_scores = eos_float_mask * global_score_values
eos_global_scores += eos_additive_float_mask
best_scores = tf.reduce_max(
tf.reshape(eos_global_scores, [num_beams, -1]), 1)
# Following operations are to finds the top num_hyps_per_beam that are
# active.
# Active ones are the ones that do not correspond to EOS termination.
# We keep num_hyps_per_beam * 2 in case every hyps is terminated by EOS id.
# Top K with eos removed.
non_eos_mask = tf.not_equal(token_ids_for_beam, eos_id)
num_candidate_hyps = num_hyps_per_beam * 2 * num_beams
index = tf.where(
non_eos_mask,
tf.reshape(tf.range(num_candidate_hyps, dtype=tf.int32),
token_scores_for_beam_shape),
num_candidate_hyps *
tf.ones(dtype=tf.int32, shape=token_scores_for_beam_shape))
# Unrolled TopK.
sorted_indices = []
# Finds the first num_hyps_per_beam unmasked indexes and stores them in
# concated_index (shape: [num_beams, num_candidate_hyps])
# This is done by iteratively record the min index in each row, and reset
# it to the max, so that next iteration reduce_min returns the 2nd minimum
# index.
for _ in range(num_hyps_per_beam):
min_index = tf.reshape(tf.reduce_min(index, [1]), [num_beams, 1])
sorted_indices.append(min_index)
# Replace position with num_candidate_hyps value.
index = tf.where(
tf.equal(index, min_index),
num_candidate_hyps *
tf.ones(dtype=tf.int32, shape=token_scores_for_beam_shape),
index)
# Post processing ops to output expected tensors.
concated_sorted_indices = tf.concat(sorted_indices, 1)
flat_sorted_indices = tf.reshape(concated_sorted_indices, [-1])
# A tensor of shape [b * k] with scores of the token selected.
out_scores = tf.reshape(
fast_gather(tf.reshape(token_scores_for_beam, [-1, 1]),
flat_sorted_indices, num_candidate_hyps), [-1, 1])
out_scores = tf.reshape(
reorder_tensor("div_to_mod", out_scores, num_beams,
num_hyps_per_beam), [-1])
# Gather the updated histories of selected hypotheses if path merging is
# enabled. Otherwise, the histories are unused, so just output in_histories.
if merge_paths:
flat_histories = tf.reshape(histories, [-1, 1])
# [num_beams, 2 * num_hyps_per_beam]
histories_for_beam = tf.reshape(
fast_gather(flat_histories, flat_candidates_indices_in_top_k,
total_num_candidates), token_scores_for_beam_shape)
out_histories = tf.reshape(
fast_gather(tf.reshape(histories_for_beam, [-1, 1]),
flat_sorted_indices, num_candidate_hyps), [-1, 1])
out_histories = tf.reshape(
reorder_tensor("div_to_mod", out_histories, num_beams,
num_hyps_per_beam), [-1])
else:
out_histories = in_histories
prev_hyps_ids = tf.reshape(
tf.reshape(
fast_gather(tf.reshape(hyps_id, [1, -1]),
flat_sorted_indices,
num_candidate_hyps,
max_value=num_hyps_per_beam,
axis=1), [num_beams, -1]) * num_beams +
tf.expand_dims(tf.range(num_beams), 1), [-1, 1])
prev_hyps_ids = reorder_tensor("div_to_mod",
prev_hyps_ids,
num_beams,
num_hyps_per_beam,
max_value=num_hyps_per_beam)
# A tensor of shape [b * k] with index to the previous hyps which was
# selected.
out_prev_hyps = tf.reshape(prev_hyps_ids, [-1])
# A tensor of shape [b * k, seq_len] which contain the attention
# probabilities over the source words against word in the previous hyps.
out_atten_probs = tf.reshape(
fast_gather(in_atten_probs, out_prev_hyps,
num_beams * num_hyps_per_beam),
[num_beams * num_hyps_per_beam, -1])
sorted_top_k_ids = fast_gather(tf.reshape(token_ids_for_beam, [1, -1]),
flat_sorted_indices,
num_candidate_hyps,
max_value=num_classes - 1,
axis=1)
sorted_top_k_ids = reorder_tensor("div_to_mod",
sorted_top_k_ids,
num_beams,
num_hyps_per_beam,
max_value=num_classes - 1,
axis=1)
# A tensor of shape [b * k] with ids of the token selected.
out_hyps = tf.reshape(sorted_top_k_ids, [-1])
# A tensor of shape [b * k]. The cumulative score of the selected hyps after
# the current decoding step.
out_cumulative_scores = tf.reshape(
fast_gather(tf.reshape(global_scores_for_beam, [-1, 1]),
flat_sorted_indices, num_candidate_hyps), [-1, 1])
out_cumulative_scores = tf.reshape(
reorder_tensor("div_to_mod", out_cumulative_scores, num_beams,
num_hyps_per_beam), [-1])
out_best_scores = tf.maximum(best_scores, in_best_scores)
# A scalar, whether decoding should terminate for all beams.
out_all_done = tf.reshape(
tf.math.logical_not(
tf.reduce_any(
tf.greater(
out_cumulative_scores,
tf.reshape(
tf.tile(
tf.reshape(out_best_scores - beam_size,
[-1, 1]), [1, num_hyps_per_beam]),
[-1])))), [])
return (out_best_scores, out_cumulative_scores, out_scores,
out_eos_scores, out_hyps, out_prev_hyps, out_done_hyps,
out_atten_probs, eos_atten_probs, out_all_done, out_histories)
def FProp(self, theta, input_batch):
"""Embeds source ids and transforms with TransformerStack.
Args:
theta: A `.NestedMap` object containing weights' values of this layer and
its children layers.
input_batch: A `.NestedMap` object containing: ids - The inputs tensor of
shape [batch, time]. paddings - The ids' paddings of shape [batch,
time].
Returns:
A '.NestedMap' object containing:
encoded - The encoded features of shape [time, batch, dim] or [batch,
time, dim], depending p.output_data_format.
padding - The encoded features' padding of shape [time, batch] or
[batch, time].
segment_id - The segmentation of packed inputs of shape [time, batch] or
[batch, time] if it is supported by the model, or None otherwise.
embedded_inputs - The embedded inputs tokens without positional
encodings of shape [time, batch, dim] or [batch, time, dim].
"""
p = self.params
with tf.name_scope(p.name):
# [batch, time]
input_ids = input_batch.ids
# [batch, time]
paddings = input_batch.paddings
# [batch, time]
segment_ids = input_batch.segment_ids if p.packed_input else None
batch = py_utils.GetShape(input_ids)[0]
time = py_utils.GetShape(input_ids)[1]
# Embedding layer.
# [batch, time, dim]
if not p.shared_emb:
input_embs = self.token_emb.EmbLookup(theta.token_emb,
input_ids)
else:
input_embs = self.softmax.EmbLookup(theta.softmax, input_ids)
orig_input_embs = input_embs
# [1, time, dim]
if p.packed_input:
positions = input_batch.segment_pos
position_embs = tf.expand_dims(
self.position_emb.FPropWithPosition(
theta.position_emb, positions), 0)
else:
position_embs = tf.expand_dims(
self.position_emb.FProp(theta.position_emb, time), 0)
# [batch, time, dim]
input_embs += tf.cast(position_embs, tf.bfloat16)
if p.input_dropout_tpl.fprop_dtype:
input_embs = tf.cast(input_embs,
p.input_dropout_tpl.fprop_dtype)
paddings = tf.cast(paddings, p.input_dropout_tpl.fprop_dtype)
input_embs = self.input_dropout.FProp(theta.input_dropout,
input_embs)
# [batch, time, dim]
transformer_input = input_embs
# Explicitly set the input shape of Transformer layers, to avoid
# unknown shape error occurred to tf.einsum on nonTPU devices.
transformer_input = tf.reshape(transformer_input,
[batch, time, p.model_dim])
# Compute self-attention segment mask once.
if p.packed_input:
segment_mask = batch_major_attention.SegmentMask(
segment_ids, segment_ids, dtype=transformer_input.dtype)
else:
segment_mask = tf.zeros([batch, 1, time, time])
encoded, padding = self.transformer_stack.FProp(
theta.transformer_stack, transformer_input, paddings,
segment_mask)
if p.final_layer_norm:
encoded = self.final_ln.FProp(theta.final_ln, encoded)
seq_lengths = tf.cast(tf.reduce_sum(1. - padding, axis=1),
tf.int32)
if p.output_data_format == 'TBC':
encoded = tf.transpose(encoded,
[1, 0, 2]) # [time, batch, dim]
padding = tf.transpose(padding) # [time, batch]
segment_ids = tf.transpose(
segment_ids) if p.packed_input else None
orig_input_embs = tf.transpose(orig_input_embs, [1, 0, 2])
return py_utils.NestedMap(
encoded=encoded,
padding=padding,
seq_lengths=seq_lengths, # used by beam_search_helper.
segment_id=segment_ids,
embedded_inputs=orig_input_embs)
def _GetDefaultPaddings(self, inputs):
"""Gets the default paddings for an input."""
return tf.zeros(
tf.concat([tf.shape(inputs)[:-1], [1]], 0), dtype=inputs.dtype)