def infer(self,
features=None,
decode_length=50,
beam_size=1,
top_beams=1,
alpha=0.0,
use_tpu=False):
"""Returns the targets and their log probabilities."""
del decode_length, beam_size, top_beams, alpha, use_tpu
assert features is not None
self._fill_problem_hparams_features(features)
# Run the model
self.hparams.force_full_predict = True
with tf.variable_scope(self.name):
logits, _ = self.model_fn(features)
assert len(logits.shape) == 5 # [batch, time, 1, 1, vocab]
logits = tf.squeeze(logits, [2, 3])
#import pdb; pdb.set_trace()
# Compute the log probabilities
log_probs = common_layers.log_prob_from_logits(logits)
targets = features["targets"]
assert len(targets.shape) == 4 # [batch, time, 1, 1]
targets = tf.squeeze(targets, [2, 3])
# Slice out the log_probs of the targets
log_probs = common_layers.index_last_dim_with_indices(log_probs, targets)
# return log-probs instead of beam-score
return {"outputs": targets, "scores": log_probs}
def testIndexLastDimWithIndices(self):
x = np.array([[2., 3., 4., 5.], [6., 7., 8., 9.]])
indices = np.array([2, 0])
x_idx = common_layers.index_last_dim_with_indices(x, indices)
expected = np.array([4., 6.])
self.assertAllEqual(expected, self.evaluate(x_idx))
def testIndexLastDimWithIndices(self):
x = np.array([[2., 3., 4., 5.],
[6., 7., 8., 9.]])
indices = np.array([2, 0])
x_idx = common_layers.index_last_dim_with_indices(x, indices)
expected = np.array([4., 6.])
self.assertAllEqual(expected, self.evaluate(x_idx))
def compute_uncertainty_reward(logits, predictions):
"""Uncertainty reward based on logits."""
# TODO(rsepassi): Add support for L1/L2 loss models. Current code only
# works for softmax models.
vocab_size = logits.shape[-1]
assert vocab_size > 1
log_probs = common_layers.log_prob_from_logits(logits)
max_log_probs = common_layers.index_last_dim_with_indices(
log_probs, predictions)
# Threshold
neg_log_prob = tf.nn.relu(-max_log_probs - 0.02)
# Sum across all but the batch dimension
reduce_dims = list(range(len(neg_log_prob.shape)))[1:]
summed = tf.reduce_sum(neg_log_prob, axis=reduce_dims)
return summed / 10
def compute_uncertainty_reward(logits, predictions):
"""Uncertainty reward based on logits."""
# TODO(rsepassi): Add support for L1/L2 loss models. Current code only
# works for softmax models.
vocab_size = logits.shape[-1]
assert vocab_size > 1
log_probs = common_layers.log_prob_from_logits(logits)
max_log_probs = common_layers.index_last_dim_with_indices(log_probs,
predictions)
# Threshold
neg_log_prob = tf.nn.relu(-max_log_probs - 0.02)
# Sum across all but the batch dimension
reduce_dims = list(range(len(neg_log_prob.shape)))[1:]
summed = tf.reduce_sum(neg_log_prob, axis=reduce_dims)
return summed / 10
