Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_dataset.py: 0%
88 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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)
29from hyper_parallel.data.indexed.indexed_pretrain_dataset import (
30 GPTDataset,
31 GPTFromMRDataset,
32 MockGPTDataset,
33)
34from hyper_parallel.data.indexed.indexed_split_builder import IndexedDatasetSplitBuilder
35from hyper_parallel.data.parallel import DataLoaderParallelContext
37DataPath: TypeAlias = str | Sequence[str]
38DatasetSplits: TypeAlias = tuple[Any | None, Any | None, Any | None]
39SplitBuilder: TypeAlias = Callable[..., DatasetSplits]
40InstructionSplitBuilder: TypeAlias = Callable[..., Sequence[Any | None]]
41logger = get_dataset_logger(__name__)
44class IndexedPretrainDatasetBuilder:
45 """Build GPT pretraining datasets from indexed token files.
47 Its ``build`` method keeps the indexed provider flow visible while each
48 private step can be implemented and reviewed independently.
49 """
51 def __init__(
52 self,
53 *,
54 data_path: DataPath,
55 data_config: Mapping[str, Any],
56 train_valid_test_num_samples: Sequence[int],
57 dataloader_context: DataLoaderParallelContext | None = None,
58 ) -> None:
59 """Store the inputs required by the indexed pretraining build.
61 Args:
62 data_path: Indexed prefix, directory, or weighted paths. May be
63 omitted when mock data is enabled.
64 data_config: Indexed pretraining build options.
65 train_valid_test_num_samples: Trainer-derived target sizes.
66 dataloader_context: DataLoader ownership and synchronization policy.
67 """
68 self.data_path = data_path
69 self.data_config = data_config
70 self.train_valid_test_num_samples = train_valid_test_num_samples
71 self.dataloader_context = dataloader_context or DataLoaderParallelContext(
72 data_index_cache=bool(data_config.get("data_index_cache", False))
73 )
74 self.split_builder = IndexedDatasetSplitBuilder(self.dataloader_context)
76 def build(self) -> DatasetSplits:
77 """Build train, validation, and test indexed datasets.
79 Returns:
80 Train, validation, and test datasets.
81 """
82 # Discover indexed files and add their default blend weights.
83 # Mock Datasets generate their samples and do not need filesystem
84 # discovery. Real indexed Datasets still require an explicit path.
85 mock_data = bool(self.data_config.get("mock_data", False))
86 logger.debug("Building indexed Dataset: mock=%s, data_path=%s", mock_data, self.data_path)
87 if mock_data:
88 data_paths = []
89 else:
90 if self.data_path is None:
91 raise ValueError("data_path is required when mock_data is false")
93 data_paths = self._resolve_data_paths()
95 # Normalize Dataset options. Sample counts remain separate and are
96 # passed directly to the split builder.
97 config = self._build_gpt_dataset_config(data_paths)
98 logger.debug(
99 "Indexed Dataset config: simple_blend=%s, lazy=%s, shared_blend=%s, independent_blends=%s",
100 config.simple_blend, config.data_lazy_load, bool(config.blend), bool(config.blend_per_split),
101 )
103 # Follow the Provider boundary: instruction and pretraining datasets
104 # share one public entry but own separate construction pipelines.
105 dataset_splits = self._build_dataset_splits(config)
106 split_types = tuple(type(split).__name__ if split is not None else None for split in dataset_splits)
107 logger.debug("Built indexed Dataset splits=%s", split_types)
108 return dataset_splits
110 def _build_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
111 """Select the instruction or indexed-pretraining Dataset pipeline."""
112 if bool(self.data_config.get("is_instruction_dataset", False)):
113 logger.debug("Selecting instruction Dataset pipeline")
114 dataset_splits = self._build_instruction_dataset_splits(config)
115 return dataset_splits
117 dataset_splits = self._build_pretrain_dataset_splits(config)
118 return dataset_splits
120 def _build_pretrain_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
121 """Build GPT pretraining splits through the selected Dataset and blend types."""
122 # Select GPTDataset / MockGPTDataset / GPTFromMRDataset.
123 dataset_type = self._select_dataset_type(config)
125 # Select the standard or simple blended Dataset builder.
126 split_builder = self._select_split_builder(config)
127 logger.debug("Selected indexed Dataset: type=%s, split_mode=%s", dataset_type.__name__, config.simple_blend)
129 # Build train, validation, and test datasets.
130 datasets = split_builder(
131 dataset_type=dataset_type,
132 train_valid_test_num_samples=self.train_valid_test_num_samples,
133 config=config,
134 )
135 return datasets
137 def _build_instruction_dataset_splits(self, config: GPTDatasetConfig) -> DatasetSplits:
138 """Build packed instruction splits through the configured implementation.
140 The concrete packed-dataset implementation remains outside this provider.
142 Raises:
143 NotImplementedError: If no instruction Dataset builder is configured.
144 ValueError: If the configured builder does not return three splits.
145 """
146 instruction_builder = self.data_config.get("instruction_dataset_builder")
147 if not callable(instruction_builder):
148 raise NotImplementedError(
149 "is_instruction_dataset=True requires an instruction_dataset_builder"
150 )
152 instruction_split_builder: InstructionSplitBuilder = instruction_builder
154 data_prefix = config.blend if config.blend is not None else self.data_path
155 datasets = instruction_split_builder(
156 data_prefix=data_prefix,
157 splits_string=config.split,
158 train_valid_test_num_samples=self.train_valid_test_num_samples,
159 seq_length=config.sequence_length,
160 seed=config.random_seed,
161 tokenizer=config.tokenizer,
162 )
163 if not isinstance(datasets, Sequence) or len(datasets) != 3:
164 raise ValueError("instruction_dataset_builder must return train, validation, and test")
166 dataset_splits = (datasets[0], datasets[1], datasets[2])
167 return dataset_splits
169 def _resolve_data_paths(self) -> list[str]:
170 """Delegate indexed path discovery to the config build module."""
171 if self.data_path is None:
172 raise ValueError("data_path is required when mock_data is false")
174 data_paths = resolve_data_paths(
175 self.data_path,
176 distributed_walk=bool(self.data_config["distributed_walk"]),
177 )
178 return data_paths
180 def _build_gpt_dataset_config(self, data_paths: Sequence[str]) -> GPTDatasetConfig:
181 """Delegate GPT Dataset option normalization."""
182 config = build_gpt_dataset_config(data_paths, self.data_config)
183 return config
185 @staticmethod
186 def _select_dataset_type(config: GPTDatasetConfig) -> type:
187 """Select GPT, Mock GPT, or MR GPT Dataset."""
188 if config.mock:
189 return MockGPTDataset
191 if config.is_dataset_from_mr:
192 return GPTFromMRDataset
194 return GPTDataset
196 def _select_split_builder(self, config: GPTDatasetConfig) -> SplitBuilder:
197 """
198 Select the standard or simple train/valid/test builder.
199 "no": BlendedDataset
200 "inter/intra": SimpleBlendedDataset
201 """
202 if config.simple_blend == "no":
203 split_builder = partial(self.split_builder.build, blend_mode="no")
204 return split_builder
206 if not config.is_dataset_from_mr:
207 raise ValueError("simple_blend values other than 'no' require is_dataset_from_mr=True")
209 if config.simple_blend not in {"inter", "intra"}:
210 raise ValueError("simple_blend must be one of 'no', 'inter', or 'intra'; "
211 f"got {config.simple_blend!r}"
212 )
214 split_builder = partial(self.split_builder.build, blend_mode=config.simple_blend)
215 return split_builder
218def build_indexed_dataset(
219 *,
220 data_path: DataPath = None,
221 data_config: Mapping[str, Any],
222 train_valid_test_num_samples: Sequence[int],
223 dataloader_context: DataLoaderParallelContext | None = None,
224) -> DatasetSplits:
225 """Build indexed pretraining datasets through the dedicated builder.
227 Args:
228 data_path: Indexed prefix, directory, or weighted paths. May be omitted
229 when mock data is enabled.
230 data_config: Indexed pretraining build options.
231 train_valid_test_num_samples: Trainer-derived target sizes.
232 dataloader_context: DataLoader ownership and synchronization policy.
234 Returns:
235 Train, validation, and test datasets.
236 """
237 builder = IndexedPretrainDatasetBuilder(
238 data_path=data_path,
239 data_config=data_config,
240 train_valid_test_num_samples=train_valid_test_num_samples,
241 dataloader_context=dataloader_context,
242 )
243 datasets = builder.build()
244 return datasets