Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / build_dataset.py: 0%
67 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"""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) -> Any:
85 """Build an Online Mapping Dataset and apply its text transform.
87 Args:
88 data_config: Online Mapping source options.
89 data_path: Optional local source path, ordered paths, or pre-split
90 train/valid/test path mapping. Use ``data_config.sources`` for
91 multiple sources.
92 transform: Plaintext or conversation sample transform.
93 training_config: Training plan providing the random seed and split sizes.
95 Returns:
96 A transformed Online Dataset on each DataLoader-owning rank.
98 Raises:
99 ValueError: If no Online text transform is configured.
100 """
101 if transform is None:
102 raise ValueError("Online Dataset requires a plaintext or conversation data_transform")
104 dataset_config = dict(data_config)
105 training_seed = getattr(training_config, "seed", None)
106 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
107 train_valid_test_num_samples = None
108 if training_config is not None:
109 train_valid_test_num_samples = _get_train_valid_test_num_samples(training_config)
110 sample_filter = transform.is_valid_sample
111 source_dataset = build_online_mapping_source(
112 data_path=data_path,
113 data_config=dataset_config,
114 sample_filter=sample_filter,
115 train_valid_test_num_samples=train_valid_test_num_samples,
116 )
117 text_dataset = MappingTransformDataset.apply(source_dataset, transform)
118 return text_dataset
121def build_online_iterable_dataset(
122 *,
123 data_config: Mapping[str, Any],
124 data_path: OnlineDataPath | None = None,
125 transform: Callable[[Any], Any] | None = None,
126 dataloader_context: DataLoaderParallelContext | None = None,
127 mesh_context: Any = None,
128 training_config: Any = None,
129) -> Any:
130 """Build an Online Iterable Dataset and apply its text transform.
132 Args:
133 data_config: Online Iterable source options.
134 data_path: Optional local source path, ordered paths, or Hub Dataset
135 ID. Use ``data_config.sources`` for multiple sources.
136 transform: Plaintext or conversation sample transform.
137 dataloader_context: Optional explicit DataLoader ownership context.
138 mesh_context: Runtime mesh used to derive Iterable source ownership.
139 training_config: Training plan providing the shuffle seed.
141 Returns:
142 A transformed Online Iterable Dataset on each DataLoader-owning rank.
144 Raises:
145 ValueError: If no Online text transform is configured.
146 """
147 if transform is None:
148 raise ValueError("Online Dataset requires a plaintext or conversation data_transform")
150 dataset_config = dict(data_config)
151 training_seed = getattr(training_config, "seed", None)
152 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
153 sample_filter = transform.is_valid_sample
154 if dataloader_context is None:
155 dataloader_context = _build_dataloader_context(mesh_context, dataset_config)
157 source_dataset = build_online_iterable_source(
158 data_path=data_path,
159 data_config=dataset_config,
160 dataloader_context=dataloader_context,
161 sample_filter=sample_filter,
162 )
163 text_dataset = IterableTransformDataset.apply(source_dataset, transform)
164 return text_dataset
167def build_indexed_text_dataset(
168 *,
169 data_config: Mapping[str, Any],
170 data_path: str | Sequence[str] | None = None,
171 tokenizer: Any = None,
172 train_valid_test_num_samples: Sequence[int] | None = None,
173 dataloader_context: DataLoaderParallelContext | None = None,
174 mesh_context: Any = None,
175 training_config: Any = None,
176) -> Any:
177 """Build train, validation, and test Datasets from Indexed token files.
179 Args:
180 data_config: Indexed Dataset and split options.
181 data_path: Indexed prefix, directory, or ordered paths.
182 tokenizer: Tokenizer providing vocabulary and EOD metadata.
183 train_valid_test_num_samples: Optional explicit split sizes.
184 dataloader_context: Optional explicit DataLoader ownership context.
185 mesh_context: Runtime mesh used to derive DataLoader ownership.
186 training_config: Training plan used to derive split sizes and seed.
188 Returns:
189 Train, validation, and test Indexed Datasets.
191 Raises:
192 ValueError: If required path or training information is missing.
193 """
194 dataset_config = dict(data_config)
195 training_seed = getattr(training_config, "seed", None)
196 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed)
197 if tokenizer is not None:
198 dataset_config["tokenizer"] = tokenizer
200 if data_path is None and not bool(dataset_config.get("mock_data", False)):
201 raise ValueError("Indexed Dataset requires data_path")
203 if train_valid_test_num_samples is None:
204 train_valid_test_num_samples = _get_train_valid_test_num_samples(training_config)
206 if dataloader_context is None:
207 dataloader_context = _build_dataloader_context(mesh_context, dataset_config)
209 indexed_datasets = _build_indexed_dataset(
210 data_path=data_path,
211 data_config=dataset_config,
212 train_valid_test_num_samples=train_valid_test_num_samples,
213 dataloader_context=dataloader_context,
214 )
215 return indexed_datasets