diff --git a/litgpt/data/__init__.py b/litgpt/data/__init__.py index 66405091c9..2598fbede3 100644 --- a/litgpt/data/__init__.py +++ b/litgpt/data/__init__.py @@ -11,6 +11,7 @@ from litgpt.data.lit_data import LitData from litgpt.data.longform import LongForm from litgpt.data.microllama import MicroLlama +from litgpt.data.multiturn_json_data import MultiturnJSON from litgpt.data.openwebtext import OpenWebText from litgpt.data.text_files import TextFiles from litgpt.data.tinyllama import TinyLlama @@ -27,6 +28,7 @@ "LitData", "DataModule", "LongForm", + "MultiturnJSON", "MultiturnSFTDataset", "OpenWebText", "SFTDataset", diff --git a/litgpt/data/multiturn_json_data.py b/litgpt/data/multiturn_json_data.py new file mode 100644 index 0000000000..5a62622719 --- /dev/null +++ b/litgpt/data/multiturn_json_data.py @@ -0,0 +1,189 @@ +# 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, MultiturnSFTDataset, get_sft_collate_fn +from litgpt.prompts import PromptStyle +from litgpt.tokenizer import Tokenizer + + +@dataclass +class MultiturnJSON(DataModule): + """Loads JSON or JSONL multi-turn conversation data for supervised finetuning. + + 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 a 'messages' key: a + list of `{"role": "system"|"user"|"assistant", "content": str}` turns, ending on an + "assistant" turn (that's the training target).""" + + json_path: Path + + prompt_style: str | PromptStyle + """The style to apply to instruction prompts. Required for multi-turn SFT. See `litgpt.prompts` for a list of available styles.""" + + mask_prompt: bool = True + """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`.""" + 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: MultiturnSFTDataset | None = field(default=None, init=False, repr=False) + test_dataset: MultiturnSFTDataset | 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) + if not self.prompt_style.supports_multiturn: + raise ValueError( + f"The prompt style {self.prompt_style.__class__.__name__} does not support multi-turn conversations. " + "Please choose a prompt style that supports multi-turn conversations." + ) + + 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 = MultiturnSFTDataset( + 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, + transform=to_messages, + ) + self.test_dataset = MultiturnSFTDataset( + 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, + transform=to_messages, + ) + + 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 + + +_SHAREGPT_ROLE_MAP = {"system": "system", "human": "user", "gpt": "assistant"} + + +def sharegpt_to_messages(example: dict) -> list[dict[str, str]]: + """Reshape a ShareGPT-style example (``{"conversations": [{"from": ..., "value": ...}]}``) into + the ``{"role": ..., "content": ...}`` turns MultiturnSFTDataset expects.""" + return [{"role": _SHAREGPT_ROLE_MAP[turn["from"]], "content": turn["value"]} for turn in example["conversations"]] + + +def to_messages(example: dict) -> list[dict[str, str]]: + """Auto-detect whether an example is OpenAI-style (``"messages"``) or ShareGPT-style + (``"conversations"``) and reshape it into ``{"role": ..., "content": ...}`` turns. Raises if + neither key is present.""" + if "messages" in example: + return example["messages"] + if "conversations" in example: + return sharegpt_to_messages(example) + raise ValueError( + "Could not determine the conversation format for example with keys " + f"{list(example.keys())}. Expected a 'messages' (OpenAI-style) or 'conversations' " + "(ShareGPT-style) key." + ) + + +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`.") diff --git a/litgpt/utils.py b/litgpt/utils.py index ec79fa0764..308bb037c9 100644 --- a/litgpt/utils.py +++ b/litgpt/utils.py @@ -13,7 +13,7 @@ import subprocess import sys import warnings -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import asdict, dataclass, is_dataclass from io import BytesIO from pathlib import Path @@ -866,27 +866,41 @@ def create_finetuning_performance_report(training_time, token_counts, device_typ return output +def _instruction_of(example: Any, transform: Callable[[Any], Any] | None = None) -> str: + """Extract a representative "instruction" string from a raw dataset example, for use as an + eval-time generation prompt. Handles both single-turn examples (``{"instruction": ...}``) and + multi-turn conversations (a list of ``{"role", "content"}`` turns, taking the last user turn) + — applying the dataset's own ``transform`` first, if any, so this works regardless of the + example's original on-disk shape (e.g. MultiturnJSON's OpenAI/ShareGPT auto-detection).""" + if transform is not None: + example = transform(example) + if isinstance(example, list): + user_turns = [turn["content"] for turn in example if turn["role"] == "user"] + return user_turns[-1] if user_turns else "" + return example["instruction"] + + def select_sft_generate_example(eval, data): if eval.evaluate_example == "first": if len(data.test_dataset.data): - instruction = data.test_dataset.data[0]["instruction"] + instruction = _instruction_of(data.test_dataset.data[0], data.test_dataset.transform) else: - instruction = data.train_dataset.data[0]["instruction"] + instruction = _instruction_of(data.train_dataset.data[0], data.train_dataset.transform) elif eval.evaluate_example == "random": if len(data.test_dataset.data): random_idx = random.randint(0, len(data.test_dataset.data) - 1) - instruction = data.test_dataset.data[random_idx]["instruction"] + instruction = _instruction_of(data.test_dataset.data[random_idx], data.test_dataset.transform) else: random_idx = random.randint(0, len(data.train_dataset.data) - 1) - instruction = data.train_dataset.data[random_idx]["instruction"] + instruction = _instruction_of(data.train_dataset.data[random_idx], data.train_dataset.transform) elif isinstance(eval.evaluate_example, int): index = eval.evaluate_example if len(data.test_dataset.data) > index: - instruction = data.test_dataset.data[index]["instruction"] + instruction = _instruction_of(data.test_dataset.data[index], data.test_dataset.transform) elif len(data.train_dataset.data) > index: - instruction = data.train_dataset.data[index]["instruction"] + instruction = _instruction_of(data.train_dataset.data[index], data.train_dataset.transform) else: raise IndexError(f"Index {index} is out of range for both test and training datasets.") diff --git a/tests/data/test_multiturn_json.py b/tests/data/test_multiturn_json.py new file mode 100644 index 0000000000..2b2792334b --- /dev/null +++ b/tests/data/test_multiturn_json.py @@ -0,0 +1,216 @@ +# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. +import json + +import pytest + +from litgpt.data import MultiturnJSON +from litgpt.data.multiturn_json_data import sharegpt_to_messages, to_messages +from litgpt.prompts import PromptStyle + + +class MultiturnStyle(PromptStyle): + supports_multiturn = True + + def apply(self, prompt, *, sys_prompt: str | None = None, add_generation_prompt: bool = True, **kwargs) -> str: + text = "".join(f"[{m['role']}]{m['content']}" for m in prompt) + if add_generation_prompt: + text += "[assistant]" + return text + + +def _decode_stripped(mock_tokenizer, row) -> str: + # Right-padded rows get pad_id=0 appended, which MockTokenizer decodes as "\x00". + return mock_tokenizer.decode(row).rstrip("\x00") + + +@pytest.mark.parametrize("as_jsonl", [False, True]) +def test_multiturn_json(as_jsonl, tmp_path, mock_tokenizer): + json_path = tmp_path / ("data.jsonl" if as_jsonl else "data.json") + mock_data = [ + {"messages": [{"role": "user", "content": "Add"}, {"role": "assistant", "content": "4"}]}, + {"messages": [{"role": "user", "content": "Subtract"}, {"role": "assistant", "content": "2"}]}, + {"messages": [{"role": "user", "content": "Multiply"}, {"role": "assistant", "content": "24"}]}, + {"messages": [{"role": "user", "content": "Divide"}, {"role": "assistant", "content": "5"}]}, + ] + + with open(json_path, "w", encoding="utf-8") as fp: + if as_jsonl: + for line in mock_data: + json.dump(line, fp) + fp.write("\n") + else: + json.dump(mock_data, fp) + + data = MultiturnJSON(json_path, val_split_fraction=0.5, prompt_style=MultiturnStyle(), num_workers=0) + data.connect(tokenizer=mock_tokenizer, batch_size=2) + data.prepare_data() # does nothing + data.setup() + + train_dataloader = data.train_dataloader() + val_dataloader = data.val_dataloader() + + assert len(train_dataloader) == 1 + assert len(val_dataloader) == 1 + + train_data = list(train_dataloader) + val_data = list(val_dataloader) + + assert train_data[0]["input_ids"].size(0) == 2 + assert val_data[0]["input_ids"].size(0) == 2 + + decoded = { + _decode_stripped(mock_tokenizer, row) for batch in (*train_data, *val_data) for row in batch["input_ids"] + } + assert decoded == { + "[user]Add[assistant]4", + "[user]Subtract[assistant]2", + "[user]Multiply[assistant]24", + "[user]Divide[assistant]5", + } + + assert isinstance(train_dataloader.dataset.prompt_style, MultiturnStyle) + assert isinstance(val_dataloader.dataset.prompt_style, MultiturnStyle) + + # has attributes from super class `LightningDataModule` + assert data.prepare_data_per_node + + +def test_multiturn_json_sharegpt_format(tmp_path, mock_tokenizer): + json_path = tmp_path / "data.json" + mock_data = [ + {"conversations": [{"from": "human", "value": "Add"}, {"from": "gpt", "value": "4"}]}, + {"conversations": [{"from": "human", "value": "Subtract"}, {"from": "gpt", "value": "2"}]}, + ] + with open(json_path, "w", encoding="utf-8") as fp: + json.dump(mock_data, fp) + + data = MultiturnJSON(json_path, val_split_fraction=0.5, prompt_style=MultiturnStyle(), num_workers=0) + data.connect(tokenizer=mock_tokenizer, batch_size=2) + data.setup() + + decoded = { + _decode_stripped(mock_tokenizer, row) + for batch in (*data.train_dataloader(), *data.val_dataloader()) + for row in batch["input_ids"] + } + assert decoded == {"[user]Add[assistant]4", "[user]Subtract[assistant]2"} + + +def test_multiturn_json_mixed_formats(tmp_path, mock_tokenizer): + # OpenAI-style and ShareGPT-style records in the same file should both work, since format + # detection happens per example. + json_path = tmp_path / "data.json" + mock_data = [ + {"messages": [{"role": "user", "content": "Add"}, {"role": "assistant", "content": "4"}]}, + {"conversations": [{"from": "human", "value": "Subtract"}, {"from": "gpt", "value": "2"}]}, + ] + with open(json_path, "w", encoding="utf-8") as fp: + json.dump(mock_data, fp) + + data = MultiturnJSON(json_path, val_split_fraction=0.5, prompt_style=MultiturnStyle(), num_workers=0) + data.connect(tokenizer=mock_tokenizer, batch_size=2) + data.setup() + + decoded = { + _decode_stripped(mock_tokenizer, row) + for batch in (*data.train_dataloader(), *data.val_dataloader()) + for row in batch["input_ids"] + } + assert decoded == {"[user]Add[assistant]4", "[user]Subtract[assistant]2"} + + +def test_multiturn_json_input_validation(tmp_path): + with pytest.raises(FileNotFoundError, match="The `json_path` must be a file or a directory"): + MultiturnJSON(tmp_path / "not exist", prompt_style=MultiturnStyle()) + + with pytest.raises(ValueError, match="`val_split_fraction` should not be set"): + MultiturnJSON(tmp_path, val_split_fraction=0.5, prompt_style=MultiturnStyle()) + + data = MultiturnJSON(tmp_path, prompt_style=MultiturnStyle()) + data.prepare_data() # does nothing + + # Empty directory + with pytest.raises(FileNotFoundError, match="must be a file or a directory containing"): + data.setup() + + # Only train.json exists + (tmp_path / "train.json").touch() + with pytest.raises(FileNotFoundError, match="must be a file or a directory containing"): + data.setup() + + # When a single file is passed without val_split_fraction, it defaults to 0.05 and warns. + with pytest.warns(UserWarning, match="Defaulting to `val_split_fraction=0.05`"): + data = MultiturnJSON(tmp_path / "train.json", val_split_fraction=None, prompt_style=MultiturnStyle()) + assert data.val_split_fraction == 0.05 + + +def test_multiturn_json_requires_multiturn_style(tmp_path): + with pytest.raises(ValueError, match="does not support multi-turn conversations"): + MultiturnJSON(tmp_path, prompt_style="alpaca") + + +@pytest.mark.parametrize("as_jsonl", [False, True]) +def test_multiturn_json_with_splits(as_jsonl, tmp_path, mock_tokenizer): + mock_train_data = [ + {"messages": [{"role": "user", "content": "Add"}, {"role": "assistant", "content": "4"}]}, + {"messages": [{"role": "user", "content": "Subtract"}, {"role": "assistant", "content": "2"}]}, + {"messages": [{"role": "user", "content": "Multiply"}, {"role": "assistant", "content": "24"}]}, + ] + mock_test_data = [ + {"messages": [{"role": "user", "content": "Divide"}, {"role": "assistant", "content": "5"}]}, + ] + + train_file = tmp_path / ("train.jsonl" if as_jsonl else "train.json") + val_file = tmp_path / ("val.jsonl" if as_jsonl else "val.json") + + with open(train_file, "w", encoding="utf-8") as fp: + if as_jsonl: + for line in mock_train_data: + json.dump(line, fp) + fp.write("\n") + else: + json.dump(mock_train_data, fp) + with open(val_file, "w", encoding="utf-8") as fp: + if as_jsonl: + for line in mock_test_data: + json.dump(line, fp) + fp.write("\n") + else: + json.dump(mock_test_data, fp) + + data = MultiturnJSON(tmp_path, prompt_style=MultiturnStyle(), num_workers=0) + data.connect(tokenizer=mock_tokenizer, batch_size=2) + data.prepare_data() # does nothing + data.setup() + + train_dataloader = data.train_dataloader() + val_dataloader = data.val_dataloader() + + assert len(train_dataloader) == 2 + assert len(val_dataloader) == 1 + + +def test_to_messages_openai_style(): + example = {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]} + assert to_messages(example) == example["messages"] + + +def test_to_messages_sharegpt_style(): + example = { + "conversations": [ + {"from": "system", "value": "sys"}, + {"from": "human", "value": "hi"}, + {"from": "gpt", "value": "hello"}, + ] + } + assert to_messages(example) == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + assert to_messages(example) == sharegpt_to_messages(example) + + +def test_to_messages_unknown_format_raises(): + with pytest.raises(ValueError, match="Could not determine the conversation format"): + to_messages({"foo": "bar"}) diff --git a/tests/test_utils.py b/tests/test_utils.py index ae10780826..047e6b000a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -25,6 +25,7 @@ _TENSORBOARD_AVAILABLE, _WANDB_AVAILABLE, ) +from litgpt.data.multiturn_json_data import to_messages from litgpt.parser_config import save_hyperparameters from litgpt.utils import ( CLI, @@ -811,6 +812,8 @@ def test_select_sft_generate_example(): data_mock.test_dataset.data = test_dataset["data"] data_mock.train_dataset.data = train_dataset["data"] + data_mock.test_dataset.transform = None + data_mock.train_dataset.transform = None # Test "first" instruction from test dataset eval_mock.evaluate_example = "first" @@ -857,3 +860,43 @@ def test_select_sft_generate_example(): eval_mock.evaluate_example = "unknown" with pytest.raises(ValueError): select_sft_generate_example(eval_mock, data_mock) + + +def test_select_sft_generate_example_multiturn(): + eval_mock = mock.MagicMock() + eval_mock.evaluate_example = "first" + data_mock = mock.MagicMock() + + # MultiturnSFTDataset used directly (no transform): .data is already a list of + # {"role", "content"} turns per example. + data_mock.test_dataset.data = [ + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}], + ] + data_mock.test_dataset.transform = None + data_mock.train_dataset.data = [] + data_mock.train_dataset.transform = None + instruction = select_sft_generate_example(eval_mock, data_mock) + assert instruction == "hi" + + # MultiturnJSON-style: raw JSON records with a transform (e.g. to_messages) that must run + # first to reshape them into {"role", "content"} turns. + data_mock.test_dataset.data = [ + {"conversations": [{"from": "human", "value": "yo"}, {"from": "gpt", "value": "sup"}]}, + ] + data_mock.test_dataset.transform = to_messages + instruction = select_sft_generate_example(eval_mock, data_mock) + assert instruction == "yo" + + # Multiple user turns: the *last* user turn should be used. + data_mock.test_dataset.data = [ + { + "messages": [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": "follow-up question"}, + ] + }, + ] + data_mock.test_dataset.transform = to_messages + instruction = select_sft_generate_example(eval_mock, data_mock) + assert instruction == "follow-up question"