code: def get_pretrain_dataset_fn(
      input_file_pattern,
      seq_length,
      max_predictions_per_seq,
      global_batch_size,
      num_replicas_in_sync,
  ):
      """Returns input dataset from input file string."""
  
      def _dataset_fn(ctx=None):
          """Returns tf.data.Dataset for distributed BERT pretraining."""
          input_patterns = input_file_pattern.split(",")
          batch_size = int(global_batch_size / num_replicas_in_sync)
          train_dataset = input_pipeline.create_pretrain_dataset(
              input_patterns,
              seq_length,
              max_predictions_per_seq,
              batch_size,
              is_training=True,
          )
          return train_dataset
  
      return _dataset_fn
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import yaml
  import os
  import sys
  import tensorflow
  from absl import app
  from absl import flags
  from absl import logging
  from utils.logs import logger
  from utils.misc import keras_utils
  from utils import bert_modeling
  from utils import bert_models
  from utils import common_flags
  from utils import input_pipeline
  from utils import bert_utils
  from autodist import AutoDist
  from autodist.strategy.all_reduce_strategy import AllReduce
  from autodist.strategy.ps_strategy import PS
  from autodist.strategy.ps_lb_strategy import PSLoadBalancing
  from autodist.strategy.parallax_strategy import Parallax
  from autodist.strategy.partitioned_ps_strategy import PartitionedPS
label_desc: Use image node_id 1 for calling a global function with example usage: get_pretrain_dataset_fn(input_file_pattern, seq_length, max_predictions_per_seq, global_batch_size, num_replicas_in_sync) and returns: _dataset_fn, train_dataset