File size: 6,634 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 149 150 151 152 153 154 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import json
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import torch
from torch.utils.data import DataLoader, random_split
from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
from litgpt.prompts import PromptStyle
from litgpt.tokenizer import Tokenizer
@dataclass
class JSON(DataModule):
"""Loads JSON or JSONL data for supervised finetuning."""
json_path: Path
"""A path to a JSON file or a directory with `train.json` and `val.json` containing the data.
The file(s) should contain a list of samples (dicts). Each dict must have the keys 'instruction' and 'output',
and can optionally have a key 'input' (see Alpaca)."""
mask_prompt: bool = False
"""Whether to mask the prompt section from the label (with ``ignore_index``)."""
val_split_fraction: float | None = None
"""The fraction of the dataset to use for the validation dataset. The rest is used for training.
Only applies if you passed in a single file to `json_path`."""
prompt_style: str | PromptStyle = "alpaca"
"""The style to apply to instruction prompts. See `litgpt.prompts` for a list of available styles."""
ignore_index: int = -100
"""The index to use for elements to be ignored in the label."""
seed: int = 42
"""The random seed for creating the train/val splits and shuffling the dataset."""
num_workers: int = 4
"""How many DataLoader processes to use for 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)
train_dataset: SFTDataset | None = field(default=None, init=False, repr=False)
val_dataset: SFTDataset | None = field(default=None, init=False, repr=False)
def __post_init__(self):
super().__init__()
if self.json_path.is_file() and self.val_split_fraction is None:
self.val_split_fraction = 0.05
warnings.warn(
"The `json_path` points to a single file and `val_split_fraction` was not set. "
"Defaulting to `val_split_fraction=0.05`. Set `val_split_fraction` explicitly "
"to use a different split percentage.",
UserWarning,
stacklevel=2,
)
if self.json_path.is_dir() and self.val_split_fraction is not None:
raise ValueError(
"If `json_path` is a directory, it must contain 'train.json' and 'val.json' files and"
f" hence `val_split_fraction` should not be set. Got `{self.val_split_fraction=}`."
)
if not self.json_path.exists():
raise FileNotFoundError(
"The `json_path` must be a file or a directory containing 'train.json' and 'val.json' files,"
f" but '{self.json_path!s}' does not exist."
)
if isinstance(self.prompt_style, str):
self.prompt_style = PromptStyle.from_name(self.prompt_style)
def connect(
self, tokenizer: Tokenizer | None = None, batch_size: int = 1, max_seq_length: int | None = None
) -> None:
self.tokenizer = tokenizer
self.batch_size = batch_size
self.max_seq_length = -1 if max_seq_length is None else max_seq_length
def setup(self, stage: str = "") -> None:
train_data, test_data = self.get_splits()
self.train_dataset = SFTDataset(
data=train_data,
tokenizer=self.tokenizer,
prompt_style=self.prompt_style,
max_seq_length=self.max_seq_length,
mask_prompt=self.mask_prompt,
ignore_index=self.ignore_index,
)
self.test_dataset = SFTDataset(
data=test_data,
tokenizer=self.tokenizer,
prompt_style=self.prompt_style,
max_seq_length=self.max_seq_length,
mask_prompt=self.mask_prompt,
ignore_index=self.ignore_index,
)
def train_dataloader(self) -> DataLoader:
return DataLoader(
self.train_dataset,
batch_size=self.batch_size,
shuffle=True,
generator=torch.Generator().manual_seed(self.seed),
num_workers=self.num_workers,
collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
)
def val_dataloader(self) -> DataLoader:
return DataLoader(
self.test_dataset,
batch_size=self.batch_size,
shuffle=False,
num_workers=self.num_workers,
collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
)
def get_splits(self) -> tuple:
# A single file (gets split into train and test)
if self.json_path.is_file():
data = load_split(self.json_path)
# Partition the dataset into train and test
train_data, test_data = random_split(
data,
[1.0 - self.val_split_fraction, self.val_split_fraction],
generator=torch.Generator().manual_seed(self.seed),
)
return train_data, test_data
# A directory containing train.json and val.json
if (train_file := self.find_split("train")) and (val_file := self.find_split("val")):
train_data = load_split(train_file)
test_data = load_split(val_file)
return train_data, test_data
raise FileNotFoundError(
"The `json_path` must be a file or a directory containing 'train.json' and 'val.json' files."
)
def find_split(self, split_name: str) -> Path | None:
for suffix in (".json", ".jsonl"):
if (file := self.json_path / f"{split_name}{suffix}").is_file():
return file
return None
def load_split(json_path: Path) -> Any:
if json_path.suffix == ".json":
with open(json_path, encoding="utf-8") as file:
return json.load(file)
if json_path.suffix == ".jsonl":
with open(json_path, encoding="utf-8") as file:
return [json.loads(line) for line in file]
else:
raise ValueError(f"Unsupported file format: {json_path.suffix}. Expected `.json` or `.jsonl`.")
|