def make_tf_data( cls, filename: type_utils.PathLike, buffer_size: Optional[int] = None, ) -> tf.data.Dataset: from riegeli.tensorflow.ops import riegeli_dataset_ops as riegeli_tf # pylint: disable=g-import-not-at-top return riegeli_tf.RiegeliDataset(filename, buffer_size=buffer_size)
def make_tf_data(cls, filename: type_utils.PathLike, buffer_size: tf.int64) -> tf.data.Dataset: """Returns TensorFlow Dataset comprising given record file.""" return riegeli_tf.RiegeliDataset(filename, buffer_size=buffer_size)
def dataset_fn(self, filenames, num_epochs=1, batch_size=None): repeat_dataset = riegeli_dataset_ops.RiegeliDataset(filenames).repeat( num_epochs) if batch_size: return repeat_dataset.batch(batch_size) return repeat_dataset