File size: 6,329 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 144 145 146 147 148 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import glob
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 TextFiles(DataModule):
"""The TextFile data module used for pretraining.
Reads in text data from plaintext files contained in a data folder
and provides training and validation dataloaders that return batches of tokens.
Every sample is set to a fixed length.
"""
train_data_path: Path
"""The path to the data directory used for training that contains .txt files"""
val_data_path: Path | None = None
"""The path to the data directory used for validation that
contains .txt files. Splits off data for validation from the
training set if None."""
seed: int = 42
"""The seed to use for shuffling the dataset."""
num_workers: int = 4
"""The number of workers to use for data loading."""
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.out_path_train = self.train_data_path / "train"
if self.val_data_path is None:
self.out_path_val = self.train_data_path / "val"
else:
self.out_path_val = Path(self.val_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 optimize
from litdata.streaming import TokensLoader
train_files = sorted(glob.glob(str(self.train_data_path / "*.txt")))
assert len(train_files) > 0, f"No .txt files found in train data {train_files}"
if self.val_data_path is not None:
self.val_data_path = Path(self.val_data_path)
val_files = sorted(glob.glob(str(self.val_data_path / "*.txt")))
assert len(val_files) > 0, f"No .txt files found in validation data {val_files}"
# train/test split. let's use only shard 0 for test split, rest train
else:
assert len(train_files) > 1, f"Expected at least two .txt files in {train_files}"
val_files, *train_files = train_files
val_files = [val_files]
# It's ok to use almost all CPUs here because this runs in a single process
num_workers = max(1, self.num_workers)
use_workers = min(num_workers, len(train_files))
if not Path(self.out_path_train).is_dir():
validate_tokenizer(self.tokenizer)
optimize(
fn=partial(tokenize, tokenizer=self.tokenizer),
inputs=train_files,
output_dir=str(self.out_path_train),
num_workers=use_workers,
chunk_bytes="50MB",
item_loader=TokensLoader(block_size=self.max_seq_length),
)
else:
print(
f"\nWarning: Preprocessed training data found in {self.out_path_train}."
" For efficiency, reprocessing is skipped. If your text input has changed since"
" the last `litgpt pretrain` command, remove the preprocessed file(s) to trigger"
f" reprocessing: `rm -rf {self.out_path_train}`\n"
)
use_workers = min(num_workers, len(val_files))
if not Path(self.out_path_val).is_dir():
validate_tokenizer(self.tokenizer)
optimize(
fn=partial(tokenize, tokenizer=self.tokenizer),
inputs=val_files,
output_dir=str(self.out_path_val),
num_workers=use_workers,
chunk_bytes="50MB",
item_loader=TokensLoader(block_size=self.max_seq_length),
)
else:
print(
f"\nWarning: Preprocessed validation data found in {self.out_path_val}."
" For efficiency, reprocessing is skipped. If your text input has changed since"
" the last `litgpt pretrain` command, remove the preprocessed file(s) to trigger"
f" reprocessing: `rm -rf {self.out_path_val}`\n"
)
def train_dataloader(self) -> DataLoader:
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader
train_dataset = StreamingDataset(
input_dir=str(self.out_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.out_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 file:
text = file.read()
text = text.strip()
yield tokenizer.encode(text, bos=True, eos=False)
def validate_tokenizer(tokenizer: Tokenizer) -> None:
if tokenizer is None:
raise ValueError(
"Tokenizer is None. If you are using this data module via `litgpt pretrain`, "
"please provide a valid `--tokenizer_dir` path."
)
|