def build_dataset(instruction_dicts,
                  dataset_from_file_fn,
                  shuffle_files=False,
                  parallel_reads=64):
  """Constructs a `tf.data.Dataset` from TFRecord files.

  Args:
    instruction_dicts: `list` of {'filepath':, 'mask':, 'offset_mask':}
      containing the information about which files and which examples to use.
      The boolean mask will be repeated and zipped with the examples from
      filepath.
    dataset_from_file_fn: function returning a `tf.data.Dataset` given a
      filename.
    shuffle_files: `bool`, Whether to shuffle the input filenames.
    parallel_reads: `int`, how many files to read in parallel.

  Returns:
    `tf.data.Dataset`
  """

  # First case: All examples are taken (No value skipped)
  if _no_examples_skipped(instruction_dicts):
    # Only use the filenames as instruction
    instruction_ds = tf.data.Dataset.from_tensor_slices([
        d["filepath"] for d in instruction_dicts
    ])
    build_ds_from_instruction = dataset_from_file_fn
  # Second case: Use the instructions to read the examples
  else:
    instruction_ds = _build_instruction_ds(instruction_dicts)
    build_ds_from_instruction = functools.partial(
        _build_ds_from_instruction,
        ds_from_file_fn=dataset_from_file_fn,
    )

  # If shuffle is True, we shuffle the instructions/shards
  if shuffle_files:
    instruction_ds = instruction_ds.shuffle(len(instruction_dicts))

  # Use interleave to parallel read files and decode records
  ds = instruction_ds.interleave(
      build_ds_from_instruction,
      cycle_length=parallel_reads,
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return ds
Constructs a `tf.data.Dataset` from TFRecord files.

  Args:
    instruction_dicts: `list` of {'filepath':, 'mask':, 'offset_mask':}
      containing the information about which files and which examples to use.
      The boolean mask will be repeated and zipped with the examples from
      filepath.
    dataset_from_file_fn: function returning a `tf.data.Dataset` given a
      filename.
    shuffle_files: `bool`, Whether to shuffle the input filenames.
    parallel_reads: `int`, how many files to read in parallel.

  Returns:
    `tf.data.Dataset`