Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / online / source.py: 12%
247 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"""Shared Online source loading and multi-source construction policy."""
17from __future__ import annotations
19import os
20from collections.abc import Callable, Mapping, Sequence
21from typing import Any, TypeAlias
23from hyper_parallel.data.constants import ONLINE_FILE_FORMATS, ONLINE_SPLIT_NAMES, ONLINE_STOPPING_STRATEGIES
25OnlinePathSpec: TypeAlias = str | Sequence[str]
26OnlineDataPath: TypeAlias = OnlinePathSpec | Mapping[str, OnlinePathSpec]
29class _OnlineSourceBuilder:
30 """Share source parsing and interleave policy across Online adapters."""
32 def __init__(self, data_path: OnlineDataPath | None, data_config: Mapping[str, Any]) -> None:
33 """Store the source and composition configuration."""
34 self.data_path = data_path
35 self.data_config = data_config
37 @property
38 def _random_seed(self) -> int:
39 """Return the common source shuffle and blend seed."""
40 return int(self.data_config.get("random_seed", 42))
42 @property
43 def _stopping_strategy(self) -> str:
44 """Return the multi-source exhaustion policy."""
45 stopping_strategy = str(self.data_config.get("stopping_strategy", "first_exhausted"))
46 if stopping_strategy not in ONLINE_STOPPING_STRATEGIES:
47 raise ValueError(
48 f"Unsupported Online stopping_strategy {stopping_strategy!r}; "
49 f"expected one of {ONLINE_STOPPING_STRATEGIES!r}"
50 )
51 return stopping_strategy
53 def _create_source_loaders(self) -> tuple[list[_OnlineLeafSourceLoader], list[float]]:
54 """Create one leaf source loader per configured source.
56 Example:
57 ``data_path="/data/train.jsonl"`` produces one source with weight
58 ``1.0``.
60 The following configuration produces two sources. Both inherit
61 ``split="98, 1, 1"``, while the second source overrides it::
63 {
64 "split": "98, 1, 1",
65 "sources": [
66 {"data_path": "/data/a.parquet", "weight": 0.7},
67 {
68 "data_path": "/data/b.parquet",
69 "weight": 0.3,
70 "split": "90, 5, 5",
71 },
72 ],
73 }
75 Returns:
76 Leaf source loaders and their corresponding unnormalized weights.
77 """
78 sources = self.data_config.get("sources")
79 if sources is None:
80 if self.data_path is None:
81 raise ValueError("Online Dataset requires data_path or data_config.sources")
82 source_loader = _OnlineLeafSourceLoader(self.data_path, self.data_config)
83 return [source_loader], [1.0]
85 if self.data_path is not None:
86 raise ValueError("Online Dataset cannot combine data_path with data_config.sources")
87 if isinstance(sources, (str, bytes)) or not isinstance(sources, Sequence) or not sources:
88 raise ValueError("Online data_config.sources must be a non-empty sequence")
90 common_config = dict(self.data_config)
91 common_config.pop("sources")
92 source_loaders = []
93 source_weights = []
94 for source in sources:
95 if not isinstance(source, Mapping):
96 raise ValueError("Each Online source must be a mapping with data_path and weight")
97 try:
98 source_path = source["data_path"]
99 source_weight = float(source["weight"])
100 except KeyError as error:
101 raise ValueError("Each Online source must define data_path and weight") from error
102 if source_weight <= 0:
103 raise ValueError("Online source weight must be positive")
105 source_config = dict(common_config)
106 for option_name, option_value in source.items():
107 if option_name not in ("data_path", "weight"):
108 source_config[option_name] = option_value
109 source_loaders.append(_OnlineLeafSourceLoader(source_path, source_config))
110 source_weights.append(source_weight)
111 return source_loaders, source_weights
113 def _interleave_sources(self, source_datasets: Sequence[Any], weights: Sequence[float]) -> Any:
114 """Interleave Mapping or Iterable sources through the Hugging Face interface."""
115 if len(source_datasets) == 1:
116 return source_datasets[0]
118 try:
119 # datasets remains optional for Indexed-only installations.
120 from datasets import interleave_datasets # pylint: disable=C0415
121 except ImportError as error:
122 raise ImportError("Online Dataset requires the optional 'datasets' package") from error
124 weight_sum = sum(weights)
125 probabilities = []
126 for weight in weights:
127 probability = weight / weight_sum
128 probabilities.append(probability)
130 interleaved_dataset = interleave_datasets(
131 list(source_datasets),
132 probabilities=probabilities,
133 seed=self._random_seed,
134 stopping_strategy=self._stopping_strategy,
135 )
136 return interleaved_dataset
139class _OnlineLeafSourceLoader:
140 """Load one local or Hub source with ``datasets.load_dataset``."""
142 def __init__(self, data_path: OnlineDataPath, data_config: Mapping[str, Any]) -> None:
143 """Store one leaf source description."""
144 self.data_path = data_path
145 self.data_config = data_config
147 def load(
148 self,
149 *,
150 streaming: bool,
151 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
152 ) -> Any:
153 """Load, restore configured splits, and filter one leaf source."""
154 try:
155 # datasets remains optional for Indexed-only installations.
156 from datasets import load_dataset # pylint: disable=C0415
157 except ImportError as error:
158 raise ImportError("Online Dataset requires the optional 'datasets' package") from error
160 # Examples: "Salesforce/wikitext", "/data/*.parquet", {"train": "/data/train.jsonl"}.
161 load_dataset_path, load_options, split_expressions = self._build_load_request(streaming=streaming)
162 ratio_split_config = self.data_config.get("split")
163 filter_before_ratio_split = ratio_split_config is not None and sample_filter is not None
164 selected_split: str | list[str] = "train"
165 if split_expressions is not None and not filter_before_ratio_split:
166 selected_splits = []
167 for split_expression in split_expressions:
168 if split_expression is not None:
169 selected_splits.append(split_expression)
170 selected_split = selected_splits
172 loaded_dataset = load_dataset(
173 load_dataset_path,
174 **load_options,
175 split=selected_split,
176 streaming=streaming,
177 cache_dir=self.data_config.get("cache_dir"),
178 )
180 if filter_before_ratio_split:
181 filtered_dataset = loaded_dataset.filter(sample_filter)
182 split_percentages = self._parse_split_percentages(ratio_split_config)
183 filtered_splits = self._split_filtered_dataset(filtered_dataset, split_percentages)
184 return filtered_splits
186 dataset_result = loaded_dataset
187 if split_expressions is not None:
188 # Example: ("train", None, "test") keeps the empty valid slot.
189 loaded_dataset_iterator = iter(loaded_dataset)
190 restored_splits = []
191 for split_expression in split_expressions:
192 if split_expression is None:
193 restored_splits.append(None)
194 continue
195 loaded_split = next(loaded_dataset_iterator)
196 restored_splits.append(loaded_split)
197 dataset_result = (restored_splits[0], restored_splits[1], restored_splits[2])
199 if sample_filter is None:
200 return dataset_result
202 # Example: (train, None, test) filters available splits and preserves the empty valid slot.
203 if isinstance(dataset_result, tuple):
204 filtered_splits = []
205 for split_dataset in dataset_result:
206 if split_dataset is None:
207 filtered_splits.append(None)
208 continue
209 filtered_split = split_dataset.filter(sample_filter)
210 filtered_splits.append(filtered_split)
211 filtered_dataset = (filtered_splits[0], filtered_splits[1], filtered_splits[2])
212 return filtered_dataset
214 filtered_dataset = dataset_result.filter(sample_filter)
215 return filtered_dataset
217 def _build_load_request(
218 self,
219 *,
220 streaming: bool,
221 ) -> tuple[str, dict[str, Any], tuple[str | None, str | None, str | None] | None]:
222 """Build the path, keyword arguments, and split selections for ``load_dataset``."""
223 split_config = self.data_config.get("split")
224 split_expressions = None
225 if split_config is not None and streaming:
226 raise ValueError("Online ratio split does not support streaming")
228 if split_config is not None and isinstance(self.data_path, Mapping):
229 raise ValueError("Online ratio split cannot be combined with pre-split paths")
231 if split_config is not None:
232 split_expressions = self._parse_split_ratios(split_config)
234 load_options: dict[str, Any] = {}
235 # Example: {"train": "/data/train.jsonl", "valid": "/data/valid.jsonl"}.
236 if isinstance(self.data_path, Mapping):
237 data_files, load_dataset_path, split_expressions = self._resolve_split_data_files(self.data_path)
238 load_options["data_files"] = data_files
239 resolved_request = (load_dataset_path, load_options, split_expressions)
240 return resolved_request
242 is_local_path = not isinstance(self.data_path, str)
243 if isinstance(self.data_path, str):
244 file_extension = os.path.splitext(self.data_path)[1].lower()
245 is_local_path = os.path.exists(self.data_path) or file_extension in ONLINE_FILE_FORMATS
247 # Examples: ["/data/a.jsonl", "/data/b.jsonl"], "/data/*.parquet".
248 if is_local_path:
249 data_files, load_dataset_path = self._collect_local_data_files(self.data_path)
250 load_options["data_files"] = data_files
251 resolved_request = (load_dataset_path, load_options, split_expressions)
252 return resolved_request
254 # Example: "Salesforce/wikitext" is a Hub dataset ID.
255 load_dataset_path = self.data_path
256 config_name = self.data_config.get("config_name")
257 if config_name is not None:
258 load_options["name"] = config_name
260 resolved_request = (load_dataset_path, load_options, split_expressions)
261 return resolved_request
263 @staticmethod
264 def _split_filtered_dataset(dataset: Any, percentages: Sequence[int]) -> tuple[Any, Any, Any]:
265 """Split one filtered Mapping Dataset using deterministic contiguous ranges."""
266 dataset_length = len(dataset)
267 train_end = int(round(percentages[0] * dataset_length / 100.0))
268 valid_end = int(round((percentages[0] + percentages[1]) * dataset_length / 100.0))
269 split_ranges = ((0, train_end), (train_end, valid_end), (valid_end, dataset_length))
271 split_datasets = []
272 for percentage, (split_start, split_end) in zip(percentages, split_ranges):
273 split_dataset = None
274 if percentage > 0:
275 split_dataset = dataset.select(range(split_start, split_end))
276 split_datasets.append(split_dataset)
277 dataset_splits = (split_datasets[0], split_datasets[1], split_datasets[2])
278 return dataset_splits
280 def _resolve_split_data_files(
281 self,
282 data_path: Mapping[str, OnlinePathSpec],
283 ) -> tuple[dict[str, list[str]], str, tuple[str | None, str | None, str | None]]:
284 """Resolve pre-split local paths into ``load_dataset`` arguments."""
285 if not data_path:
286 raise ValueError("Online split data_path mapping must not be empty")
288 split_data_files: dict[str, list[str]] = {}
289 loader_formats = set()
290 for split_name, split_path in data_path.items():
291 if split_name not in ONLINE_SPLIT_NAMES:
292 raise ValueError(
293 f"Unsupported Online data_path split {split_name!r}; "
294 f"expected one of {ONLINE_SPLIT_NAMES!r}"
295 )
297 data_files, loader_format = self._collect_local_data_files(split_path)
298 split_data_files[split_name] = data_files
299 loader_formats.add(loader_format)
301 if len(loader_formats) != 1:
302 raise ValueError("Online train/valid/test paths must use one common file format")
304 train_split = None
305 if "train" in split_data_files:
306 train_split = "train"
308 valid_split = None
309 if "valid" in split_data_files:
310 valid_split = "valid"
312 test_split = None
313 if "test" in split_data_files:
314 test_split = "test"
316 selected_splits = (train_split, valid_split, test_split)
317 loader_format = loader_formats.pop()
318 return split_data_files, loader_format, selected_splits
320 @staticmethod
321 def _collect_local_data_files(data_path: OnlinePathSpec) -> tuple[list[str], str]:
322 """Collect local files and select their ``load_dataset`` format."""
323 if isinstance(data_path, str):
324 configured_paths = [data_path]
325 else:
326 configured_paths = list(data_path)
328 data_files = []
329 for configured_path in configured_paths:
330 if not os.path.isdir(configured_path):
331 data_files.append(configured_path)
332 continue
334 for filename in sorted(os.listdir(configured_path)):
335 file_extension = os.path.splitext(filename)[1].lower()
336 if file_extension in ONLINE_FILE_FORMATS:
337 data_files.append(os.path.join(configured_path, filename))
339 if not data_files:
340 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
342 first_extension = os.path.splitext(data_files[0])[1].lower()
343 loader_format = ONLINE_FILE_FORMATS.get(first_extension)
344 for data_file in data_files[1:]:
345 file_extension = os.path.splitext(data_file)[1].lower()
346 if ONLINE_FILE_FORMATS.get(file_extension) != loader_format:
347 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
349 if loader_format is None:
350 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
352 return data_files, loader_format
354 @staticmethod
355 def _parse_split_percentages(split: Any) -> tuple[int, int, int]:
356 """Parse and validate train, validation, and test percentages."""
357 try:
358 percentages = []
359 for split_part in split.split(","):
360 percentages.append(int(split_part.strip()))
361 except (AttributeError, ValueError) as error:
362 raise ValueError("Online split must contain three integers such as '98, 1, 1'") from error
364 if len(percentages) != 3:
365 raise ValueError("Online split must contain three integers such as '98, 1, 1'")
367 if sum(percentages) != 100:
368 raise ValueError("Online split percentages must sum to 100")
369 if any(percentage < 0 for percentage in percentages):
370 raise ValueError("Online split percentages must be non-negative")
372 split_percentages = (percentages[0], percentages[1], percentages[2])
373 return split_percentages
375 @classmethod
376 def _parse_split_ratios(cls, split: Any) -> tuple[str | None, str | None, str | None]:
377 """Translate train/valid/test percentages into ``load_dataset`` split expressions."""
378 train_percentage, valid_percentage, test_percentage = cls._parse_split_percentages(split)
380 train_end = train_percentage
381 valid_end = train_percentage + valid_percentage
382 split_ranges = ((0, train_end), (train_end, valid_end), (valid_end, 100))
384 split_expressions: list[str | None] = []
385 for split_begin, split_end in split_ranges:
386 if split_begin == split_end:
387 split_expression = None
388 elif split_begin == 0 and split_end == 100:
389 split_expression = "train"
390 elif split_begin == 0:
391 split_expression = f"train[:{split_end}%]"
392 elif split_end == 100:
393 split_expression = f"train[{split_begin}%:]"
394 else:
395 split_expression = f"train[{split_begin}%:{split_end}%]"
396 split_expressions.append(split_expression)
398 online_split_expressions = (split_expressions[0], split_expressions[1], split_expressions[2])
399 return online_split_expressions
402__all__ = [
403 "OnlineDataPath",
404 "OnlinePathSpec",
405]