dataset

Functions

load(path, schema, collator, batch[, ...])

Build a DataLoader for the given dataset.

Classes

ChunkedSampler(dataset, batch, workers[, ...])

A chunked distributed sampler.

DataModule(training, validation, schema, ...)

A DataModule that wraps ParquetDatasets.

ParquetDataset(source[, memory])

A parquet-backed dataset featuring random access and caching.
class undertale.models.dataset.ParquetDataset(source: str, memory: int = 4096)

Bases: Dataset

A parquet-backed dataset featuring random access and caching.

Loads data from a parquet dataset in one or more shards on disk. Caches shard reads in memory to optimize for high-locality access. Cache size is controlled by the memory parameter.

If the source dataset is smaller than the cache limit, then this is essentially just a lazy-loaded dataset. If the dataset is larger than the cache limit, this is a locality-optimized cached dataset reader.

Parameters:

  • source – Path to a single parquet file or a directory of several parquet files.

  • memory – Cache size limit in megabytes. Controls the maximum amount of memory the shard cache will use.

validate(schema: Type[Dataset]) → None

Validate this dataset against a given schema.

Parameters:

schema – A schema class to validate the dataset against.

Raises:

SchemaError – If the dataset does not conform to schema.

class undertale.models.dataset.ChunkedSampler(dataset: Dataset, batch: int, workers: int, ranks: int | None = None, rank: int | None = None)

Bases: Sampler

A chunked distributed sampler.

The default DistributedSampler from pytorch uses a strided sampling approach that does not suit the locality constraints of the caching approach in ParquetDataset. This sampler guarantees per rank and per worker contiguity.

If not provided, this will attempt to discover rank and world size from torch.distributed state. If not running in a distributed environment a world size of 1 and rank index of 0 will be assumed - i.e., non-distributed training.

Note

This makes some assumptions about relatively stable, but undocumented pytorch internals - in particular that the worker distribution strategy in DataLoader is round-robin. If this changes in the future this code will break.

Parameters:

  • dataset – The dataset from which to sample.

  • batch – Batch size.

  • workers – Number of parallel dataset workers. By default, this will spawn no dataset workers and fetch data in the main process.

  • ranks – The number of distributed ranks.

  • rank – The index of this distributed rank.

undertale.models.dataset.load(path: str, schema: Type[Dataset], collator: Callable, batch: int, workers: int = 0, memory: int = 4096) → DataLoader

Build a DataLoader for the given dataset.

Parameters:

  • path – Path to the dataset to process.

  • schema – Expected dataset schema.

  • collator – Dataset collator.

  • batch – Batch size.

  • workers – Number of parallel dataset workers. By default, this will spawn no dataset workers and fetch data in the main process.

  • memory – Shard cache memory limit in megabytes.

Returns:

A DataLoader for the given dataset at path with the given parameters.

class undertale.models.dataset.DataModule(training: str, validation: str | None, schema: Type[Dataset], collator: Callable, batch: int, workers: int = 0, memory: int = 4096)

Bases: LightningDataModule

A DataModule that wraps ParquetDatasets.

Parameters:

  • training – Path to the training dataset.

  • validation – Path to the validation dataset. If None validation will be skipped.

  • schema – Expected dataset schema.

  • collator – Dataset collator.

  • batch – Batch size.

  • workers – Number of parallel dataset workers. By default, this will spawn no dataset workers and fetch data in the main process.

  • memory – Shard cache memory limit in megabytes.