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
  
      assert features.issubset(
          {"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:
          assert isinstance(
              extra_features, dict
          ), "extra_features should be a dictionary"
          for k, v in extra_features.items():
              assert isinstance(k, six.string_types)
              assert isinstance(v, 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:
          raise ValueError(
              "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
  
imports: from zoo.tfpark.estimator import *
  from bert import modeling
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