composer.datasets.lm_datasets#

Generic dataset class for self-supervised training of autoregressive and masked language models.

Hparams

These classes are used with yahp for YAML-based configuration.

LMDatasetHparams

Defines a generic dataset class for self-supervised training of autoregressive and masked language models.
class composer.datasets.lm_datasets.LMDatasetHparams(is_train=True, drop_last=True, shuffle=True, datadir=<factory>, split=None, tokenizer_name=None, use_masked_lm=None, num_tokens=0, mlm_probability=0.15, seed=5, subsample_ratio=1.0, train_sequence_length=1024, val_sequence_length=1024)[source]#

Bases: composer.datasets.hparams.DatasetHparams

Defines a generic dataset class for self-supervised training of autoregressive and masked language models.

Parameters

  • datadir (list) โ€“ The path to the data directory.

  • is_train (bool) โ€“ Whether to load the training data or validation data. Default: True.

  • drop_last (bool) โ€“ If the number of samples is not divisible by the batch size, whether to drop the last batch or pad the last batch with zeros. Default: True.

  • shuffle (bool) โ€“ Whether to shuffle the dataset. Default: True.

  • datadir โ€“ List containing the string of the path to the HuggingFace Datasets directory.

  • split (str) โ€“ Whether to use 'train', 'test', or 'validation' split.

  • tokenizer_name (str) โ€“ The name of the HuggingFace tokenizer to preprocess text with. See HuggingFace documentation.

  • use_masked_lm (bool) โ€“ Whether the dataset should be encoded with masked language modeling or not.

  • num_tokens (int, optional) โ€“ Number of tokens to train on. 0 will train on all tokens in the dataset. Default: 0.

  • mlm_probability (float, optional) โ€“ If using masked language modeling, the probability with which tokens will be masked. Default: 0.15.

  • seed (int, optional) โ€“ Random seed for generating train and validation splits. Default: 5.

  • subsample_ratio (float, optional) โ€“ Proportion of the dataset to use. Default: 1.0.

  • train_sequence_length (int, optional) โ€“ Sequence length for training dataset. Default: 1024.

  • val_sequence_length (int, optional) โ€“ Sequence length for validation dataset. Default: 1024.

initialize_object(batch_size, dataloader_hparams)[source]#

Creates a DataLoader or DataSpec for this dataset.

Parameters

  • batch_size (int) โ€“ The size of the batch the dataloader should yield. This batch size is device-specific and already incorporates the world size.

  • dataloader_hparams (DataLoaderHparams) โ€“ The dataset-independent hparams for the dataloader.

Returns

DataLoader or DataSpec โ€“ The DataLoader, or if the dataloader yields batches of custom types, a DataSpec.

validate()[source]#

Validate that the hparams are of the correct types. Recurses through sub-hparams.

Raises

TypeError โ€“ Raises a TypeError if any fields are an incorrect type.