def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
schema_guided_dst = SchemaGuidedDST(bert_config,
use_one_hot_embeddings)
outputs = schema_guided_dst.define_model(features, is_training)
if is_training:
total_loss = schema_guided_dst.define_loss(features, outputs)
else:
total_loss = tf.constant(0.0)
tvars = tf.trainable_variables()
scaffold_fn = None
if init_checkpoint:
assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
if use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(init_checkpoint,
assignment_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
output_spec = None
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = optimization.create_optimizer(total_loss, learning_rate,
num_train_steps,
num_warmup_steps, use_tpu)
global_step = tf.train.get_or_create_global_step()
logged_tensors = {
"global_step": global_step,
"total_loss": total_loss,
}
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn,
training_hooks=[
tf.train.LoggingTensorHook(logged_tensors, every_n_iter=5)
])
elif mode == tf.estimator.ModeKeys.EVAL:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, loss=total_loss, scaffold_fn=scaffold_fn)
else: # mode == tf.estimator.ModeKeys.PREDICT
predictions = schema_guided_dst.define_predictions(
features, outputs)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)
return output_spec
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
input_ids = features["input_ids"]
input_mask = features["input_mask"]
input_type_ids = features["input_type_ids"]
model = modeling.BertModel(
config=bert_config,
is_training=False,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=input_type_ids,
use_one_hot_embeddings=use_one_hot_embeddings)
if mode != tf.estimator.ModeKeys.PREDICT:
raise ValueError("Only PREDICT modes are supported: %s" % (mode))
tvars = tf.trainable_variables()
scaffold_fn = None
assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
if use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
all_layers = model.get_all_encoder_layers()
predictions = {}
# Add input features containing identity of the input sequence.
for feature in [
"embedding_tensor_name", "service_id", "intent_or_slot_id",
"value_id"
]:
predictions[feature] = features[feature]
# Use the embedding obtained from the final layer.
predictions["final_layer"] = all_layers[-1]
output_spec = tf.contrib.tpu.TPUEstimatorSpec(mode=mode,
predictions=predictions,
scaffold_fn=scaffold_fn)
return output_spec
