def _sequence_correct(labels, predictions):
"""Computes a per-example sequence accuracy."""
target_decode_steps = decode_utils.decode_steps_from_labels(
labels, trim_start_symbol=True)
predicted_decode_steps = decode_utils.decode_steps_from_predictions(
predictions)
decode_utils.assert_shapes_match(target_decode_steps,
predicted_decode_steps)
equal_tokens = decode_utils.compare_decode_steps(target_decode_steps,
predicted_decode_steps)
target_len = labels["target_len"] - 1
loss_mask = tf.sequence_mask(lengths=tf.to_int32(target_len),
maxlen=tf.to_int32(tf.shape(equal_tokens)[1]))
equal_tokens = tf.logical_or(equal_tokens, tf.logical_not(loss_mask))
all_equal = tf.cast(tf.reduce_all(equal_tokens, 1), tf.float32)
return all_equal
def model_fn(features, labels, mode, params=None):
"""Model function for use with tf.learn.Estimator."""
del params # unused. model_fn is batch-size agnostic.
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
pretrained_variable_names = None
scaffold_fn = None
bert_config = bert_utils.get_bert_config(
model_config.model_parameters.pretrained_bert_dir,
reinitialize_type_embeddings=model_config.model_parameters.
use_segment_ids)
input_embeddings = embeddings.get_input_embeddings(
model_config,
bert_config,
features,
is_training,
use_one_hot_embeddings=use_tpu)
source_len = tf.to_int32(features[constants.SOURCE_LEN_KEY])
output_embeddings_table = embeddings.get_output_vocab_embeddings_table(
model_config, output_vocab_filepath)
output_vocab_size = embeddings.get_output_vocab_size(
output_vocab_filepath)
clean_output_mask = None
if clean_output_vocab_path:
clean_output_mask_list = embeddings.get_clean_output_mask(
output_vocab_filepath, clean_output_vocab_path)
clean_output_mask = tf.convert_to_tensor(clean_output_mask_list)
# For inference, just compute the inference predictions and return.
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = transformer.infer(
model_config,
input_embeddings,
source_len,
output_vocab_size,
output_embeddings_table,
mode,
input_copy_mask=features[constants.COPIABLE_INPUT_KEY],
clean_output_mask=clean_output_mask,
beam_size=beam_size)
if use_tpu:
return tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
predictions=predictions,
scaffold_fn=scaffold_fn)
else:
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions)
with tpu_utils.rewire_summary_calls(use_tpu):
# Get training predictions.
train_decode_steps = decode_utils.decode_steps_from_labels(
labels, trim_end_symbol=True)
logits, predictions = transformer.train(
model_config,
input_embeddings,
source_len,
output_vocab_size,
output_embeddings_table,
train_decode_steps,
mode,
input_copy_mask=features[constants.COPIABLE_INPUT_KEY])
# Calculate loss.
weights = labels[constants.WEIGHT_KEY]
loss_decode_steps = decode_utils.decode_steps_from_labels(
labels, trim_start_symbol=True)
# Account for removed start symbol.
target_len = tf.to_int32(labels[constants.TARGET_LEN_KEY])
target_len -= 1
batch_loss = _compute_loss(logits, loss_decode_steps, target_len,
weights, output_vocab_size,
model_config)
if mode == tf.estimator.ModeKeys.TRAIN:
pretrained_variable_names, scaffold_fn = load_from_checkpoint.init_model_from_checkpoint(
model_config.model_parameters.pretrained_bert_dir,
use_tpu=use_tpu,
checkpoint_file="bert_model.ckpt",
reinitialize_type_embeddings=model_config.model_parameters.
use_segment_ids)
train_op = adam_weight_decay.build_train_op_with_pretraining(
batch_loss, model_config, pretrained_variable_names,
use_tpu)
if use_tpu:
return tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=batch_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
else:
return tf.estimator.EstimatorSpec(mode=mode,
loss=batch_loss,
train_op=train_op)
eval_metrics = metrics.create_metrics_ops(labels=labels,
predictions=predictions)
return tf.estimator.EstimatorSpec(mode=mode,
loss=batch_loss,
eval_metric_ops=eval_metrics)