wehub-resource-sync
107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
# 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
|