Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / build_dataset.py: 59%
68 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"""Public Online and Indexed text Dataset builders."""
17from __future__ import annotations
19from collections.abc import Callable, Mapping, Sequence
20from typing import Any
22from hyper_parallel.data.dataset_logging import get_dataset_logger
23from hyper_parallel.data.indexed.indexed_dataset import (
24 build_indexed_dataset as _build_indexed_dataset,
25)
26from hyper_parallel.data.online import (
27 IterableTransformDataset,
28 MappingTransformDataset,
29 OnlineDataPath,
30 build_online_iterable_source,
31 build_online_mapping_source,
32)
33from hyper_parallel.data.parallel import (
34 DataLoaderParallelContext,
35 create_dataloader_parallel_context,
36)
38logger = get_dataset_logger(__name__)
41def _build_dataloader_context(
42 mesh_context: Any,
43 data_config: Mapping[str, Any],
44) -> DataLoaderParallelContext | None:
45 """Build DataLoader ownership from the runtime mesh when available."""
46 if mesh_context is None:
47 return None
49 dataloader_context = create_dataloader_parallel_context(
50 mesh_context,
51 data_index_cache=bool(data_config.get("data_index_cache", False)),
52 shared_storage=not bool(data_config.get("no_shared_storage", False)),
53 )
54 return dataloader_context
57def _get_train_valid_test_num_samples(training_config: Any) -> tuple[int, int, int]:
58 """Calculate train, validation, and test sizes from the training plan."""
59 if training_config is None:
60 raise ValueError("Dataset sample planning requires a training configuration")
62 global_batch_size = training_config.global_batch_size
63 if training_config.train_iters is not None:
64 train_iters = training_config.train_iters
65 elif training_config.train_samples:
66 train_iters = training_config.train_samples // global_batch_size
67 else:
68 raise ValueError("training.train_iters and training.train_samples cannot both be None")
70 train_samples = training_config.train_samples or train_iters * global_batch_size
71 eval_iters = training_config.eval_iters
72 valid_iters = (train_iters // eval_iters + 1) * eval_iters if eval_iters else 0
73 split_sizes = (train_samples, valid_iters * global_batch_size, eval_iters * global_batch_size)
74 logger.debug("Dataset target sizes: train=%d, validation=%d, test=%d", *split_sizes)
75 return split_sizes
78def build_online_text_mapping_dataset(
79 *,
80 data_config: Mapping[str, Any],
81 data_path: OnlineDataPath | None = None,
82 transform: Callable[[Any], Any] | None = None,
83 training_config: Any = None,
84 sampler_type: str = "single",
85) -> Any:
86 """Build an Online Mapping Dataset and apply its text transform.
88 Args:
89 data_config: Online Mapping source options.
90 data_path: Optional local source path, ordered paths, or pre-split
91 train/valid/test path mapping. Use ``data_config.sources`` for
92 multiple sources.
93 transform: Plaintext or conversation sample transform.
94 training_config: Training plan providing the random seed and split sizes.
95 sampler_type: Mapping sampler policy controlling source-level shuffle.
97 Returns:
98 A transformed Online Dataset on each DataLoader-owning rank.
100 Raises:
101 ValueError: If no Online text transform is configured.
102 """
103 if transform is None:
104 raise ValueError("Online Dataset requires a plaintext or conversation data_transform")
106 dataset_config = dict(data_config)
107 training_seed = getattr(training_config, "seed", None)
108 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
109 train_valid_test_num_samples = None
110 if training_config is not None:
111 train_valid_test_num_samples = _get_train_valid_test_num_samples(training_config)
112 sample_filter = transform.is_valid_sample
113 source_dataset = build_online_mapping_source(
114 data_path=data_path,
115 data_config=dataset_config,
116 sample_filter=sample_filter,
117 sampler_type=sampler_type,
118 train_valid_test_num_samples=train_valid_test_num_samples,
119 )
120 text_dataset = MappingTransformDataset.apply(source_dataset, transform)
121 return text_dataset
124def build_online_iterable_dataset(
125 *,
126 data_config: Mapping[str, Any],
127 data_path: OnlineDataPath | None = None,
128 transform: Callable[[Any], Any] | None = None,
129 dataloader_context: DataLoaderParallelContext | None = None,
130 mesh_context: Any = None,
131 training_config: Any = None,
132) -> Any:
133 """Build an Online Iterable Dataset and apply its text transform.
135 Args:
136 data_config: Online Iterable source options.
137 data_path: Optional local source path, ordered paths, or Hub Dataset
138 ID. Use ``data_config.sources`` for multiple sources.
139 transform: Plaintext or conversation sample transform.
140 dataloader_context: Optional explicit DataLoader ownership context.
141 mesh_context: Runtime mesh used to derive Iterable source ownership.
142 training_config: Training plan providing the shuffle seed.
144 Returns:
145 A transformed Online Iterable Dataset on each DataLoader-owning rank.
147 Raises:
148 ValueError: If no Online text transform is configured.
149 """
150 if transform is None:
151 raise ValueError("Online Dataset requires a plaintext or conversation data_transform")
153 dataset_config = dict(data_config)
154 training_seed = getattr(training_config, "seed", None)
155 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
156 sample_filter = transform.is_valid_sample
157 if dataloader_context is None:
158 dataloader_context = _build_dataloader_context(mesh_context, dataset_config)
160 source_dataset = build_online_iterable_source(
161 data_path=data_path,
162 data_config=dataset_config,
163 dataloader_context=dataloader_context,
164 sample_filter=sample_filter,
165 )
166 text_dataset = IterableTransformDataset.apply(source_dataset, transform)
167 return text_dataset
170def build_indexed_text_dataset(
171 *,
172 data_config: Mapping[str, Any],
173 data_path: str | Sequence[str] | None = None,
174 tokenizer: Any = None,
175 train_valid_test_num_samples: Sequence[int] | None = None,
176 dataloader_context: DataLoaderParallelContext | None = None,
177 mesh_context: Any = None,
178 training_config: Any = None,
179) -> Any:
180 """Build train, validation, and test Datasets from Indexed token files.
182 Args:
183 data_config: Indexed Dataset and split options.
184 data_path: Indexed prefix, directory, or ordered paths.
185 tokenizer: Tokenizer providing vocabulary and EOD metadata.
186 train_valid_test_num_samples: Optional explicit split sizes.
187 dataloader_context: Optional explicit DataLoader ownership context.
188 mesh_context: Runtime mesh used to derive DataLoader ownership.
189 training_config: Training plan used to derive split sizes and seed.
191 Returns:
192 Train, validation, and test Indexed Datasets.
194 Raises:
195 ValueError: If required path or training information is missing.
196 """
197 dataset_config = dict(data_config)
198 dataset_config.setdefault("distributed_walk", True)
199 training_seed = getattr(training_config, "seed", None)
200 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
201 if tokenizer is not None:
202 dataset_config["tokenizer"] = tokenizer
204 if data_path is None and not bool(dataset_config.get("mock_data", False)):
205 raise ValueError("Indexed Dataset requires data_path")
207 if train_valid_test_num_samples is None:
208 train_valid_test_num_samples = _get_train_valid_test_num_samples(training_config)
210 if dataloader_context is None:
211 dataloader_context = _build_dataloader_context(mesh_context, dataset_config)
213 indexed_datasets = _build_indexed_dataset(
214 data_path=data_path,
215 data_config=dataset_config,
216 train_valid_test_num_samples=train_valid_test_num_samples,
217 dataloader_context=dataloader_context,
218 )
219 return indexed_datasets