(bert_config, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings)
return model_fn
Returns `model_fn` closure for TPUEstimator.
Returns `model_fn` closure for TPUEstimator.
Returns model_fn closure for TPUEstimator .
def model_fn_builder(bert_config, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for TPUEstimator."""

    tf.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    masked_lm_positions = features["masked_lm_positions"]
    masked_lm_ids = features["masked_lm_ids"]
    masked_lm_weights = features["masked_lm_weights"]
    next_sentence_labels = features["next_sentence_labels"]

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

    model = modeling.BertModel(
        config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_mask,
        token_type_ids=segment_ids,
        use_one_hot_embeddings=use_one_hot_embeddings)

    (masked_lm_loss,
     masked_lm_example_loss, masked_lm_log_probs) = get_masked_lm_output(
         bert_config, model.get_sequence_output(), model.get_embedding_table(),
         masked_lm_positions, masked_lm_ids, masked_lm_weights)

    (next_sentence_loss, next_sentence_example_loss,
     next_sentence_log_probs) = get_next_sentence_output(
         bert_config, model.get_pooled_output(), next_sentence_labels)

    total_loss = masked_lm_loss + next_sentence_loss

    tvars = tf.trainable_variables()

    initialized_variable_names = {}
    scaffold_fn = None
    if init_checkpoint:
      (assignment_map, initialized_variable_names
      ) = 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)

    tf.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                      init_string)

    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)

      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op,
          scaffold_fn=scaffold_fn)
    elif mode == tf.estimator.ModeKeys.EVAL:

      def metric_fn(masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids,
                    masked_lm_weights, next_sentence_example_loss,
                    next_sentence_log_probs, next_sentence_labels):
        """Computes the loss and accuracy of the model."""
        masked_lm_log_probs = tf.reshape(masked_lm_log_probs,
                                         [-1, masked_lm_log_probs.shape[-1]])
        masked_lm_predictions = tf.argmax(
            masked_lm_log_probs, axis=-1, output_type=tf.int32)
        masked_lm_example_loss = tf.reshape(masked_lm_example_loss, [-1])
        masked_lm_ids = tf.reshape(masked_lm_ids, [-1])
        masked_lm_weights = tf.reshape(masked_lm_weights, [-1])
        masked_lm_accuracy = tf.metrics.accuracy(
            labels=masked_lm_ids,
            predictions=masked_lm_predictions,
            weights=masked_lm_weights)
        masked_lm_mean_loss = tf.metrics.mean(
            values=masked_lm_example_loss, weights=masked_lm_weights)

        next_sentence_log_probs = tf.reshape(
            next_sentence_log_probs, [-1, next_sentence_log_probs.shape[-1]])
        next_sentence_predictions = tf.argmax(
            next_sentence_log_probs, axis=-1, output_type=tf.int32)
        next_sentence_labels = tf.reshape(next_sentence_labels, [-1])
        next_sentence_accuracy = tf.metrics.accuracy(
            labels=next_sentence_labels, predictions=next_sentence_predictions)
        next_sentence_mean_loss = tf.metrics.mean(
            values=next_sentence_example_loss)

        return {
            "masked_lm_accuracy": masked_lm_accuracy,
            "masked_lm_loss": masked_lm_mean_loss,
            "next_sentence_accuracy": next_sentence_accuracy,
            "next_sentence_loss": next_sentence_mean_loss,
        }

      eval_metrics = (metric_fn, [
          masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids,
          masked_lm_weights, next_sentence_example_loss,
          next_sentence_log_probs, next_sentence_labels
      ])
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metrics=eval_metrics,
          scaffold_fn=scaffold_fn)
    else:
      raise ValueError("Only TRAIN and EVAL modes are supported: %s" % (mode))

    return output_spec

  return model_fn
def model_fn_builder ( bert_config , init_checkpoint , learning_rate , num_train_steps , num_warmup_steps , use_tpu , use_one_hot_embeddings ) : def model_fn ( features , labels , mode , params ) : # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" tf . logging . info ( "*** Features ***" ) for name in sorted ( features . keys ( ) ) : tf . logging . info ( "  name = %s, shape = %s" % ( name , features [ name ] . shape ) ) input_ids = features [ "input_ids" ] input_mask = features [ "input_mask" ] segment_ids = features [ "segment_ids" ] masked_lm_positions = features [ "masked_lm_positions" ] masked_lm_ids = features [ "masked_lm_ids" ] masked_lm_weights = features [ "masked_lm_weights" ] next_sentence_labels = features [ "next_sentence_labels" ] is_training = ( mode == tf . estimator . ModeKeys . TRAIN ) model = modeling . BertModel ( config = bert_config , is_training = is_training , input_ids = input_ids , input_mask = input_mask , token_type_ids = segment_ids , use_one_hot_embeddings = use_one_hot_embeddings ) ( masked_lm_loss , masked_lm_example_loss , masked_lm_log_probs ) = get_masked_lm_output ( bert_config , model . get_sequence_output ( ) , model . get_embedding_table ( ) , masked_lm_positions , masked_lm_ids , masked_lm_weights ) ( next_sentence_loss , next_sentence_example_loss , next_sentence_log_probs ) = get_next_sentence_output ( bert_config , model . get_pooled_output ( ) , next_sentence_labels ) total_loss = masked_lm_loss + next_sentence_loss tvars = tf . trainable_variables ( ) initialized_variable_names = { } scaffold_fn = None if init_checkpoint : ( assignment_map , initialized_variable_names ) = 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 ) tf . logging . info ( "**** Trainable Variables ****" ) for var in tvars : init_string = "" if var . name in initialized_variable_names : init_string = ", *INIT_FROM_CKPT*" tf . logging . info ( "  name = %s, shape = %s%s" , var . name , var . shape , init_string ) 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 ) output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , loss = total_loss , train_op = train_op , scaffold_fn = scaffold_fn ) elif mode == tf . estimator . ModeKeys . EVAL : def metric_fn ( masked_lm_example_loss , masked_lm_log_probs , masked_lm_ids , masked_lm_weights , next_sentence_example_loss , next_sentence_log_probs , next_sentence_labels ) : """Computes the loss and accuracy of the model.""" masked_lm_log_probs = tf . reshape ( masked_lm_log_probs , [ - 1 , masked_lm_log_probs . shape [ - 1 ] ] ) masked_lm_predictions = tf . argmax ( masked_lm_log_probs , axis = - 1 , output_type = tf . int32 ) masked_lm_example_loss = tf . reshape ( masked_lm_example_loss , [ - 1 ] ) masked_lm_ids = tf . reshape ( masked_lm_ids , [ - 1 ] ) masked_lm_weights = tf . reshape ( masked_lm_weights , [ - 1 ] ) masked_lm_accuracy = tf . metrics . accuracy ( labels = masked_lm_ids , predictions = masked_lm_predictions , weights = masked_lm_weights ) masked_lm_mean_loss = tf . metrics . mean ( values = masked_lm_example_loss , weights = masked_lm_weights ) next_sentence_log_probs = tf . reshape ( next_sentence_log_probs , [ - 1 , next_sentence_log_probs . shape [ - 1 ] ] ) next_sentence_predictions = tf . argmax ( next_sentence_log_probs , axis = - 1 , output_type = tf . int32 ) next_sentence_labels = tf . reshape ( next_sentence_labels , [ - 1 ] ) next_sentence_accuracy = tf . metrics . accuracy ( labels = next_sentence_labels , predictions = next_sentence_predictions ) next_sentence_mean_loss = tf . metrics . mean ( values = next_sentence_example_loss ) return { "masked_lm_accuracy" : masked_lm_accuracy , "masked_lm_loss" : masked_lm_mean_loss , "next_sentence_accuracy" : next_sentence_accuracy , "next_sentence_loss" : next_sentence_mean_loss , } eval_metrics = ( metric_fn , [ masked_lm_example_loss , masked_lm_log_probs , masked_lm_ids , masked_lm_weights , next_sentence_example_loss , next_sentence_log_probs , next_sentence_labels ] ) output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , loss = total_loss , eval_metrics = eval_metrics , scaffold_fn = scaffold_fn ) else : raise ValueError ( "Only TRAIN and EVAL modes are supported: %s" % ( mode ) ) return output_spec return model_fn