Create a dataset containing individual instruction for each shard.

  Each instruction is a dict:
  ```
  {
      "filepath": tf.Tensor(shape=(), dtype=tf.string),
      "mask_offset": tf.Tensor(shape=(), dtype=tf.int64),
      "mask": tf.Tensor(shape=(100,), dtype=tf.bool),
  }
  ```

  Args:
    instructions: `list[dict]`, the list of instruction dict

  Returns:
    instruction_ds: The dataset containing the instruction. The dataset size is
      the number of shard.

def _build_instruction_ds(instructions):
  """Create a dataset containing individual instruction for each shard.

  Each instruction is a dict:
  ```
  {
      "filepath": tf.Tensor(shape=(), dtype=tf.string),
      "mask_offset": tf.Tensor(shape=(), dtype=tf.int64),
      "mask": tf.Tensor(shape=(100,), dtype=tf.bool),
  }
  ```

  Args:
    instructions: `list[dict]`, the list of instruction dict

  Returns:
    instruction_ds: The dataset containing the instruction. The dataset size is
      the number of shard.
  """
  # Transpose the list[dict] into dict[list]
  tensor_inputs = {
      # offset_mask need to be converted to int64 explicitly
      k: np.array(vals, dtype=np.int64) if k == "mask_offset" else list(vals)
      for k, vals in utils.zip_dict(*instructions)
  }
  return tf.data.Dataset.from_tensor_slices(tensor_inputs)