feat: add MultiturnJSON datamodule + full multi-turn SFT support

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",

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`.")

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.")