Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_dataset.py: 81%
83 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Indexed pretraining LLM datasets stored in ``.idx/.bin`` files."""
17from __future__ import annotations
19from collections.abc import Callable, Mapping, Sequence
20from functools import partial
21from typing import Any, TypeAlias
23from hyper_parallel.data.dataset_logging import get_dataset_logger
24from hyper_parallel.data.indexed.indexed_data_config import (
25 GPTDatasetConfig,
26 build_gpt_dataset_config,
27 resolve_data_paths,
28 resolve_dataset_type,
29)
30from hyper_parallel.data.indexed.indexed_pretrain_dataset import SFTFromMRDataset
31from hyper_parallel.data.indexed.indexed_split_builder import IndexedDatasetSplitBuilder
32from hyper_parallel.data.parallel import DataLoaderParallelContext
34DataPath: TypeAlias = str | Sequence[str]
35DatasetSplits: TypeAlias = tuple[Any | None, Any | None, Any | None]
36SplitBuilder: TypeAlias = Callable[..., DatasetSplits]
37InstructionSplitBuilder: TypeAlias = Callable[..., Sequence[Any | None]]
38logger = get_dataset_logger(__name__)
41class IndexedPretrainDatasetBuilder:
42 """Build GPT pretraining datasets from indexed token files.
44 Its ``build`` method keeps the indexed provider flow visible while each
45 private step can be implemented and reviewed independently.
46 """
48 def __init__(
49 self,
50 *,
51 data_path: DataPath,
52 data_config: Mapping[str, Any],
53 train_valid_test_num_samples: Sequence[int],
54 dataloader_context: DataLoaderParallelContext | None = None,
55 ) -> None:
56 """Store the inputs required by the indexed pretraining build.
58 Args:
59 data_path: Indexed prefix, directory, or weighted paths. May be
60 omitted when mock data is enabled.
61 data_config: Indexed pretraining build options.
62 train_valid_test_num_samples: Trainer-derived target sizes.
63 dataloader_context: DataLoader ownership and synchronization policy.
64 """
65 self.data_path = data_path
66 self.data_config = data_config
67 self.train_valid_test_num_samples = train_valid_test_num_samples
68 self.dataloader_context = dataloader_context or DataLoaderParallelContext(
69 data_index_cache=bool(data_config.get("data_index_cache", False))
70 )
71 self.split_builder = IndexedDatasetSplitBuilder(self.dataloader_context)
73 def build(self) -> DatasetSplits:
74 """Build train, validation, and test indexed datasets.
76 Returns:
77 Train, validation, and test datasets.
78 """
79 # Discover indexed files and add their default blend weights.
80 # Mock Datasets generate their samples and do not need filesystem
81 # discovery. Real indexed Datasets still require an explicit path.
82 mock_data = bool(self.data_config.get("mock_data", False))
83 logger.debug("Building indexed Dataset: mock=%s, data_path=%s", mock_data, self.data_path)
84 if mock_data:
85 data_paths = []
86 else:
87 if self.data_path is None:
88 raise ValueError("data_path is required when mock_data is false")
90 data_paths = self._resolve_data_paths()
92 # Normalize Dataset options. Sample counts remain separate and are
93 # passed directly to the split builder.
94 config = self._build_gpt_dataset_config(data_paths)
95 logger.debug(
96 "Indexed Dataset config: simple_blend=%s, lazy=%s, shared_blend=%s, independent_blends=%s",
97 config.simple_blend, config.data_lazy_load, bool(config.blend), bool(config.blend_per_split),
98 )
100 # Follow the Provider boundary: instruction and pretraining datasets
101 # share one public entry but own separate construction pipelines.
102 dataset_splits = self._build_dataset_splits(config)
103 split_types = tuple(type(split).__name__ if split is not None else None for split in dataset_splits)
104 logger.debug("Built indexed Dataset splits=%s", split_types)
105 return dataset_splits
107 def _build_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
108 """Select the instruction or indexed-pretraining Dataset pipeline."""
109 if bool(self.data_config.get("is_instruction_dataset", False)):
110 logger.debug("Selecting instruction Dataset pipeline")
111 dataset_splits = self._build_instruction_dataset_splits(config)
112 return dataset_splits
114 dataset_splits = self._build_pretrain_dataset_splits(config)
115 return dataset_splits
117 def _build_pretrain_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
118 """Build GPT pretraining splits through the selected Dataset and blend types."""
119 # Select GPTDataset / MockGPTDataset / GPTFromMRDataset.
120 dataset_type = self._select_dataset_type(config)
122 # Select the standard or simple blended Dataset builder.
123 split_builder = self._select_split_builder(config)
124 logger.debug("Selected indexed Dataset: type=%s, split_mode=%s", dataset_type.__name__, config.simple_blend)
126 # Build train, validation, and test datasets.
127 datasets = split_builder(
128 dataset_type=dataset_type,
129 train_valid_test_num_samples=self.train_valid_test_num_samples,
130 config=config,
131 )
132 return datasets
134 def _build_instruction_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
135 """Build packed instruction splits through the configured implementation.
137 The concrete packed-dataset implementation remains outside this provider.
139 Raises:
140 NotImplementedError: If no instruction Dataset builder is configured.
141 ValueError: If the configured builder does not return three splits.
142 """
143 instruction_builder = self.data_config.get("instruction_dataset_builder")
144 if not callable(instruction_builder):
145 raise NotImplementedError(
146 "is_instruction_dataset=True requires an instruction_dataset_builder"
147 )
149 instruction_split_builder: InstructionSplitBuilder = instruction_builder
151 data_prefix = config.blend if config.blend is not None else self.data_path
152 datasets = instruction_split_builder(
153 data_prefix=data_prefix,
154 splits_string=config.split,
155 train_valid_test_num_samples=self.train_valid_test_num_samples,
156 seq_length=config.sequence_length,
157 seed=config.random_seed,
158 tokenizer=config.tokenizer,
159 )
160 if not isinstance(datasets, Sequence) or len(datasets) != 3:
161 raise ValueError("instruction_dataset_builder must return train, validation, and test")
163 dataset_splits = (datasets[0], datasets[1], datasets[2])
164 return dataset_splits
166 def _resolve_data_paths(self) -> list[str]:
167 """Delegate indexed path discovery to the config build module."""
168 if self.data_path is None:
169 raise ValueError("data_path is required when mock_data is false")
171 dataset_type = resolve_dataset_type(self.data_config["low_level_dataset"])
172 data_paths = resolve_data_paths(
173 self.data_path,
174 distributed_walk=bool(self.data_config["distributed_walk"]),
175 file_suffix=".arrow" if issubclass(dataset_type, SFTFromMRDataset) else ".bin",
176 )
177 return data_paths
179 def _build_gpt_dataset_config(self, data_paths: Sequence[str]) -> GPTDatasetConfig:
180 """Delegate GPT Dataset option normalization."""
181 config = build_gpt_dataset_config(data_paths, self.data_config)
182 return config
184 @staticmethod
185 def _select_dataset_type(config: GPTDatasetConfig) -> type:
186 """Select GPT, Mock GPT, or MR GPT Dataset."""
187 return config.low_level_dataset
189 def _select_split_builder(self, config: GPTDatasetConfig) -> SplitBuilder:
190 """
191 Select the standard or simple train/valid/test builder.
192 "no": BlendedDataset
193 "inter/intra": SimpleBlendedDataset
194 """
195 if config.simple_blend == "no":
196 split_builder = partial(self.split_builder.build, blend_mode="no")
197 return split_builder
199 if config.simple_blend not in {"inter", "intra"}:
200 raise ValueError("simple_blend must be one of 'no', 'inter', or 'intra'; "
201 f"got {config.simple_blend!r}"
202 )
204 split_builder = partial(self.split_builder.build, blend_mode=config.simple_blend)
205 return split_builder
208def build_indexed_dataset(
209 *,
210 data_path: DataPath = None,
211 data_config: Mapping[str, Any],
212 train_valid_test_num_samples: Sequence[int],
213 dataloader_context: DataLoaderParallelContext | None = None,
214) -> DatasetSplits:
215 """Build indexed pretraining datasets through the dedicated builder.
217 Args:
218 data_path: Indexed prefix, directory, or weighted paths. May be omitted
219 when mock data is enabled.
220 data_config: Indexed pretraining build options.
221 train_valid_test_num_samples: Trainer-derived target sizes.
222 dataloader_context: DataLoader ownership and synchronization policy.
224 Returns:
225 Train, validation, and test datasets.
226 """
227 builder = IndexedPretrainDatasetBuilder(
228 data_path=data_path,
229 data_config=data_config,
230 train_valid_test_num_samples=train_valid_test_num_samples,
231 dataloader_context=dataloader_context,
232 )
233 datasets = builder.build()
234 return datasets