code: def bert_input_fn(
      rdd,
      max_seq_length,
      batch_size,
      features={"input_ids", "input_mask", "token_type_ids"},
      extra_features=None,
      labels=None,
      label_size=None,
  ):
      """
      Takes an RDD to create the input function for BERT related TFEstimators.
      For training and evaluation, each element in rdd should be a tuple:
      (dict of features, a single label or dict of labels)
      Note that currently only integer or integer array labels are supported.
      For prediction, each element in rdd should be a dict of features.
  
      Features in each RDD element should contain "input_ids", "input_mask" and "token_type_ids",
      each of shape max_seq_length.
      If you have other extra features in your dict of features, you need to explicitly specify
      the argument `extra_features`, which is supposed to be the dict with feature name as key
      and tuple of (dtype, shape) as its value.
      """
      import tensorflow as tf
  
      invalidInputError(
          features.issubset(
              {"input_ids", "input_mask", "token_type_ids"}
          ),
          "features should be subset of {input_ids, input_mask, token_type_ids}",
      )
      features_dict = {}
      for feature in features:
          features_dict[feature] = (tf.int32, [max_seq_length])
      if extra_features is not None:
          invalidInputError(
              isinstance(extra_features, dict),
              "extra_features should be a dictionary",
          )
          for k, v in extra_features.items():
              invalidInputError(
                  isinstance(
                      k, six.string_types, "expect k is string type"
                  )
              )
              invalidInputError(
                  isinstance(v, tuple), "expect v is tuple"
              )
              features_dict[k] = v
      if label_size is None:
          label_size = []
      else:
          label_size = [label_size]
      if labels is None:
          res_labels = (tf.int32, label_size)
      elif isinstance(labels, list) or isinstance(labels, set):
          labels = set(labels)
          if len(labels) == 1:
              res_labels = (tf.int32, label_size)
          else:
              res_labels = {}
              for label in labels:
                  res_labels[label] = (tf.int32, label_size)
      else:
          invalidInputError(
              False,
              "Wrong labels. labels should be a set of label names"
              " if you have multiple labels",
          )
  
      def input_fn(mode):
          if mode == tf.estimator.ModeKeys.TRAIN:
              return TFDataset.from_rdd(
                  rdd,
                  features=features_dict,
                  labels=res_labels,
                  batch_size=batch_size,
              )
          elif mode == tf.estimator.ModeKeys.EVAL:
              return TFDataset.from_rdd(
                  rdd,
                  features=features_dict,
                  labels=res_labels,
                  batch_per_thread=batch_size // rdd.getNumPartitions(),
              )
          else:
              return TFDataset.from_rdd(
                  rdd,
                  features=features_dict,
                  batch_per_thread=batch_size // rdd.getNumPartitions(),
              )
  
      return input_fn
  
file_path: BigDL/python/orca/src/bigdl/orca/tfpark/text/estimator/bert_base.py
imports: from bigdl.orca.tfpark.estimator import *
  from bert import modeling
  from bigdl.dllib.utils.log4Error import *
label_desc: Use image node_id 2 for calling a global function with example usage: bert_input_fn(rdd, max_seq_length, batch_size, features, extra_features, labels, label_size) and returns: input_fn, TFDataset, TFDataset, TFDataset