File size: 2,894 Bytes
236083b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
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 # Increase by one because we need the next token as well
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
|