File size: 5,790 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import glob
import json
import os
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from torch.utils.data import DataLoader
from tqdm import tqdm
from litgpt.data import DataModule
from litgpt.data.alpaca import download_if_missing
from litgpt.data.text_files import validate_tokenizer
from litgpt.tokenizer import Tokenizer
@dataclass
class TinyStories(DataModule):
"""The TinyStories data module: https://huggingface.co/datasets/roneneldan/TinyStories
Provides training and validation dataloaders that return batches of tokens. Every sample is set to a fixed length.
"""
data_path: Path = Path("data/tinystories")
"""The path to the data directory, containing two folders 'train' and 'val'
which are the output of the preprocessing step."""
seed: int = 42
"""The seed to use for shuffling the dataset."""
num_workers: int = 8
"""The number of workers to use for the dataloaders."""
tokenizer: Tokenizer | None = field(default=None, init=False, repr=False)
batch_size: int = field(default=1, init=False, repr=False)
max_seq_length: int = field(default=-1, init=False, repr=False)
def __post_init__(self) -> None:
super().__init__()
self.data_path_train = self.data_path / "train"
self.data_path_val = self.data_path / "val"
def connect(self, tokenizer: Tokenizer | None = None, batch_size: int = 1, max_seq_length: int = -1) -> None:
self.tokenizer = tokenizer
self.batch_size = batch_size
self.max_seq_length = max_seq_length + 1 # Increase by one because we need the next token as well
def prepare_data(self) -> None:
from litdata import TokensLoader, optimize
download(self.data_path)
files = sorted(glob.glob(str(self.data_path / "TinyStories_all_data" / "*.json")))
assert len(files) > 0, f"No json files found in {files}"
assert len(files) > 1, f"Expected at least two json files in {files}"
# train/test split. let's use only shard 0 for test split, rest train
val_file, *train_files = files
num_workers = os.cpu_count() - 1
if not Path(self.data_path_train).is_dir():
validate_tokenizer(self.tokenizer)
optimize(
fn=partial(tokenize, tokenizer=self.tokenizer),
inputs=train_files,
output_dir=str(self.data_path_train),
num_workers=num_workers,
chunk_bytes="200MB",
item_loader=TokensLoader(),
)
if not Path(self.data_path_val).is_dir():
validate_tokenizer(self.tokenizer)
optimize(
fn=partial(tokenize, tokenizer=self.tokenizer),
inputs=[val_file],
output_dir=str(self.data_path_val),
num_workers=1, # there's only 1 file
chunk_bytes="200MB",
item_loader=TokensLoader(),
)
def train_dataloader(self) -> DataLoader:
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader
train_dataset = StreamingDataset(
input_dir=str(self.data_path_train),
item_loader=TokensLoader(block_size=self.max_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=str(self.data_path_val),
item_loader=TokensLoader(block_size=self.max_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
def tokenize(filename: str, tokenizer: Tokenizer):
with open(filename, encoding="utf-8") as f:
data = json.load(f)
global_rank = int(os.environ["DATA_OPTIMIZER_GLOBAL_RANK"])
num_workers = int(os.environ["DATA_OPTIMIZER_NUM_WORKERS"])
local_rank = global_rank % num_workers
for example in tqdm(data, position=local_rank):
text = example["story"]
text = text.strip() # get rid of leading/trailing whitespace
tokens = tokenizer.encode(text, bos=True, eos=False) # encode the text, use BOS
yield tokens
_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories_all_data.tar.gz"
def download(data_dir: Path):
data_dir.mkdir(exist_ok=True, parents=True)
data_tar = data_dir / "TinyStories_all_data.tar.gz"
data_dir = data_dir / "TinyStories_all_data"
shard_filenames = sorted(glob.glob(str(data_dir / "*.json")))
if shard_filenames:
print(f"{data_dir} already exists, skipping unpacking...")
return
# download the TinyStories dataset, unless it's already downloaded
download_if_missing(data_tar, _URL, stream=True, mode="wb")
# unpack the tar.gz file into all the data shards (json files)
data_dir.mkdir(exist_ok=False)
tar_command = f"tar -xzf {data_tar} -C {data_dir}"
print(tar_command)
os.system(tar_command)
shard_filenames = sorted(glob.glob(str(data_dir / "*.json")))
print(f"Number of shards: {len(shard_filenames)}")
|