(input_files,
                     max_seq_length,
                     max_predictions_per_seq,
                     is_training,
                     num_cpu_threads=4)
return input_fn
Creates an `input_fn` closure to be passed to TPUEstimator.
Creates an `input_fn` closure to be passed to TPUEstimator.
Creates an input_fn closure to be passed to TPUEstimator .
def input_fn_builder(input_files,
                     max_seq_length,
                     max_predictions_per_seq,
                     is_training,
                     num_cpu_threads=4):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    name_to_features = {
        "input_ids":
            tf.FixedLenFeature([max_seq_length], tf.int64),
        "input_mask":
            tf.FixedLenFeature([max_seq_length], tf.int64),
        "segment_ids":
            tf.FixedLenFeature([max_seq_length], tf.int64),
        "masked_lm_positions":
            tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
        "masked_lm_ids":
            tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
        "masked_lm_weights":
            tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
        "next_sentence_labels":
            tf.FixedLenFeature([1], tf.int64),
    }

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
      d = d.repeat()
      d = d.shuffle(buffer_size=len(input_files))

      # `cycle_length` is the number of parallel files that get read.
      cycle_length = min(num_cpu_threads, len(input_files))

      # `sloppy` mode means that the interleaving is not exact. This adds
      # even more randomness to the training pipeline.
      d = d.apply(
          tf.contrib.data.parallel_interleave(
              tf.data.TFRecordDataset,
              sloppy=is_training,
              cycle_length=cycle_length))
      d = d.shuffle(buffer_size=100)
    else:
      d = tf.data.TFRecordDataset(input_files)
      # Since we evaluate for a fixed number of steps we don't want to encounter
      # out-of-range exceptions.
      d = d.repeat()

    # We must `drop_remainder` on training because the TPU requires fixed
    # size dimensions. For eval, we assume we are evaluating on the CPU or GPU
    # and we *don't* want to drop the remainder, otherwise we wont cover
    # every sample.
    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            num_parallel_batches=num_cpu_threads,
            drop_remainder=True))
    return d

  return input_fn
def input_fn_builder ( input_files , max_seq_length , max_predictions_per_seq , is_training , num_cpu_threads = 4 ) : def input_fn ( params ) : """The actual input function.""" batch_size = params [ "batch_size" ] name_to_features = { "input_ids" : tf . FixedLenFeature ( [ max_seq_length ] , tf . int64 ) , "input_mask" : tf . FixedLenFeature ( [ max_seq_length ] , tf . int64 ) , "segment_ids" : tf . FixedLenFeature ( [ max_seq_length ] , tf . int64 ) , "masked_lm_positions" : tf . FixedLenFeature ( [ max_predictions_per_seq ] , tf . int64 ) , "masked_lm_ids" : tf . FixedLenFeature ( [ max_predictions_per_seq ] , tf . int64 ) , "masked_lm_weights" : tf . FixedLenFeature ( [ max_predictions_per_seq ] , tf . float32 ) , "next_sentence_labels" : tf . FixedLenFeature ( [ 1 ] , tf . int64 ) , } # For training, we want a lot of parallel reading and shuffling. # For eval, we want no shuffling and parallel reading doesn't matter. if is_training : d = tf . data . Dataset . from_tensor_slices ( tf . constant ( input_files ) ) d = d . repeat ( ) d = d . shuffle ( buffer_size = len ( input_files ) ) # `cycle_length` is the number of parallel files that get read. cycle_length = min ( num_cpu_threads , len ( input_files ) ) # `sloppy` mode means that the interleaving is not exact. This adds # even more randomness to the training pipeline. d = d . apply ( tf . contrib . data . parallel_interleave ( tf . data . TFRecordDataset , sloppy = is_training , cycle_length = cycle_length ) ) d = d . shuffle ( buffer_size = 100 ) else : d = tf . data . TFRecordDataset ( input_files ) # Since we evaluate for a fixed number of steps we don't want to encounter # out-of-range exceptions. d = d . repeat ( ) # We must `drop_remainder` on training because the TPU requires fixed # size dimensions. For eval, we assume we are evaluating on the CPU or GPU # and we *don't* want to drop the remainder, otherwise we wont cover # every sample. d = d . apply ( tf . contrib . data . map_and_batch ( lambda record : _decode_record ( record , name_to_features ) , batch_size = batch_size , num_parallel_batches = num_cpu_threads , drop_remainder = True ) ) return d return input_fn