code: def create_dataset_from_tf_record_files(
      input_file_pattern, pre_batch_size, batch_size, is_training=True
  ):
      """Creates dataset from (tf)records files for training/evaluation."""
  
      files = tf.data.Dataset.list_files(
          input_file_pattern, shuffle=is_training
      )
  
      def make_dataset(files_dataset, shard_index):
          """Returns dataset for sharded tf record files."""
          if pre_batch_size != batch_size:
              raise ValueError(
                  "Pre-batch ({}) size is not equal to batch "
                  "size ({})".format(pre_batch_size, batch_size)
              )
          files_dataset = files_dataset.shard(NUM_SHARDS, shard_index)
          dataset = files_dataset.interleave(tf.data.TFRecordDataset)
          decode_fn = functools.partial(
              data_pipeline.DatasetManager.deserialize,
              batch_size=pre_batch_size,
              is_training=is_training,
          )
          dataset = dataset.map(
              decode_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          return dataset
  
      dataset = tf.data.Dataset.range(NUM_SHARDS)
      map_fn = functools.partial(make_dataset, files)
      dataset = dataset.interleave(
          map_fn,
          cycle_length=NUM_SHARDS,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
      return dataset
  
file_path: autodist/examples/benchmark/utils/recommendation/ncf_input_pipeline.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import tensorflow.compat.v2
  from utils.recommendation import constants
  from utils.recommendation import movielens
  from utils.recommendation import data_pipeline
label_desc: Use image node_id 1 for calling a global function with example usage: create_dataset_from_tf_record_files(input_file_pattern, pre_batch_size, batch_size, is_training) and returns: dataset, dataset