wehub-resource-sync
103 lines
4.2 KiB
Python
103 lines
4.2 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
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 TinyLlama(DataModule):
|
|
"""The TinyLlama data module is composed of a mix of SlimPajama and Starcoder data.
|
|
|
|
Provides training and validation streaming dataloaders that return batches of tokens.
|
|
"""
|
|
|
|
data_path: str | Path = Path("data/")
|
|
"""The path to the data directory, containing two folders 'slimpajama' and 'starcoder'
|
|
which are the output of the preprocessing step done in advance. See the `tutorial/pretrain_tinyllama.md`
|
|
for instructions. The path can also be a remote path (e.g., s3://)."""
|
|
seed: int = 42
|
|
"""The random seed for shuffling the dataset."""
|
|
num_workers: int = 8
|
|
"""How many DataLoader processes to use for loading."""
|
|
use_starcoder: bool = True
|
|
"""Toggle for using Starcoder data."""
|
|
|
|
batch_size: int = field(init=False, repr=False, default=1)
|
|
seq_length: int = field(init=False, repr=False, default=2048)
|
|
|
|
def __post_init__(self):
|
|
super().__init__()
|
|
# Could be a remote path (s3://) or a local path
|
|
self.slimpajama_train = str(self.data_path).rstrip("/") + "/slimpajama/train"
|
|
self.slimpajama_val = str(self.data_path).rstrip("/") + "/slimpajama/val"
|
|
self.required_paths = [self.slimpajama_train, self.slimpajama_val]
|
|
|
|
if self.use_starcoder:
|
|
self.starcoder_train = str(self.data_path).rstrip("/") + "/starcoder"
|
|
self.required_paths += [self.starcoder_train]
|
|
|
|
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 prepare_data(self) -> None:
|
|
for path in self.required_paths:
|
|
if not path.startswith("s3://") and not Path(path).is_dir():
|
|
raise FileNotFoundError(
|
|
"The data path for TinyLlama is expected to be the directory containing these subdirectories:"
|
|
f" `slimpajama/train`, `slimpajama/val`, `starcoder`. The directory {path} does not exist."
|
|
" Set it via `--data.data_path=...`"
|
|
)
|
|
|
|
def train_dataloader(self) -> DataLoader:
|
|
from litdata.streaming import CombinedStreamingDataset, StreamingDataLoader, StreamingDataset, TokensLoader
|
|
|
|
slim_train_data = StreamingDataset(
|
|
input_dir=self.slimpajama_train,
|
|
item_loader=TokensLoader(block_size=self.seq_length),
|
|
shuffle=True,
|
|
drop_last=True,
|
|
)
|
|
train_data = slim_train_data
|
|
|
|
if self.use_starcoder:
|
|
train_datasets = [
|
|
slim_train_data,
|
|
StreamingDataset(
|
|
input_dir=self.starcoder_train,
|
|
item_loader=TokensLoader(block_size=self.seq_length),
|
|
shuffle=True,
|
|
drop_last=True,
|
|
),
|
|
]
|
|
|
|
# Mix SlimPajama data and Starcoder data with these proportions:
|
|
weights = (0.693584, 0.306416)
|
|
train_data = CombinedStreamingDataset(
|
|
datasets=train_datasets, seed=self.seed, weights=weights, iterate_over_all=False
|
|
)
|
|
|
|
train_dataloader = StreamingDataLoader(
|
|
train_data, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_workers, drop_last=True
|
|
)
|
|
return train_dataloader
|
|
|
|
def val_dataloader(self) -> DataLoader:
|
|
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader
|
|
|
|
val_dataset = StreamingDataset(
|
|
input_dir=self.slimpajama_val,
|
|
item_loader=TokensLoader(block_size=self.seq_length),
|
|
shuffle=True,
|
|
)
|
|
val_dataloader = StreamingDataLoader(
|
|
val_dataset, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_workers, drop_last=True
|
|
)
|
|
return val_dataloader
|