def Encode(self, text):
"""Converts string `text` to integer ids and the encoded string.
Encoding includes prefixing the beginning-of-word token to each word.
Returns:
ids: the encoded integer ids.
tokens: the encoded string.
"""
words = tf.sparse.to_dense(tf.strings.split([text]),
default_value='')[0]
num_words = tf.size(words)
ids_ta = tf.TensorArray(tf.int32, 0, dynamic_size=True)
def _WordsToIds(i, words, ids_ta):
encoded_ids = self._EncodeToIds(BOW_STR + words[i])
ids_ta = ids_ta.scatter(
tf.range(ids_ta.size(),
ids_ta.size() + tf.size(encoded_ids)), encoded_ids)
return i + 1, words, ids_ta
_, _, ids_ta = tf.while_loop(lambda i, *_: i < num_words,
_WordsToIds,
loop_vars=(tf.constant(0, tf.int32),
words, ids_ta),
parallel_iterations=30,
back_prop=False)
ids = ids_ta.stack()
return ids, self._TokenToString(ids)
def while_loop(self, cond, body, loop_vars, **kwargs):
"""Wrapper for tf.while_loop that adds summary_tensors to loop vars."""
ctx = TpuSummaryContext.current()
assert ctx is not None, 'must be inside TpuSummaryContext'
assert self.tf_while_loop, 'must be inside self RewriteLoopContext'
outer_summary_tensors = ctx.summary_tensors
ctx.summary_tensors = []
accumulators = [
tf.constant(0, dtype=self.dtype) for _ in range(self.max_loop_vars)
]
size = kwargs['maximum_iterations'] # as if not None
arrays = [
tf.TensorArray(self.dtype, size) for _ in range(self.max_loop_vars)
]
loop_count = tf.constant(0, dtype=tf.int32)
def loop_body(loop_vars, accumulators, arrays, loop_count):
loop_vars = body(*loop_vars)
del ctx.summary_tensors[self.max_loop_vars:]
for i, x in enumerate(ctx.summary_tensors):
if x.while_loop_reduce == 'stack':
arrays[i] = arrays[i].write(loop_count,
tf.cast(x.value, self.dtype))
else:
accumulators[i] += tf.cast(x.value, self.dtype)
loop_count += 1
return loop_vars, accumulators, arrays, loop_count
def loop_cond(loop_vars, accumulators, arrays, loop_count):
del accumulators, arrays, loop_count
return cond(*loop_vars)
loop_vars, accumulators, arrays, loop_count = self.tf_while_loop(
cond=loop_cond,
body=loop_body,
loop_vars=(loop_vars, accumulators, arrays, loop_count),
**kwargs)
for i, x in enumerate(ctx.summary_tensors):
if x.while_loop_reduce == 'stack':
x.value = arrays[i].stack()
elif x.while_loop_reduce == 'mean':
denominator = tf.cast(tf.math.maximum(1, loop_count),
self.dtype)
x.value = accumulators[i] / denominator
else:
x.value = accumulators[i]
ctx.summary_tensors = outer_summary_tensors + ctx.summary_tensors
return loop_vars
def _OutfeedDequeueLoop(self, per_example_tensors, num_loops, num_devices):
"""Process all per-example tensor outfeed data for a TPU sess.run.
Args:
per_example_tensors: dict of key -> tensor as generated by TpuTrainStep.
num_loops: number of times that TpuTrainStep will be executed by TpuTrain.
num_devices: number of TPU cores assigned to this process.
Returns:
A dict of per-example tensors from the latest TpuTrainStep.
"""
if not per_example_tensors:
return tf.no_op()
tensor_shapes = [
py_utils.GetShape(per_example_tensors[key])
for key in sorted(per_example_tensors)
]
tensor_types = [
tf.as_dtype(per_example_tensors[key].dtype)
for key in sorted(per_example_tensors)
]
def LoopBody(i, *input_arrays):
"""Process outfeed data for a single TpuTrainStep.
Args:
i: current loop index.
*input_arrays: One tf.TensorArray per outfeed tensor.
Returns:
i+1 (new index) plus post-write tf.TensorArray handles.
"""
# Outfeed ops execute on each JF node, so they must be located on the
# nodes.
outfeed_devices = []
device_assignment = py_utils.GetTpuDeviceAssignment()
assert device_assignment
for replica in range(device_assignment.num_replicas):
for core in range(device_assignment.num_cores_per_replica):
with tf.device(device_assignment.host_device(
replica, core)):
outfeed_devices.append(
tpu_ops.outfeed_dequeue_tuple(
tensor_types,
tensor_shapes,
device_ordinal=device_assignment.tpu_ordinal(
replica, core)))
offset = i * num_devices
output_arrays = list(input_arrays)
# Each output_array holds a different per-example tensor. We get results
# for each tensor from each TPU for each TpuTrainStep call.
for j in range(len(output_arrays)):
for k in range(len(outfeed_devices)):
output_arrays[j] = output_arrays[j].write(
offset + k, outfeed_devices[k][j])
return tuple([i + 1] + output_arrays)
def LoopCond(i, *output_arrays):
del output_arrays
return i < num_loops
output_arrays = []
for i in range(len(tensor_shapes)):
output_arrays.append(
tf.TensorArray(tensor_types[i],
size=num_loops * num_devices,
element_shape=tensor_shapes[i]))
# Loop once for each time that TpuTrainStep runs.
output_arrays = tf.while_loop(LoopCond,
LoopBody, [0] + output_arrays,
parallel_iterations=1)[1:]
concatenated_arrays = [array.concat() for array in output_arrays]
return dict(zip(sorted(per_example_tensors), concatenated_arrays))
def FarthestPointSampler(points,
padding,
num_sampled_points,
precomputed_squared_distance=None,
num_seeded_points=0,
random_seed=None):
"""Samples num_sampled_points from points using farthest point sampling.
Algorithm:
1. Start by selecting a random point and adding to a selected set.
2. For all remaining points, find the furthest point from those selected.
3. Add furthest point to selected.
4. Repeat 2-3 until num_sampled_points are selected.
More details at https://en.wikipedia.org/wiki/Farthest-first_traversal
This output of this function can be used with tf.batch_gather to extract the
desired points, for example: tf.batch_gather(points, sampled_idx)
Args:
points: floating point tf.Tensor of shape [N, P1, dims]
padding: A floating point tf.Tensor of shape [N, P1] with 0 if the point is
real, and 1 otherwise.
num_sampled_points: integer number of points to sample.
precomputed_squared_distance: optional tf.Tensor of shape [N, P1, P1] of
distances between each point. if None, distances will be computed on the
fly.
num_seeded_points: If num_seeded_points > 0, then the first
num_seeded_points in points are considered to be seeded in the FPS
sampling. Note that we assume that these points are *not* padded, and do
not check padding when seeding them.
random_seed: optional integer random seed to use with all the random ops.
Returns:
A tuple of tf.Tensors (sampled_idx, closest_idx) of types
(tf.int32, tf.int32).
sampled_idx is of shape [N, num_sampled_points] representing the indices
selected using the sampler. This will have range of [0, P1].
closest_idx is of shape [N, P1] representing the indices of the closest
sampled points for each input point. closest_idx is used in PCNN as part of
the pooling operation: each point is assigned to the closest sampled point
and a max is taken over them. This will have a range of [0, P2] with the
index of the closest sampled point that remains.
"""
points = py_utils.HasRank(points, 3)
batch_size, num_points, dims = py_utils.GetShape(points, 3)
points = py_utils.with_dependencies(
[py_utils.assert_greater_equal(num_points, num_sampled_points)],
points)
# Add a tiny bit of noise to the distance matrix or points so all
# points are unique. This will also ensure true repeated points
# like padded points are only selected after all valid points are selected.
if precomputed_squared_distance is not None:
precomputed_squared_distance = py_utils.HasShape(
precomputed_squared_distance, [batch_size, num_points, num_points])
precomputed_squared_distance += tf.random.uniform(
(batch_size, num_points, 1),
minval=1e-6,
maxval=1e-5,
dtype=tf.float32,
seed=random_seed)
else:
points += tf.random.uniform((batch_size, num_points, dims),
minval=1e-6,
maxval=1e-5,
dtype=tf.float32,
seed=random_seed)
# TensorArray to store the sampled indices in the loop.
sampled_idx = tf.TensorArray(tf.int32, num_sampled_points)
# Initialize distance_to_selected to inf for all points.
distance_to_selected = float('inf') * tf.ones((batch_size, num_points))
# For tracking the index to the closest selected point.
closest_idx = tf.zeros((batch_size, num_points), dtype=tf.int32)
# Current loop index counter.
curr_idx = tf.constant(0, dtype=tf.int32)
# Get number of valid points (1 is padded, so num_points - num_padded).
num_valid_points = tf.cast(tf.cast(num_points, dtype=tf.float32) -
tf.reduce_sum(padding, axis=1),
dtype=tf.int32)
def _BodyFn(curr_idx, distance_to_selected, sampled_idx, closest_idx):
"""Loop body for farthest point sampler."""
def _GetRandomRealPoint():
"""Select the first point.
For the first point, we want any random real (non padded) point, so we
create a random values per point, and then set all padded ones to
some large value (more than the maxval). We then take the min per batch
element to get the first points.
Returns:
Tensor containing the index of a random point selected for each example
in the batch.
"""
random_values = tf.random.uniform((batch_size, num_points),
minval=0,
maxval=1,
dtype=tf.float32,
seed=random_seed)
random_values = tf.where(tf.equal(padding, 0.0), random_values,
padding * 10)
return tf.argmin(random_values, axis=1, output_type=tf.int32)
def _GetFurthestPoint():
"""Get point that is furthest from those already selected.
We also bias the sampling towards real points by setting the distance
to padded points negative until we are out of real points.
Returns:
Tensor containing the index of the next farthest point selected for each
example in the batch.
"""
# Set padded points distance to negative so they aren't selected.
padding_masked_distance_to_selected = tf.where(
tf.equal(padding, 0.0), distance_to_selected, -1.0 * tf.ones(
(batch_size, num_points), dtype=tf.float32))
# But only do this when we still have valid points left.
padding_masked_distance_to_selected = tf.where(
tf.less(curr_idx, num_valid_points),
padding_masked_distance_to_selected, distance_to_selected)
return tf.argmax(padding_masked_distance_to_selected,
axis=-1,
output_type=tf.int32)
def _GetSeededPoint():
"""Select a seeded point.
Seeded points are assumed to be at the beginning of the original points.
Returns:
Tensor containing the index of the next seeded point to select for each
example in the batch.
"""
return tf.ones((batch_size, ), dtype=tf.int32) * curr_idx
# Select indices for this loop iteration.
def _Seeded():
return tf.cond(tf.less(curr_idx, num_seeded_points),
_GetSeededPoint, _GetFurthestPoint)
def _Real():
return tf.cond(tf.equal(curr_idx, 0), _GetRandomRealPoint,
_GetFurthestPoint)
new_selected = tf.cond(tf.greater(num_seeded_points, 0), _Seeded,
_Real)
sampled_idx = sampled_idx.write(curr_idx, new_selected)
# Extract the distance to the latest point selected to update
# distance_to_selected.
new_selected_gather_idx = tf.stack(
[tf.range(batch_size), new_selected], axis=1)
if precomputed_squared_distance is not None:
new_distance = tf.gather_nd(precomputed_squared_distance,
new_selected_gather_idx)
else:
new_points = tf.reshape(
tf.gather_nd(points, new_selected_gather_idx),
[batch_size, 1, dims])
new_distance = tf.reshape(
SquaredDistanceMatrix(points, new_points),
[batch_size, num_points])
is_newly_closest = tf.less(new_distance, distance_to_selected)
distance_to_selected = tf.minimum(distance_to_selected, new_distance)
# Track the index to the closest selected point.
new_selected_tiled = tf.tile([[curr_idx]], [batch_size, num_points])
closest_idx = tf.cond(
tf.equal(curr_idx, 0),
# At the first loop iteration, the init points are the closest.
lambda: new_selected_tiled,
# Otherwise, update with the new points based on the distances.
lambda: tf.where(is_newly_closest, new_selected_tiled, closest_idx)
)
return curr_idx + 1, distance_to_selected, sampled_idx, closest_idx
_, _, sampled_idx, closest_idx = tf.while_loop(
lambda curr_idx, *args: tf.less(curr_idx, num_sampled_points),
_BodyFn,
loop_vars=(curr_idx, distance_to_selected, sampled_idx, closest_idx),
back_prop=False,
maximum_iterations=num_sampled_points)
sampled_idx = sampled_idx.stack() # num_sampled_points x n
sampled_idx = tf.transpose(sampled_idx, [1, 0])
if isinstance(batch_size, int) and isinstance(num_sampled_points, int):
sampled_idx.set_shape((batch_size, num_sampled_points))
return sampled_idx, closest_idx
def flat_beam_search(batch_size,
beam_size,
max_steps,
dec_callback,
dec_state,
bos_id=1,
eos_id=2,
length_norm_alpha=0.8,
beam_gap=3.0,
top_k_fn=tf.math.top_k,
prefix=None,
prefix_len=None,
fprop_dtype=tf.float32,
ext_size=0,
nbest_size=None,
debug=True):
"""Flat beam search.
Args:
batch_size: batch size
beam_size: beam size limit in number of hyps
max_steps: max steps
dec_callback: decoder callback (see above)
dec_state: decoder state
bos_id: <s> token id
eos_id: </s> token id
length_norm_alpha: length normalization parameter
beam_gap: early stopping threshold; None to disable
top_k_fn: top_k function to call
prefix: (optional) int32 tensor [batch_size, prefix_max]
prefix_len: (optional) int32 tensor [batch_size]
fprop_dtype: fprop dtype
ext_size: int >= beam_size, extension buffer size
nbest_size: number of returned hyps, default is beam_size
debug: log intermediate vlaues with tpu_summary.tensor()
Returns:
(loop_vars, dec_state, nbest) where
nbest = (topk_ids, topk_len, topk_score)
"""
assert beam_size > 0
assert batch_size > 0
assert max_steps > 0
buf_size = beam_size * max_steps
output_len = max_steps
if prefix is None:
assert prefix_len is None
prefix = tf.zeros([batch_size, beam_size], dtype=tf.int32)
prefix += tf.one_hot(0, beam_size, dtype=tf.int32) * bos_id
prefix_len = tf.ones([batch_size], dtype=tf.int32)
else:
assert int(prefix.shape[0]) == batch_size, (batch_size, prefix.shape)
assert int(prefix_len.shape[0]) == batch_size, (batch_size,
prefix_len.shape)
output_len += int(prefix.shape[1])
if debug:
tpu_summary.tensor('prefix', prefix)
tpu_summary.tensor('prefix_len', prefix_len)
with tf.name_scope('init_state'):
t = tf.constant(0)
tgt_id = tf.zeros([batch_size, beam_size], dtype=tf.int32)
tgt_id += bos_id
tgt_pos = tf.zeros([batch_size, beam_size], dtype=tf.int32)
tgt_mask = tf.zeros([batch_size, beam_size, buf_size],
dtype=fprop_dtype)
tgt_mask += tf.one_hot(tf.range(beam_size),
buf_size,
dtype=fprop_dtype)
hyp_score = tf.zeros([batch_size, beam_size], dtype=fprop_dtype)
# penalize all hyps except the first
hyp_score -= tf.cast(tf.range(beam_size, dtype=tf.float32) * 1e5,
dtype=fprop_dtype)
nbest_size = nbest_size or beam_size
nbest_score = tf.zeros([batch_size, nbest_size], dtype=fprop_dtype)
nbest_score -= 1e9
nbest_score_norm = nbest_score
nbest_mask = tf.zeros([batch_size, nbest_size, buf_size],
dtype=fprop_dtype)
with tf.name_scope('init_ext'):
# Initialize the extension buffer.
#
# Extension buffer stores a (potentially large) set of 'extensions',
# which consist of a hypothesis (represented by ext_mask) and next token
# (represented by ext_id). At each decoder iteration, top_k extensions
# from each hypothesis are added to the buffer and sorted by score.
#
# Then top beam_size extensions are removed from the buffer and used
# in the next decoder iteration. And top 'ext_size' remaining extensions
# are carried over to be possibly evaluated at a later step.
#
# As a result of this manipulation, the decoder is no longer restricted
# to always compare hyps of the same token length at each iteration.
# In particular, for a fixed length N it can generate more than beam_size
# terminated hyps.
#
# Setting ext_size = 0 disables this feautre.
if ext_size:
ext_id = tf.zeros([batch_size, ext_size], dtype=tf.int32)
ext_score = tf.zeros([batch_size, ext_size], dtype=fprop_dtype)
ext_score -= 1e9
ext_mask = tf.zeros([batch_size, ext_size, buf_size],
dtype=fprop_dtype)
else:
ext_size = ext_id = ext_score = ext_mask = 0
with tf.name_scope('init_prefix'):
# rename prefix->pfx for shorter variables
pfx = tf.cast(prefix, tf.int32)
pfx_len = tf.cast(prefix_len, tf.int32)
del prefix, prefix_len
# Before the first call to dec_callback() the prefix shall be packed into
# the tgt_id buffer as follows:
#
# [ P P P P P P - - - - - - P* - - - ] ^
# [ P P P P P P P P P P - - P* - - - ] | batch
# [ P - - - - - - - - - - - P* - - - ] V
# |<---- prefix len ----> |<-- beam -->
#
# The last meaningful token in the prefix (P*)
# must be located at the same position in all batch rows.
#
# We then make one dec_callback() with full prefix (minus P*)
# which will populate the initial dec_state
# (for transformer -- self-attention key/value cache)
#
# The last block [batch, beam] then becomes the first tgt_id for the loop.
pfx_max = int(pfx.shape[1])
pfx_mul = pfx_max // beam_size
assert pfx_max == pfx_mul * beam_size, (pfx_max, pfx_mul, beam_size)
pfx_time = tf.range(pfx_max)
pfx_pad = tf.cast(
tf.less(tf.expand_dims(pfx_time, 0),
tf.expand_dims(pfx_len - 1, 1)), tf.int32)
pfx_id = pfx * pfx_pad
pfx_last = einsum_i32(
'BT,BT->B', pfx, tf.one_hot(pfx_len - 1,
pfx_max,
dtype=fprop_dtype))
buf_time = tf.range(buf_size)
pfx_time_mask = tf.cast(
tf.less_equal(tf.expand_dims(buf_time, 0),
tf.expand_dims(pfx_time, 1)), fprop_dtype)
pfx_mask = tf.einsum('BQ,QK->BQK', tf.cast(pfx_pad, fprop_dtype),
pfx_time_mask)
pfx_segment_id = pfx_pad
pfx_pos = pfx_time * pfx_pad
if debug:
tpu_summary.tensor('pfx_id', pfx_id)
tpu_summary.tensor('pfx_len', pfx_len)
tpu_summary.tensor('pfx_pos', pfx_pos)
tpu_summary.tensor('pfx_last', pfx_last)
# Now call decoder with prefix minus P*:
# 'dec_state' now shall contain the key/value cache for prefix tokens
# (for transformer models), and 'logits' we can either discard or
# roll into the initial hyp_score. Discard is simpler.
with tf.name_scope('prefix_fprop'):
# TODO(krikun): remove extra type checks
assert (pfx_id.dtype == tf.int32), (pfx_id.dtype)
assert (pfx_segment_id.dtype == tf.int32), (pfx_segment_id.dtype)
assert (pfx_pos.dtype == tf.int32), (pfx_pos.dtype)
assert (pfx_mask.dtype == fprop_dtype), (pfx_mask.dtype)
assert (t.dtype == tf.int32), (t.dtype)
logits, dec_state = dec_callback(pfx_id, pfx_segment_id, pfx_pos,
pfx_mask, dec_state, t)
del logits
# Now construct the initial state for the rest of the beam search loop.
# 'tgt_id' is simply 'pfx_last' padded to [batch, beam] shape
# 'tgt_pos' is different for each batch row and is equal to prefix_len
# 'tgt_segment_id' always 1 (no packing)
# 'hyp_score' is 0 for beam=0 and negative for beam>=1
tgt_id = tf.zeros([batch_size, beam_size], tf.int32) + tf.expand_dims(
pfx_last, 1)
tgt_pos = tf.zeros([batch_size, beam_size], tf.int32) + tf.expand_dims(
(pfx_len - 1), 1)
hyp_score = tf.zeros(
[batch_size, beam_size], dtype=fprop_dtype) - tf.cast(
tf.range(beam_size, dtype=tf.float32) * 1e5, dtype=fprop_dtype)
# TODO(krikun) Here we make initial 't' constant and determined by the
# shape of the prefix tensor 'pfx_max'. It is possible to make it dynamic
# as t ~ max(pfx_len) / beam_size and this will more steps for beam search
# however 'max' results in a very slow all-to-all for 'max' on 16x16
# and variable number of decoder steps may result in bad latency.
t = tf.cast(tf.math.ceil(pfx_max / beam_size), tf.int32)
# Initial tgt_mask is such that each token P* has attention on itself
# (as usual) and on all prefix tokens before it, which are not padding.
tgt_mask = tf.zeros([batch_size, beam_size, buf_size],
dtype=fprop_dtype)
tgt_mask += tf.cast(
tf.expand_dims(
tf.pad(pfx_pad, [[0, 0], [0, (buf_size - pfx_max)]]), 1),
fprop_dtype)
tgt_mask += tf.one_hot(tf.range(beam_size) + t * beam_size,
buf_size,
dtype=fprop_dtype)
if debug:
tpu_summary.tensor('tgt_id', tgt_id)
tpu_summary.tensor('tgt_pos', tgt_pos)
tpu_summary.tensor('tgt_mask', tgt_mask)
tpu_summary.tensor('t', t)
with tf.name_scope('init_hist'):
# h_tgt_id is used to recover topk_ids from nbest_mask
h_tgt_id = tf.TensorArray(dtype=tf.int32, size=max_steps)
h_tgt_pos = tf.TensorArray(dtype=tf.int32, size=max_steps)
# When non-trivial prefix is present we also write prefix ids to
# h_tgt_id so that the full sequence including prefix can be recovered
# by unmask() below. When prefix is empty, pfx_id shape is [batch, 0]
# and the loop below becomes a no-op.
# TODO(krikun): maybe a tf.while_loop is more appropriate here.
for i, x_i in enumerate(tf.split(pfx_id, pfx_mul, 1)):
h_tgt_id = h_tgt_id.write(i, x_i)
for i, x_i in enumerate(tf.split(pfx_pos, pfx_mul, 1)):
h_tgt_pos = h_tgt_pos.write(i, x_i)
hist = (h_tgt_id, h_tgt_pos)
tf.logging.info('hist=%r', hist)
nbest_hyps = (nbest_mask, nbest_score, nbest_score_norm)
tf.logging.info('nbest_hyps=%r', nbest_hyps)
ext = (ext_id, ext_score, ext_mask)
tf.logging.info('ext=%r', ext)
loop_vars = (t, tgt_id, tgt_pos, tgt_mask, hyp_score, nbest_hyps, ext,
hist)
tf.logging.info('loop_vars=%r', loop_vars)
def loop_step(loop_vars, dec_state): # pylint: disable=missing-docstring
tf.logging.info('loop_vars=%r', loop_vars)
tf.logging.info('dec_state=%r', dec_state)
(t, tgt_id, tgt_pos, tgt_mask, hyp_score, nbest_hyps, ext,
hist) = loop_vars
(ext_id, ext_score, ext_mask) = ext
(h_tgt_id, h_tgt_pos) = hist
h_tgt_id = h_tgt_id.write(t, tgt_id, name='h_tgt_id')
h_tgt_pos = h_tgt_pos.write(t, tgt_pos, name='h_tgt_pos')
# not using tf.ones() here because of XLA compilation error
tgt_segment_id = tgt_id * 0 + 1
logits, dec_state = dec_callback(tgt_id, tgt_segment_id, tgt_pos,
tgt_mask, dec_state, t)
# take predicted EOS score for each hyp and compute normalized score
eos_score = hyp_score + tf.cast(logits[:, :, eos_id], hyp_score.dtype)
def length_norm(t):
t = tf.cast(t, fprop_dtype)
alpha = length_norm_alpha
tf.logging.info('length_norm.alpha=%r', alpha)
return tf.math.pow((t + 5.) / 5., alpha)
hyp_len = tgt_pos - tf.expand_dims((pfx_len - 1), -1)
eos_score_norm = eos_score / length_norm(hyp_len)
# update the n-best list
nbest_hyps = update_nbest(nbest_hyps,
(tgt_mask, hyp_score, eos_score_norm))
if debug:
tpu_summary.tensor('eos_score', eos_score)
tpu_summary.tensor('hyp_len', hyp_len)
# take top k tokens for each hyp
k = beam_size
with tf.name_scope('topk1'):
top_score, top_id = top_k_fn(logits, k)
top_score = tf.cast(top_score, fprop_dtype)
top_score += tf.expand_dims(hyp_score, -1)
top_score -= 1e9 * tf.cast(tf.equal(top_id, eos_id), fprop_dtype)
top_score = tf.reshape(top_score, [batch_size, beam_size * k])
top_id = tf.reshape(top_id, [batch_size, beam_size * k])
top_mask = tf.repeat(tgt_mask, beam_size, 1)
if debug:
tpu_summary.tensor('top_id', top_id)
tpu_summary.tensor('top_score', top_score)
# tpu_summary.tensor('top_mask', top_mask)
with tf.name_scope('update_ext'):
# combine top k tokens with extension buffer (if any)
if ext_size:
ext_id = tf.concat([ext_id, top_id], 1)
ext_score = tf.concat([ext_score, top_score], 1)
ext_mask = tf.concat([ext_mask, top_mask], 1)
else:
ext_id, ext_score, ext_mask = top_id, top_score, top_mask
# sort by score
ext_score, i = tf.math.top_k(ext_score, ext_size + beam_size)
i1 = tf.one_hot(i, ext_size + beam_size * k, dtype=fprop_dtype)
ext_mask = tf.einsum('bkt,bjk->bjt', ext_mask, i1)
ext_id = einsum_i32('bk,bjk->bj', ext_id, i1)
# pick top beam_size extensions to evaluate at next iteration
if ext_size:
hyp_score = ext_score[:, :beam_size]
ext_score = ext_score[:, beam_size:]
tgt_id = ext_id[:, :beam_size]
ext_id = ext_id[:, beam_size:]
tgt_mask = ext_mask[:, :beam_size]
ext_mask = ext_mask[:, beam_size:]
else:
hyp_score, tgt_id, tgt_mask = ext_score, ext_id, ext_mask
ext_score = ext_id = ext_mask = 0
tgt_pos = tf.reduce_sum(tgt_mask, -1)
tgt_pos = tf.cast(tgt_pos, tf.int32)
t += 1
with tf.name_scope('tgt_mask_extend'):
tgt_mask += tf.one_hot(tf.range(beam_size) + t * beam_size,
buf_size,
dtype=fprop_dtype)
ext = (ext_id, ext_score, ext_mask)
hist = (h_tgt_id, h_tgt_pos)
loop_vars = (t, tgt_id, tgt_pos, tgt_mask, hyp_score, nbest_hyps, ext,
hist)
tf.logging.info('loop_vars=%r', loop_vars)
tf.logging.info('dec_state=%r', dec_state)
return loop_vars, dec_state
def loop_cond(loop_vars, dec_state): # pylint: disable=missing-docstring
tf.logging.info('loop_vars=%r', loop_vars)
tf.logging.info('dec_state=%r', dec_state)
if beam_gap is None:
(t, _, _, _, _, _, _, _) = loop_vars
return t < max_steps
else:
(t, _, _, _, _, nbest_hyps, _, _) = loop_vars
(_, nbest_score, _) = nbest_hyps
# stop early if all current hyps are significantly worse than nbest
diff = tf.reduce_min(
tf.reduce_min(nbest_score, -1) - tf.reduce_max(hyp_score, -1))
return tf.math.logical_and(t < max_steps, diff < beam_gap)
with tf.name_scope('flat_beam_search_loop'):
(loop_vars, dec_state) = tf.while_loop(loop_cond,
loop_step,
loop_vars=(loop_vars,
dec_state),
back_prop=False,
swap_memory=False,
maximum_iterations=max_steps)
# flatten all tensorarrays into tensors
(t, tgt_id, tgt_pos, tgt_mask, hyp_score, nbest_hyps, ext,
hist) = loop_vars
(nbest_mask, nbest_score, nbest_score_norm) = nbest_hyps
(h_tgt_id, h_tgt_pos) = hist
h_tgt_id = h_tgt_id.stack()
h_tgt_pos = h_tgt_pos.stack()
hist = (h_tgt_id, h_tgt_pos)
loop_vars = (t, tgt_id, tgt_pos, tgt_mask, hyp_score, nbest_hyps, ext,
hist)
# recover topk_ids from nbest_mask and tgt_id history
h = tf.transpose(h_tgt_id, [1, 0, 2])
h = tf.reshape(h, [batch_size, buf_size])
def unmask(h, m):
with tf.name_scope('unmask'):
tpu_summary.tensor('unmask_h', h)
tpu_summary.tensor('unmask_m', m)
t = tf.cumsum(m, -1) * m - 1
mh = einsum_i32('bkt,bt->bkt', m, h)
t2 = tf.one_hot(tf.cast(t, tf.int32),
output_len,
dtype=fprop_dtype)
x = einsum_i32('bkt,bktT->bkT', mh, t2)
return tf.cast(x, h.dtype)
topk_ids = unmask(h, nbest_mask)
topk_len = tf.reduce_sum(nbest_mask, -1)
topk_len = tf.cast(topk_len, tf.int32)
# add eos, because nbest_mask does not encode eos
topk_ids += eos_id * tf.one_hot(topk_len, output_len, dtype=tf.int32)
topk_len += 1
topk_len = tf.minimum(topk_len, output_len)
topk_score = nbest_score_norm
nbest = (topk_ids, topk_len, topk_score)
return loop_vars, dec_state, nbest
def _StringsToIdsImpl(self, strs, max_length, append_eos, languages):
"""Takes a tensor of strings and returns id/padding tensors.
This generates `token_ids`, `target_ids`, and `paddings` in the format that
is expected for tokenizers. This performs padding to a fixed length and
appends the end-of-sentence token as appropriate.
Args:
strs: a string Tensor.
max_length: a python integer. The second dimension of the returned arrays.
All sequences are padded or truncated to that length.
append_eos: a python bool. See `BaseTokenizer` for explanation.
languages: A vector of strings with the same length as `strs`.
Returns:
A tuple of 3 tensors:
- token_ids: a tensor of sequences of WPM ids starting with SOS. Sequences
always end with EOS unless the sequence exceeds the maximum length.
Always padded with EOS.
- target_ids: a tensor of sequences of WPM ids not starting with SOS
but ending with EOS. Always padded with EOS.
- paddings: a tensor of floats indicating, at each position, whether
the corresponding position is padded.
"""
p = self.params
if append_eos is None:
append_eos = p.append_eos
batch_size = py_utils.GetShape(strs)[0]
token_ids_ta = tf.TensorArray(tf.int32, batch_size)
target_ids_ta = tf.TensorArray(tf.int32, batch_size)
paddings_ta = tf.TensorArray(tf.float32, batch_size)
def _TokenizeOneSentence(i, strs, token_ids_ta, target_ids_ta,
paddings_ta):
"""Tokenizes a single sentence."""
ids, _ = self._wpm_encoder.Encode(strs[i])
if append_eos:
ids = tf.concat([ids, [self.eos_id]], axis=0)
# This truncates after the eos is added, so some sentences might
# not have </s> at the end.
token_ids_ta = token_ids_ta.write(
i,
py_utils.PadOrTrimTo(tf.concat([[self.sos_id], ids], axis=0),
[max_length], self.eos_id))
target_ids_ta = target_ids_ta.write(
i, py_utils.PadOrTrimTo(ids, [max_length], self.eos_id))
paddings_ta = paddings_ta.write(
i,
py_utils.PadOrTrimTo(tf.zeros_like(ids, dtype=tf.float32),
[max_length], 1.))
return i + 1, strs, token_ids_ta, target_ids_ta, paddings_ta
_, _, token_ids_ta, target_ids_ta, paddings_ta = tf.while_loop(
lambda i, *_: i < batch_size,
_TokenizeOneSentence,
loop_vars=(tf.constant(0, tf.int32), strs, token_ids_ta,
target_ids_ta, paddings_ta),
parallel_iterations=30,
back_prop=False)
token_ids = token_ids_ta.stack()
target_ids = target_ids_ta.stack()
paddings = paddings_ta.stack()
if not p.pad_to_max_length:
maxlen = tf.cast(
tf.round(tf.reduce_max(tf.reduce_sum(1.0 - paddings, axis=1))),
tf.int32)
token_ids = token_ids[:, :maxlen]
target_ids = target_ids[:, :maxlen]
paddings = paddings[:, :maxlen]
return token_ids, target_ids, paddings
def Sample(self,
decoder_theta,
encoder_outputs,
random_seed,
init_state_callback,
pre_step_callback,
post_step_callback,
init_step_ids=None):
"""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.
init_step_ids: [batch], optional init step ids, default to SOS.
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
assert p.top_k >= 0
assert p.num_hyps_per_beam >= 1
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,
p.num_hyps_per_beam)
# 'recurrent_theta' represents all cross-timestep information used by the
# recurrent loop below, including layer theta and encoder outputs.
recurrent_theta = py_utils.NestedMap(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=init_step_ids if init_step_ids is not None else tf.fill(
[batch], tf.cast(p.target_sos_id, tf.int32)),
bs_state=bs_state)
if p.use_recurrent:
inputs = py_utils.NestedMap(
dummy=tf.zeros([p.target_seq_len, batch]))
else:
inputs = py_utils.NestedMap(
ids=tf.TensorArray(dtype=tf.int32, size=p.target_seq_len),
logits=tf.TensorArray(dtype=bs_result.log_probs.dtype,
size=p.target_seq_len),
)
def Step(recurrent_theta, state0, inputs):
"""Computes one decoder step."""
if p.use_recurrent:
del inputs
with tf.name_scope('single_sampler_step'):
# Compute logits and states.
bs_result, bs_state1 = pre_step_callback(
decoder_theta,
recurrent_theta.encoder_outputs,
tf.expand_dims(state0.ids, 1), # [batch, 1].
state0.bs_state,
num_hyps_per_beam=p.num_hyps_per_beam)
batch = tf.shape(bs_result.log_probs)[0]
state1 = py_utils.NestedMap(timestep=state0.timestep + 1)
state1.logits = bs_result.log_probs
if p.top_k > 0:
topk_logits, topk_ids = tf.math.top_k(state1.logits,
k=p.top_k)
sample_logits = tf.nn.log_softmax(
topk_logits) if p.top_k_renormalize else topk_logits
else:
sample_logits = state1.logits
# Sample ids from logits. [batch].
ids = tf.reshape(
tf.random.stateless_categorical(
sample_logits / p.temperature,
num_samples=1,
seed=tf.stack(
[recurrent_theta.random_seed, state0.timestep]),
dtype=state0.ids.dtype,
name='sample_next_id'), [batch])
state1.ids = tf.gather(topk_ids, ids, axis=1,
batch_dims=1) if p.top_k > 0 else ids
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(
decoder_theta, recurrent_theta.encoder_outputs, state1.ids,
bs_state1)
if p.use_recurrent:
return state1, py_utils.NestedMap()
else:
inputs.ids = inputs.ids.write(state0.timestep, state1.ids)
inputs.logits = inputs.logits.write(state0.timestep,
state1.logits)
return (recurrent_theta, state1, inputs)
if p.use_recurrent:
def StopFn(t, theta, state):
del t, theta # Unused: this stop function only uses the state ids.
return tf.equal(state.ids, p.target_eos_id)
else:
def StopFn(recurrent_theta, state, inputs):
del recurrent_theta, inputs
return tf.logical_not(
tf.reduce_all(tf.equal(state.ids, p.target_eos_id)))
if p.use_stop_fn:
stop_fn = StopFn
else:
stop_fn = None
if p.use_recurrent:
accumulated_states, _ = recurrent.Recurrent(
recurrent_theta,
recurrent_state0,
inputs,
Step,
stop_fn=stop_fn,
allow_implicit_capture=True)
else:
loop_vars = (recurrent_theta, recurrent_state0, inputs)
(_, _, accumulated_states) = tf.while_loop(
StopFn,
Step,
loop_vars=loop_vars,
shape_invariants=_GetShapes(loop_vars, none_shapes=True),
back_prop=False,
maximum_iterations=p.target_seq_len)
accumulated_states.ids = accumulated_states.ids.stack()
accumulated_states.logits = accumulated_states.logits.stack()
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 ComputePredictions(self,
encoder_outputs,
pronunciations,
is_inference=False):
"""Computes the predictions from the encoder_outputs, updating losses.
Despite the name, this function does the bulk of the decoding and loss
computation, incrementing the loss at each time step.
Args:
encoder_outputs: a NestedMap consisting of outputs of the
FeatureNeighborhoodEncoder with encoded - encoding of the input
spelling
neighbor_pronunciations_encoded - encodings of the neighbor prons
neighbor_pronunciations_encoded - encodings of the neighbor spellings
state - encoder state to which has been added dec_input - seed output
for the decoder [*, 1] tensor consisting of sentence start indices
(corresponding to "<s>")
pronunciations: NestedMap with pronunciations - [*, max_pronunciation_len]
tensor of pronunciations
is_inference: If False then uses teacher forcing else does autoregression.
Returns:
NestedMap with loss, per_sequence_losses,labels, a
[*, max_pronunciation_len] tensor of predictions, and attention
([*, max_pronunciation_len, max_spelling_len]), and
neighbor_attention ([*, max_pronunciation_len, max_neighbors])
tensors, along with the raw batch passed through from the encoder.
"""
p = self.params
targets = pronunciations.pronunciations
t_len = int(targets.get_shape().as_list()[1])
t_idx = tf.constant(0)
attention = tf.TensorArray(dtype=tf.float32, size=t_len)
neighbor_attention = tf.TensorArray(dtype=tf.float32, size=t_len)
outputs = tf.TensorArray(dtype=tf.float32, size=t_len)
loop_cond = lambda t_idx, ts, *_: tf.less(t_idx, t_len)
dec_input = tf.convert_to_tensor([p.start] * p.input.batch_size)
state = encoder_outputs.state
# pylint: disable=missing-docstring
def loop_body(t_idx, dec_input, attention, neighbor_attention, state,
outputs):
decoder_result = self.Decode(encoder_outputs, dec_input, state)
outputs = outputs.write(t_idx, decoder_result.predictions)
attention = attention.write(t_idx,
decoder_result.attention_weights)
neighbor_attention = neighbor_attention.write(
t_idx,
tf.cast(decoder_result.neighbor_attention_weights,
dtype=tf.float32))
if is_inference:
dec_input = tf.cast(tf.argmax(decoder_result.predictions, 1),
tf.int32)
else:
dec_input = targets[:, t_idx]
t_idx = t_idx + 1
state = decoder_result.state
return t_idx, dec_input, attention, neighbor_attention, state, outputs
_, _, attention, neighbor_attention, state, outputs = tf.while_loop(
loop_cond,
loop_body,
loop_vars=[
t_idx, dec_input, attention, neighbor_attention, state, outputs
])
outputs = tf.transpose(outputs.stack(), [1, 0, 2])
labels = tf.argmax(outputs, axis=-1)
mask = tf.cast(tf.math.logical_not(tf.math.equal(targets, 0)),
dtype=tf.float32)
loss = self._loss_object(targets, outputs, sample_weight=mask)
loss = tf.reduce_sum(loss, axis=1)
per_sequence_losses = (loss / t_len)
loss = tf.reduce_mean(per_sequence_losses)
predictions = py_utils.NestedMap()
predictions.loss = loss
predictions.per_sequence_losses = per_sequence_losses
predictions.labels = labels
predictions.attention = tf.transpose(tf.squeeze(attention.stack()),
perm=[1, 0, 2])
if p.use_neighbors:
predictions.neighbor_attention = tf.transpose(tf.squeeze(
neighbor_attention.stack()),
perm=[1, 0, 2])
else:
predictions.neighbor_attention = tf.squeeze(
neighbor_attention.stack())
# Expose this for subsequent data analysis
predictions.batch = encoder_outputs.batch
return predictions
def _StringsToIdsImpl(self, strs, max_length, append_eos, languages):
del languages
p = self.params
if append_eos is None:
append_eos = p.append_eos
batch_size = py_utils.GetShape(strs)[0]
token_ids_ta = tf.TensorArray(tf.int32, batch_size)
target_ids_ta = tf.TensorArray(tf.int32, batch_size)
paddings_ta = tf.TensorArray(tf.float32, batch_size)
def _TokenizeOneSentence(i, text, token_ids_ta, target_ids_ta,
paddings_ta):
"""Tokenizes a single sentence."""
if tf.is_tensor(i):
text_i = tf.gather(text, i)
else:
text_i = text[i]
ids = self._tokenizer.tokenize(text_i).merge_dims(0, -1)
ids.set_shape([None])
if append_eos:
ids = tf.concat([ids, [self.eos_id]], axis=0)
sos_ids = tf.concat([[self.sos_id], ids], axis=0)
if p.prepend_sos:
ids = sos_ids
# This truncates after the EOS is added, so some sentences might
# not have EOS at the end.
token_ids_ta = token_ids_ta.write(
i, py_utils.PadOrTrimTo(sos_ids, [max_length], 0))
target_ids_ta = target_ids_ta.write(
i, py_utils.PadOrTrimTo(ids, [max_length], 0))
paddings_ta = paddings_ta.write(
i,
py_utils.PadOrTrimTo(tf.zeros_like(ids, dtype=tf.float32),
[max_length], 1.))
return i + 1, strs, token_ids_ta, target_ids_ta, paddings_ta
_, _, token_ids_ta, target_ids_ta, paddings_ta = tf.while_loop(
lambda i, *_: i < batch_size,
_TokenizeOneSentence,
loop_vars=(tf.constant(0, tf.int32), strs, token_ids_ta,
target_ids_ta, paddings_ta),
parallel_iterations=30,
back_prop=False)
token_ids = token_ids_ta.stack()
target_ids = target_ids_ta.stack()
paddings = paddings_ta.stack()
if not p.pad_to_max_length:
maxlen = tf.cast(
tf.round(tf.reduce_max(tf.reduce_sum(1.0 - paddings, axis=1))),
tf.int32)
token_ids = token_ids[:, :maxlen]
target_ids = target_ids[:, :maxlen]
paddings = paddings[:, :maxlen]
return token_ids, target_ids, paddings
