(bert_config, init_checkpoint, layer_indexes, 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, layer_indexes, 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."""

    unique_ids = features["unique_ids"]
    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,
     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)

    all_layers = model.get_all_encoder_layers()

    predictions = {
        "unique_id": unique_ids,
    }

    for (i, layer_index) in enumerate(layer_indexes):
      predictions["layer_output_%d" % i] = all_layers[layer_index]

    output_spec = tf.contrib.tpu.TPUEstimatorSpec(
        mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)
    return output_spec

  return model_fn
def model_fn_builder ( bert_config , init_checkpoint , layer_indexes , use_tpu , use_one_hot_embeddings ) : def model_fn ( features , labels , mode , params ) : # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" unique_ids = features [ "unique_ids" ] 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 , 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 ) all_layers = model . get_all_encoder_layers ( ) predictions = { "unique_id" : unique_ids , } for ( i , layer_index ) in enumerate ( layer_indexes ) : predictions [ "layer_output_%d" % i ] = all_layers [ layer_index ] output_spec = tf . contrib . tpu . TPUEstimatorSpec ( mode = mode , predictions = predictions , scaffold_fn = scaffold_fn ) return output_spec return model_fn