wehub-resource-sync
90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
|
|
from litgpt.data.alpaca import download_if_missing
|
|
from litgpt.prompts import PromptStyle
|
|
from litgpt.tokenizer import Tokenizer
|
|
|
|
_URL = "https://raw.githubusercontent.com/akoksal/LongForm/main/dataset"
|
|
|
|
|
|
@dataclass
|
|
class LongForm(DataModule):
|
|
"""LongForm data module for supervised finetuning."""
|
|
|
|
mask_prompt: bool = False
|
|
"""Whether to mask the prompt section from the label (with ``ignore_index``)."""
|
|
prompt_style: str | PromptStyle = "longform"
|
|
"""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 shuffling the dataset."""
|
|
num_workers: int = 4
|
|
"""How many DataLoader processes to use for loading."""
|
|
download_dir: Path = Path("./data/longform")
|
|
"""The directory in which the downloaded dataset gets saved."""
|
|
|
|
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)
|
|
test_dataset: SFTDataset | None = field(default=None, init=False, repr=False)
|
|
|
|
def __post_init__(self) -> None:
|
|
super().__init__()
|
|
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 prepare_data(self) -> None:
|
|
self.download_dir.mkdir(parents=True, exist_ok=True)
|
|
download_if_missing(self.download_dir / "train.json", f"{_URL}/train.json")
|
|
download_if_missing(self.download_dir / "val.json", f"{_URL}/val.json")
|
|
|
|
def train_dataloader(self):
|
|
return self._dataloader("train")
|
|
|
|
def val_dataloader(self):
|
|
return self._dataloader("val")
|
|
|
|
def _dataloader(self, split: str) -> DataLoader:
|
|
with open(self.download_dir / f"{split}.json", encoding="utf-8") as file:
|
|
data = json.load(file)
|
|
|
|
dataset = SFTDataset(
|
|
data=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,
|
|
transform=_transform,
|
|
)
|
|
return DataLoader(
|
|
dataset=dataset,
|
|
batch_size=self.batch_size,
|
|
shuffle=(split == "train"),
|
|
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 _transform(item: dict) -> dict:
|
|
item["instruction"] = item.pop("input")
|
|
return item
|