def test_compute_compute_match_stats(self):
with tf.Graph().as_default():
predictions = tf.to_int32([1, 2, 3, 9])
labels = tf.to_int32([[0, 1, 5, 0, 0], [0, 2, 0, 0, 0],
[7, 3, 0, 0, 0], [0, 4, 8, 0, 0]])
non_null_match, non_null_gold, non_null_predictions = (
nq_long_utils.compute_match_stats(predictions=predictions,
labels=labels))
with tf.Session("") as sess:
tf_non_null_match, tf_non_null_gold, tf_non_null_predictions = (
sess.run(
[non_null_match, non_null_gold, non_null_predictions]))
self.assertAllClose(tf_non_null_match, [1.0, 0.0, 1.0, 0.0])
self.assertAllClose(tf_non_null_gold, [1.0, 0.0, 1.0, 1.0])
self.assertAllClose(tf_non_null_predictions,
[1.0, 1.0, 1.0, 1.0])
def model_function(features, labels, mode, params, embeddings):
"""A model function satisfying the tf.estimator API.
Args:
features: Dictionary of feature tensors with keys:
- question_tok: <string> [batch_size, max_question_len]
- context_tok: <string> [batch_size, max_num_context, max_context_len]
- question_tok_len: <int32> [batch_size]
- num_context: <int32> [batch_size]
- context_tok_len: <int32> [batch_size]
- question_tok_wid: <int32> [batch_size, max_question_len]
- context_tok_wid: <int32> [batch_size, max_num_context,
max_context_len]
- long_answer_indices: <int32> [batch_size]
labels: <int32> [batch_size] for answer index (-1 = NULL).
mode: One of the keys from tf.estimator.ModeKeys.
params: Dictionary of hyperparameters.
embeddings: An embedding_utils.PretrainedWordEmbeddings object.
Returns:
estimator_spec: A tf.estimator.EstimatorSpec object.
"""
del params # Unused.
if mode == tf.estimator.ModeKeys.PREDICT:
# Add a dummy batch dimension if we are exporting the predictor.
features = {k: tf.expand_dims(v, 0) for k, v in features.items()}
embedding_weights, embedding_scaffold = embeddings.get_params(
trainable=False)
# Features.
question_tok_len = features["question_tok_len"]
question_tok_wid = features["question_tok_wid"]
context_tok_wid = features["context_tok_wid"]
num_context = features["num_context"]
context_tok_len = features["context_tok_len"]
# Truncate the contexts and labels to a certain maximum length.
context_tok_wid, num_context, context_tok_len = (
nq_long_utils.truncate_contexts(context_token_ids=context_tok_wid,
num_contexts=num_context,
context_len=context_tok_len,
max_contexts=FLAGS.max_contexts,
max_context_len=FLAGS.max_context_len))
non_null_context_scores = nq_long_decatt_model.build_model(
question_tok_wid=question_tok_wid,
question_lens=question_tok_len,
context_tok_wid=context_tok_wid,
context_lens=context_tok_len,
embedding_weights=embedding_weights,
mode=mode)
# Mask out contexts that are padding.
num_context_mask = tf.log(
tf.sequence_mask(num_context,
tensor_utils.shape(non_null_context_scores, 1),
dtype=tf.float32))
non_null_context_scores += num_context_mask
# <float> [batch_size, 1]
null_score = tf.zeros([tf.shape(question_tok_wid)[0], 1])
# Offset everything by 1 to account for null context.
# [batch_size, 1 + max_contexts]
context_scores = tf.concat([null_score, non_null_context_scores], 1)
if mode != tf.estimator.ModeKeys.PREDICT:
labels = nq_long_utils.truncate_labels(labels, FLAGS.max_contexts)
# In the data, NULL is given index -1 but this is not compatible with
# softmax so shift by 1.
labels = labels + 1
# Reweight null examples.
weights = nq_long_utils.compute_null_weights(labels, FLAGS.null_weight)
# When computing the loss we take only the first label.
loss_labels = labels[:, 0]
# []
loss = tf.losses.sparse_softmax_cross_entropy(labels=loss_labels,
logits=context_scores,
weights=weights)
optimizer = tf.train.AdagradOptimizer(
learning_rate=FLAGS.learning_rate)
train_op = optimizer.minimize(loss=loss,
global_step=tf.train.get_global_step())
# <int32> [batch_size]
eval_predictions = tf.to_int32(tf.argmax(context_scores, 1))
non_null_match, non_null_gold, non_null_predictions = (
nq_long_utils.compute_match_stats(eval_predictions, labels))
precision, precision_op = (tf.metrics.mean(
non_null_match, weights=non_null_predictions))
recall, recall_op = (tf.metrics.mean(non_null_match,
weights=non_null_gold))
f1, f1_op = (nq_long_utils.f1_metric(precision=precision,
precision_op=precision_op,
recall=recall,
recall_op=recall_op))
# Bogus metric until we figure out how to connect Ming Wei's eval code.
eval_metric_ops = {
"precision": (precision, precision_op),
"recall": (recall, recall_op),
"f1": (f1, f1_op)
}
else:
loss = None
train_op = None
eval_metric_ops = {}
# In the export, we never predict NULL since the eval metric will compute the
# best possible F1.
export_long_answer_idx = tf.to_int32(tf.argmax(non_null_context_scores, 1))
export_long_answer_score = tf.reduce_max(non_null_context_scores, 1)
predictions = dict(idx=export_long_answer_idx,
score=export_long_answer_score)
if mode == tf.estimator.ModeKeys.PREDICT:
# Remove the dummy batch dimension if we are exporting the predictor.
predictions = {k: tf.squeeze(v, 0) for k, v in predictions.items()}
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
predictions=predictions,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
scaffold=embedding_scaffold)
return estimator_spec