def load_dataset(
    patterns: Union[str, Sequence[str]],
    decode_fn: Callable[[bytes], Dict[str, tf.Tensor]],
    preprocess_fn: Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]],
    collater_fn: Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]],
    is_training: bool,
    per_device_batch_size: int,
    local_device_count: int,
    host_count: int,
    host_id: int,
    seed: int = 0,
    pad_eval: bool = False,
) -> tf.data.Dataset:
  """Create TF dataset from one or more file patterns.

  Args:
    patterns: one or more file pattern strings from which to load data.
    decode_fn: decodes bytes into samples.
    preprocess_fn: performs preprocessing on examples.
    collater_fn: performs preprocessing on the whole batch.
    is_training: if true, shuffles.
    per_device_batch_size: total batch size per host.
    local_device_count: number of devices, used for batching.
    host_count: number of hosts, used for sharding.
    host_id: ID of current host, used for sharding.
    seed: random seed for non-deterministic operations.
    pad_eval: if true, pad rather than discard last batch for evaluation.

  Returns:
    tf.data.Dataset for training.
  """

  dataset = tf.data.Dataset.list_files(patterns, seed=seed)
  num_files = len(dataset)
  if num_files // host_count >= 10:
    shard_first = True
  else:
    shard_first = False

  if is_training:
    if shard_first:
      dataset = dataset.shard(num_shards=host_count, index=host_id)
    dataset = dataset.shuffle(buffer_size=num_files, seed=seed)
    dataset = dataset.repeat()
    dataset = dataset.interleave(
        tf.data.TFRecordDataset,
        cycle_length=num_files,
        num_parallel_calls=tf.data.AUTOTUNE,
        deterministic=True,
    )
    if not shard_first:
      dataset = dataset.shard(num_shards=host_count, index=host_id)
    dataset = dataset.shuffle(buffer_size=5000, seed=seed)
  else:
    dataset = tf.data.TFRecordDataset(dataset)

  # Decode tf example
  dataset = dataset.map(
      decode_fn,
      num_parallel_calls=tf.data.AUTOTUNE,
      deterministic=True,
  )
  # Flatten results to account for multiple samples per example
  dataset = dataset.unbatch()

  # We may not want to discard remainders or repeat for evaluation. Instead we
  # pad. If we shard first, one host might receive fewer batches than others, so
  # we first pad a full-size batch, then shard.
  if not is_training:
    if pad_eval:
      total_bsz = host_count * local_device_count * per_device_batch_size
      dataset = dataset.batch(
          batch_size=total_bsz,
          drop_remainder=False,
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=True,
      )
      dataset = dataset.map(
          functools.partial(pad_fn, bsz=total_bsz),
          num_parallel_calls=tf.data.AUTOTUNE,
          deterministic=True,
      )
      dataset = dataset.unbatch()
    dataset = dataset.shard(num_shards=host_count, index=host_id)

  # Preprocess individual samples
  dataset = dataset.map(
      preprocess_fn,
      num_parallel_calls=tf.data.AUTOTUNE,
      deterministic=True,
  )

  # Generate device batches
  dataset = dataset.batch(
      batch_size=per_device_batch_size,
      drop_remainder=True,
      num_parallel_calls=tf.data.AUTOTUNE,
      deterministic=True,
  )

  dataset = dataset.map(
      collater_fn,
      num_parallel_calls=tf.data.AUTOTUNE,
      deterministic=True,
  )

  # Host batch contains one device batch per local device
  dataset = dataset.batch(
      batch_size=local_device_count,
      drop_remainder=True,
      num_parallel_calls=tf.data.AUTOTUNE,
      deterministic=True,
  )

  dataset = dataset.prefetch(tf.data.AUTOTUNE)
  return dataset
Create TF dataset from one or more file patterns.

  Args:
    patterns: one or more file pattern strings from which to load data.
    decode_fn: decodes bytes into samples.
    preprocess_fn: performs preprocessing on examples.
    collater_fn: performs preprocessing on the whole batch.
    is_training: if true, shuffles.
    per_device_batch_size: total batch size per host.
    local_device_count: number of devices, used for batching.
    host_count: number of hosts, used for sharding.
    host_id: ID of current host, used for sharding.
    seed: random seed for non-deterministic operations.
    pad_eval: if true, pad rather than discard last batch for evaluation.

  Returns:
    tf.data.Dataset for training.
  
language/mentionmemory/utils/data_utils.py