File size: 5,518 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 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Implementation derived from https://github.com/tloen/alpaca-lora"""
import os
from dataclasses import dataclass, field
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 LIMA(DataModule):
"""LIMA data module for supervised finetuning."""
mask_prompt: bool = False
"""Whether to mask the prompt section from the label (with ``ignore_index``)."""
val_split_fraction: float = 0.1
"""The fraction of the dataset to use for the validation dataset. The rest is used for training."""
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."""
include_multiturn_conversations: bool = False
"""Whether to include multi-turn conversations in the dataset."""
repo_id: str = "GAIR/lima"
"""The Hugging Face dataset repository ID from where to download the data."""
access_token: str | None = field(repr=False, default=os.getenv("HF_TOKEN"))
"""The Hugging Face API token to use for authentication. Can also be set through the
`HF_TOKEN` environment variable."""
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):
super().__init__()
if self.access_token is None:
raise ValueError(
"LIMA requires authentication, please set the `HF_TOKEN=your_token` environment"
" variable or pass --access_token=your_token. You can find your token by visiting"
" https://huggingface.co/settings/tokens"
)
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:
from datasets import load_dataset
load_dataset(self.repo_id, token=self.access_token)
def setup(self, stage: str = "") -> None:
from datasets import load_dataset
dataset = load_dataset(self.repo_id, token=self.access_token)
data = format_dataset(dataset["train"], self.include_multiturn_conversations)
# 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),
)
train_data, test_data = list(train_data), list(test_data)
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 format_dataset(dataset_partition: dict, include_multi_turn_conversations: bool) -> list[dict]:
formatted_ds = []
for entry in dataset_partition:
convo = entry["conversations"]
if include_multi_turn_conversations:
for i in range(0, len(convo) - 1, 2):
formatted_ds.append({"instruction": convo[i], "input": "", "output": convo[i + 1]})
else:
formatted_ds.append({"instruction": convo[0], "input": "", "output": convo[1]})
return formatted_ds
|