code: def input_fn_builder(
      input_files,
      batch_size,
      max_seq_length,
      max_predictions_per_seq,
      is_training,
      input_context=None,
      num_cpu_threads=4,
      num_eval_steps=1,
  ):
      """Creates an `input_fn` closure to be passed to TPUEstimator."""
  
      def input_fn(params, input_context=None):
          """The actual input function."""
          # For yarn mode only
          # import os
          # import subprocess
          # os.environ["CLASSPATH"] = subprocess.check_output(["hadoop", "classpath", "--glob"]).decode()
  
          # 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)
              )
              if input_context:
                  tf.logging.info(
                      "Sharding the dataset: input_pipeline_id=%d num_input_pipelines=%d"
                      % (
                          input_context.input_pipeline_id,
                          input_context.num_input_pipelines,
                      )
                  )
                  d = d.shard(
                      input_context.num_input_pipelines,
                      input_context.input_pipeline_id,
                  )
              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.data.experimental.parallel_interleave(
                      tf.data.TFRecordDataset,
                      sloppy=is_training,
                      cycle_length=cycle_length,
                  )
              )
              d = d.shuffle(buffer_size=1000)
              d = d.repeat()
          else:
              d = tf.data.TFRecordDataset(input_files)
              d = d.take(batch_size * num_eval_steps)
              # 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.data.experimental.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
  
file_path: BigDL/python/orca/example/learn/tf/bert/run_pretraining.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import absl
  import sys
  import modeling
  import optimization
  import tensorflow.compat.v1
  from bigdl.orca.learn.tf.tf_estimator import TFEstimator
label_desc: Use image node_id 5 for calling a global function with example usage: input_fn_builder(input_files, batch_size, max_seq_length, max_predictions_per_seq, is_training, input_context, num_cpu_threads, num_eval_steps) and returns: input_fn, d