code: def gen_lm1b_train_dataset(file_pattern, num_step):
      """
      Returns: The training dataset (tf.data.Dataset) that has been repeated
      and shuffled
      """
      file_names = []
      for file_name in glob.glob(file_pattern):
          file_names.append(file_name)
      if not file_names:
          raise ValueError
      # create dataset ops
      BUFFER_SIZE = 100000
  
      # TODO(Hao): have to use v1 APIs
      d = (
          tf.compat.v1.data.TextLineDataset(file_names)
          .map(lambda string: tf.strings.split([string]).values)
          .flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x))
          .window(num_step, 1, 1, True)
          .flat_map(lambda x: x.batch(num_step))
          .window(2, 1, 1, True)
          .flat_map(lambda x: x.batch(2))
          .shuffle(BUFFER_SIZE, reshuffle_each_iteration=True)
          .repeat()
      )
      return d
  
file_path: autodist/examples/lm1b/lm1b_train.py
imports: import glob
  import os
  import time
  import tensorflow
  from absl import app
  from absl import flags
  import language_model
  from autodist import AutoDist
  from autodist.strategy import PS, PSLoadBalancing, PartitionedPS, AllReduce, Parallax
label_desc: Use image node_id 1 for calling a global function with example usage: gen_lm1b_train_dataset(file_pattern, num_step) and returns: d