File size: 4,445 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import os
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from torch.utils.data import DataLoader
from litgpt.data import DataModule
from litgpt.tokenizer import Tokenizer
@dataclass
class OpenWebText(DataModule):
"""The OpenWebText data module for pretraining."""
data_path: str | Path = Path("data/openwebtext")
"""The path to the data directory, containing two folders 'train' and 'val'
which are the output of the preprocessing step. The path can also be a remote path (e.g., s3://)."""
val_split_fraction: float = 0.0005
"""The fraction of data that should be put aside for validation."""
seed: int = 42
"""The seed to use for shuffling the training data."""
num_workers: int = 8
"""The number of workers to use for the dataloaders."""
tokenizer: Tokenizer | None = field(default=None, repr=False, init=False)
batch_size: int = field(default=1, repr=False, init=False)
seq_length: int = field(default=2048, repr=False, init=False)
def __post_init__(self) -> None:
super().__init__()
# Could be a remote path (s3://) or a local path
self.data_path_train = str(self.data_path).rstrip("/") + "/train"
self.data_path_val = str(self.data_path).rstrip("/") + "/val"
def connect(
self, tokenizer: Tokenizer | None = None, batch_size: int = 1, max_seq_length: int | None = 2048
) -> None:
self.tokenizer = tokenizer
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:
from datasets import Dataset, load_dataset
from litdata import optimize
if str(self.data_path).startswith("s3://"):
print(f"The OpenWebText data path points to an S3 location: {self.data_path}. Skipping preprocessing.")
return
if Path(self.data_path_train).is_dir() and Path(self.data_path_val).is_dir():
print(f"Found OpenWebText train and val dir: {self.data_path}. Skipping preprocessing.")
return
dataset = load_dataset("openwebtext", num_proc=(os.cpu_count() // 2), trust_remote_code=True)
# Split the data in training and validation
split_dataset = dataset["train"].train_test_split(
test_size=self.val_split_fraction, seed=self.seed, shuffle=True
)
split_dataset["val"] = split_dataset.pop("test") # rename the test split to val
def tokenize(data: Dataset, index: int):
yield self.tokenizer.encode(data[index]["text"], eos=True)
optimize(
fn=partial(tokenize, split_dataset["train"]),
inputs=list(range(len(split_dataset["train"]))),
output_dir=self.data_path_train,
num_workers=min(64, os.cpu_count() - 1),
chunk_bytes="200MB",
)
optimize(
fn=partial(tokenize, split_dataset["val"]),
inputs=list(range(len(split_dataset["val"]))),
output_dir=self.data_path_val,
num_workers=min(8, os.cpu_count() - 1),
chunk_bytes="200MB",
)
def train_dataloader(self) -> DataLoader:
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader
train_dataset = StreamingDataset(
input_dir=self.data_path_train,
item_loader=TokensLoader(block_size=self.seq_length),
shuffle=True,
)
train_dataloader = StreamingDataLoader(
train_dataset, 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.data_path_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
|