def as_dataset(self,
                 split,
                 shuffle_files=False,
                 read_config=None,
                 **unused_kwargs):
    """Returns a parsed and preprocessed `tf.data.Dataset`."""
    if tf.io.gfile.exists(os.path.join(self._data_dir, f'{split}.tfrecord')):
      data_filenames = os.path.join(self._data_dir, f'{split}.tfrecord')
    else:
      data_filenames = os.path.join(self._data_dir, f'{split}-*-of-*.tfrecord')
    assert tf.io.gfile.glob(data_filenames), ('No data files matched %s!' %
                                              data_filenames)
    shuffle_seed = None if read_config is None else read_config.shuffle_seed
    files = tf.data.Dataset.list_files(
        data_filenames, shuffle=shuffle_files, seed=shuffle_seed)

    dataset_type = functools.partial(
        tf.data.TFRecordDataset,
        buffer_size=TFRECORD_BUFFER_SIZE,
        num_parallel_reads=tf.data.experimental.AUTOTUNE)
    dataset = files.interleave(
        dataset_type,
        cycle_length=-1,
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
        deterministic=True)

    dataset = dataset.map(
        self._parse, num_parallel_calls=tf.data.experimental.AUTOTUNE)

    dataset = dataset.apply(
        tf.data.experimental.assert_cardinality(self.num_examples[split]))

    return dataset
Returns a parsed and preprocessed `tf.data.Dataset`.
language/gscan/xattn_model/dataset/gscan_dataset.py