data_utils.py 7.05 KB
Newer Older
chenych's avatar
chenych committed
1
# Copyright 2025 the LlamaFactory team.
chenych's avatar
chenych committed
2
3
4
5
6
7
8
9
10
11
12
13
14
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

chenych's avatar
chenych committed
15
import json
chenych's avatar
chenych committed
16
from enum import Enum, unique
chenych's avatar
chenych committed
17
from typing import TYPE_CHECKING, Any, Optional, TypedDict, Union
chenych's avatar
chenych committed
18

chenych's avatar
chenych committed
19
import fsspec
chenych's avatar
chenych committed
20
21
from datasets import DatasetDict, concatenate_datasets, interleave_datasets

luopl's avatar
luopl committed
22
from ..extras import logging
chenych's avatar
chenych committed
23
24
25
26
27
28
29
30


if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset

    from ..hparams import DataArguments


luopl's avatar
luopl committed
31
logger = logging.get_logger(__name__)
chenych's avatar
chenych committed
32
33


chenych's avatar
chenych committed
34
SLOTS = list[Union[str, set[str], dict[str, str]]]
chenych's avatar
chenych committed
35
36
37
38
39
40
41
42
43
44
45
46
47


@unique
class Role(str, Enum):
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"
    FUNCTION = "function"
    OBSERVATION = "observation"


class DatasetModule(TypedDict):
    train_dataset: Optional[Union["Dataset", "IterableDataset"]]
chenych's avatar
chenych committed
48
    eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]]
chenych's avatar
chenych committed
49
50
51


def merge_dataset(
chenych's avatar
chenych committed
52
    all_datasets: list[Union["Dataset", "IterableDataset"]], data_args: "DataArguments", seed: int
chenych's avatar
chenych committed
53
) -> Union["Dataset", "IterableDataset"]:
chenych's avatar
chenych committed
54
    r"""Merge multiple datasets to a unified dataset."""
chenych's avatar
chenych committed
55
56
    if len(all_datasets) == 1:
        return all_datasets[0]
chenych's avatar
chenych committed
57

chenych's avatar
chenych committed
58
59
    elif data_args.mix_strategy == "concat":
        if data_args.streaming:
luopl's avatar
luopl committed
60
            logger.warning_rank0_once("The samples between different datasets will not be mixed in streaming mode.")
chenych's avatar
chenych committed
61
62

        return concatenate_datasets(all_datasets)
chenych's avatar
chenych committed
63

chenych's avatar
chenych committed
64
65
    elif data_args.mix_strategy.startswith("interleave"):
        if not data_args.streaming:
luopl's avatar
luopl committed
66
            logger.warning_rank0_once("We recommend using `mix_strategy=concat` in non-streaming mode.")
chenych's avatar
chenych committed
67
68
69
70
71
72
73

        return interleave_datasets(
            datasets=all_datasets,
            probabilities=data_args.interleave_probs,
            seed=seed,
            stopping_strategy="first_exhausted" if data_args.mix_strategy.endswith("under") else "all_exhausted",
        )
chenych's avatar
chenych committed
74

chenych's avatar
chenych committed
75
    else:
luopl's avatar
luopl committed
76
        raise ValueError(f"Unknown mixing strategy: {data_args.mix_strategy}.")
chenych's avatar
chenych committed
77
78
79


def split_dataset(
chenych's avatar
chenych committed
80
    dataset: Optional[Union["Dataset", "IterableDataset"]],
chenych's avatar
chenych committed
81
    eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]],
chenych's avatar
chenych committed
82
83
    data_args: "DataArguments",
    seed: int,
shihm's avatar
uodata  
shihm committed
84
85
) -> tuple[dict, dict]:
    r"""Split the dataset and returns two dicts containing train set and validation set.
luopl's avatar
luopl committed
86

chenych's avatar
chenych committed
87
    Support both map dataset and iterable dataset.
shihm's avatar
uodata  
shihm committed
88
89
90
91

    Returns:
        train_dict: Dictionary containing training data with key "train"
        eval_dict: Dictionary containing evaluation data with keys "validation" or "validation_{name}"
chenych's avatar
chenych committed
92
    """
chenych's avatar
chenych committed
93
94
95
    if eval_dataset is not None and data_args.val_size > 1e-6:
        raise ValueError("Cannot specify `val_size` if `eval_dataset` is not None.")

shihm's avatar
uodata  
shihm committed
96
97
98
    # the train and eval better to in dict dtype and separately return for cpode clearly and good handle outside
    train_dict, eval_dict = {}, {}

chenych's avatar
chenych committed
99
100
101
102
103
104
    if dataset is not None:
        if data_args.streaming:
            dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=seed)

        if data_args.val_size > 1e-6:
            if data_args.streaming:
shihm's avatar
uodata  
shihm committed
105
106
                eval_dict["validation"] = dataset.take(int(data_args.val_size))
                train_dict["train"] = dataset.skip(int(data_args.val_size))
chenych's avatar
chenych committed
107
108
            else:
                val_size = int(data_args.val_size) if data_args.val_size > 1 else data_args.val_size
shihm's avatar
uodata  
shihm committed
109
110
111
                split_result = dataset.train_test_split(test_size=val_size, seed=seed)
                train_dict["train"] = split_result["train"]
                eval_dict["validation"] = split_result["test"]
chenych's avatar
chenych committed
112
        else:
shihm's avatar
uodata  
shihm committed
113
            train_dict["train"] = dataset
chenych's avatar
chenych committed
114
115
116

    if eval_dataset is not None:
        if isinstance(eval_dataset, dict):
shihm's avatar
uodata  
shihm committed
117
118
            for name, data in eval_dataset.items():
                eval_dict[f"validation_{name}"] = data
chenych's avatar
chenych committed
119
120
121
122
        else:
            if data_args.streaming:
                eval_dataset = eval_dataset.shuffle(buffer_size=data_args.buffer_size, seed=seed)

shihm's avatar
uodata  
shihm committed
123
            eval_dict["validation"] = eval_dataset
chenych's avatar
chenych committed
124

shihm's avatar
uodata  
shihm committed
125
    return train_dict, eval_dict
chenych's avatar
chenych committed
126
127
128


def get_dataset_module(dataset: Union["Dataset", "DatasetDict"]) -> "DatasetModule":
chenych's avatar
chenych committed
129
130
    r"""Convert dataset or dataset dict to dataset module."""
    dataset_module: DatasetModule = {}
chenych's avatar
chenych committed
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
    if isinstance(dataset, DatasetDict):  # dataset dict
        if "train" in dataset:
            dataset_module["train_dataset"] = dataset["train"]

        if "validation" in dataset:
            dataset_module["eval_dataset"] = dataset["validation"]
        else:
            eval_dataset = {}
            for key in dataset.keys():
                if key.startswith("validation_"):
                    eval_dataset[key[len("validation_") :]] = dataset[key]

            if len(eval_dataset):
                dataset_module["eval_dataset"] = eval_dataset

    else:  # single dataset
        dataset_module["train_dataset"] = dataset

    return dataset_module
chenych's avatar
chenych committed
150
151


chenych's avatar
chenych committed
152
153
def setup_fs(path: str, anon: bool = False) -> "fsspec.AbstractFileSystem":
    r"""Set up a filesystem object based on the path protocol."""
chenych's avatar
chenych committed
154
155
156
157
158
159
    storage_options = {"anon": anon} if anon else {}
    if path.startswith("s3://"):
        fs = fsspec.filesystem("s3", **storage_options)
    elif path.startswith(("gs://", "gcs://")):
        fs = fsspec.filesystem("gcs", **storage_options)
    else:
chenych's avatar
chenych committed
160
161
162
163
164
        raise ValueError(f"Unsupported protocol in path: {path}. Use 's3://' or 'gs://'.")

    if not fs.exists(path):
        raise ValueError(f"Path does not exist: {path}.")

chenych's avatar
chenych committed
165
166
167
    return fs


chenych's avatar
chenych committed
168
169
170
171
172
173
174
175
176
177
178
def _read_json_with_fs(fs: "fsspec.AbstractFileSystem", path: str) -> list[Any]:
    r"""Helper function to read JSON/JSONL files using fsspec."""
    with fs.open(path, "r") as f:
        if path.endswith(".jsonl"):
            return [json.loads(line) for line in f if line.strip()]
        else:
            return json.load(f)


def read_cloud_json(cloud_path: str) -> list[Any]:
    r"""Read a JSON/JSONL file from cloud storage (S3 or GCS).
chenych's avatar
chenych committed
179
180

    Args:
chenych's avatar
chenych committed
181
        cloud_path: str
chenych's avatar
chenych committed
182
183
184
185
186
            Cloud path in the format:
            - 's3://bucket-name/file.json' for AWS S3
            - 'gs://bucket-name/file.jsonl' or 'gcs://bucket-name/file.jsonl' for Google Cloud Storage
    """
    try:
chenych's avatar
chenych committed
187
        fs = setup_fs(cloud_path, anon=True)  # try with anonymous access first
chenych's avatar
chenych committed
188
    except Exception:
chenych's avatar
chenych committed
189
        fs = setup_fs(cloud_path)  # try again with credentials
chenych's avatar
chenych committed
190

chenych's avatar
chenych committed
191
192
193
194
195
    # filter out non-JSON files
    files = [x["Key"] for x in fs.listdir(cloud_path)] if fs.isdir(cloud_path) else [cloud_path]
    files = filter(lambda file: file.endswith(".json") or file.endswith(".jsonl"), files)
    if not files:
        raise ValueError(f"No JSON/JSONL files found in the specified path: {cloud_path}.")
chenych's avatar
chenych committed
196

chenych's avatar
chenych committed
197
    return sum([_read_json_with_fs(fs, file) for file in files], [])