|
|
| import os
|
| from dataclasses import dataclass, field
|
| from pathlib import Path
|
|
|
| from torch.utils.data import DataLoader
|
|
|
| from litgpt.data import DataModule
|
| from litgpt.tokenizer import Tokenizer
|
|
|
|
|
| @dataclass
|
| class LitData(DataModule):
|
| """Loads data using LitData's StreamingDataset given a path to a folder of preprocessed data (chunks)."""
|
|
|
| data_path: str | Path = Path("data/")
|
| """The path to the data directory containing the preprocessed chunks for the streaming dataset
|
| The path can also be a remote path (e.g., s3://). See also ``split_names`` if this path contains subfolders
|
| for training- and validation splits."""
|
| split_names: tuple[str, str] | None = None
|
| """Optional tuple for names of subfolders for training and validation under ``data_path``. If not provided,
|
| all data under data_path will be used for training, and the validation dataloader will be identical to the
|
| train dataloader."""
|
| seed: int = 42
|
| """The random seed for shuffling the dataset."""
|
| num_workers: int = 8
|
| """How many DataLoader processes to use for loading."""
|
|
|
| batch_size: int = field(init=False, repr=False, default=1)
|
| seq_length: int = field(init=False, repr=False, default=2048)
|
|
|
| def __post_init__(self) -> None:
|
| super().__init__()
|
| if self.split_names is not None and len(self.split_names) != 2:
|
| raise ValueError("If provided `split_names` must be a tuple of two strings, for example: ('train', 'val').")
|
|
|
| def connect(
|
| self, tokenizer: Tokenizer | None = None, batch_size: int = 1, max_seq_length: int | None = None
|
| ) -> None:
|
| self.batch_size = batch_size
|
| self.seq_length = max_seq_length + 1
|
|
|
| def train_dataloader(self) -> DataLoader:
|
| input_dir = os.path.join(self.data_path, self.split_names[0]) if self.split_names else str(self.data_path)
|
| return self._dataloader(input_dir=input_dir, train=True)
|
|
|
| def val_dataloader(self) -> DataLoader:
|
| input_dir = os.path.join(self.data_path, self.split_names[1]) if self.split_names else str(self.data_path)
|
| return self._dataloader(input_dir=input_dir, train=False)
|
|
|
| def _dataloader(self, input_dir: str, train: bool):
|
| from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader
|
|
|
| dataset = StreamingDataset(
|
| input_dir=input_dir,
|
| item_loader=TokensLoader(block_size=self.seq_length),
|
| shuffle=train,
|
| seed=self.seed,
|
| )
|
| dataloader = StreamingDataLoader(
|
| dataset, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_workers, drop_last=True
|
| )
|
| return dataloader
|
|
|
