Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / online / source.py: 80%
280 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"""Shared Online source loading and multi-source construction policy."""
17from __future__ import annotations
19import glob
20import os
21from collections.abc import Callable, Mapping, Sequence
22from typing import Any, TypeAlias
24from hyper_parallel.data.constants import (
25 ONLINE_FILE_FORMATS, ONLINE_SOURCE_PATH_KEY, ONLINE_SPLIT_NAMES, ONLINE_STOPPING_STRATEGIES,
26)
28OnlinePathSpec: TypeAlias = str | Sequence[str]
29OnlineDataPath: TypeAlias = OnlinePathSpec | Mapping[str, OnlinePathSpec]
32class _OnlineSourceBuilder:
33 """Share source parsing and interleave policy across Online adapters."""
35 def __init__(self, data_path: OnlineDataPath | None, data_config: Mapping[str, Any]) -> None:
36 """Store the source and composition configuration."""
37 self.data_path = data_path
38 self.data_config = data_config
40 @property
41 def _random_seed(self) -> int:
42 """Return the common source shuffle and blend seed."""
43 return int(self.data_config.get("random_seed", 42))
45 @property
46 def _stopping_strategy(self) -> str:
47 """Return the multi-source exhaustion policy."""
48 stopping_strategy = str(self.data_config.get("stopping_strategy", "first_exhausted"))
49 if stopping_strategy not in ONLINE_STOPPING_STRATEGIES:
50 raise ValueError(
51 f"Unsupported Online stopping_strategy {stopping_strategy!r}; "
52 f"expected one of {ONLINE_STOPPING_STRATEGIES!r}"
53 )
54 return stopping_strategy
56 def _create_source_loaders(self) -> tuple[list[_OnlineLeafSourceLoader], list[float]]:
57 """Create one leaf source loader per configured source.
59 Example:
60 ``data_path="/data/train.jsonl"`` produces one source with weight
61 ``1.0``.
63 The following configuration produces two sources. Both inherit
64 ``split="98, 1, 1"``, while the second source overrides it::
66 {
67 "split": "98, 1, 1",
68 "sources": [
69 {"data_path": "/data/a.parquet", "weight": 0.7},
70 {
71 "data_path": "/data/b.parquet",
72 "weight": 0.3,
73 "split": "90, 5, 5",
74 },
75 ],
76 }
78 Returns:
79 Leaf source loaders and their corresponding unnormalized weights.
80 """
81 sources = self.data_config.get("sources")
82 if sources is None:
83 if self.data_path is None:
84 raise ValueError("Online Dataset requires data_path or data_config.sources")
85 source_loader = _OnlineLeafSourceLoader(self.data_path, self.data_config)
86 return [source_loader], [1.0]
88 if self.data_path is not None:
89 raise ValueError("Online Dataset cannot combine data_path with data_config.sources")
90 if isinstance(sources, (str, bytes)) or not isinstance(sources, Sequence) or not sources:
91 raise ValueError("Online data_config.sources must be a non-empty sequence")
93 common_config = dict(self.data_config)
94 common_config.pop("sources")
95 source_loaders = []
96 source_weights = []
97 for source in sources:
98 if not isinstance(source, Mapping):
99 raise ValueError("Each Online source must be a mapping with data_path and weight")
100 try:
101 source_path = source["data_path"]
102 source_weight = float(source["weight"])
103 except KeyError as error:
104 raise ValueError("Each Online source must define data_path and weight") from error
105 if source_weight <= 0:
106 raise ValueError("Online source weight must be positive")
108 source_config = dict(common_config)
109 for option_name, option_value in source.items():
110 if option_name not in ("data_path", "weight"):
111 source_config[option_name] = option_value
112 source_loaders.append(_OnlineLeafSourceLoader(source_path, source_config))
113 source_weights.append(source_weight)
114 return source_loaders, source_weights
116 def _interleave_sources(self, source_datasets: Sequence[Any], weights: Sequence[float]) -> Any:
117 """Interleave Mapping or Iterable sources through the Hugging Face interface."""
118 if len(source_datasets) == 1:
119 return source_datasets[0]
121 try:
122 # datasets remains optional for Indexed-only installations.
123 from datasets import interleave_datasets # pylint: disable=C0415
124 except ImportError as error:
125 raise ImportError("Online Dataset requires the optional 'datasets' package") from error
127 weight_sum = sum(weights)
128 probabilities = []
129 for weight in weights:
130 probability = weight / weight_sum
131 probabilities.append(probability)
133 interleaved_dataset = interleave_datasets(
134 list(source_datasets),
135 probabilities=probabilities,
136 seed=self._random_seed,
137 stopping_strategy=self._stopping_strategy,
138 )
139 return interleaved_dataset
142class _OnlineLeafSourceLoader:
143 """Load one local or Hub source with ``datasets.load_dataset``."""
145 def __init__(self, data_path: OnlineDataPath, data_config: Mapping[str, Any]) -> None:
146 """Store one leaf source description."""
147 self.data_path = data_path
148 self.data_config = data_config
150 def load(
151 self,
152 *,
153 streaming: bool,
154 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
155 ) -> Any:
156 """Load, restore configured splits, and filter one leaf source."""
157 try:
158 # datasets remains optional for Indexed-only installations.
159 from datasets import load_dataset # pylint: disable=C0415
160 except ImportError as error:
161 raise ImportError("Online Dataset requires the optional 'datasets' package") from error
163 # Examples: "Salesforce/wikitext", "/data/*.parquet", {"train": "/data/train.jsonl"}.
164 load_dataset_path, load_options, split_expressions = self._build_load_request(streaming=streaming)
165 ratio_split_config = self.data_config.get("split")
166 filter_before_ratio_split = ratio_split_config is not None and sample_filter is not None
167 selected_split: str | list[str] = "train"
168 if split_expressions is not None and not filter_before_ratio_split:
169 selected_splits = []
170 for split_expression in split_expressions:
171 if split_expression is not None:
172 selected_splits.append(split_expression)
173 selected_split = selected_splits
175 if load_dataset_path in ("json", "parquet") and "data_files" in load_options and not streaming:
176 loaded_dataset = self._load_local_files(load_dataset_path, load_options["data_files"], selected_split)
177 else:
178 loaded_dataset = load_dataset(
179 load_dataset_path,
180 **load_options,
181 split=selected_split,
182 streaming=streaming,
183 cache_dir=self.data_config.get("cache_dir"),
184 )
186 if filter_before_ratio_split:
187 filtered_dataset = loaded_dataset.filter(sample_filter)
188 split_percentages = self.parse_split_percentages(ratio_split_config)
189 filtered_splits = self._split_filtered_dataset(filtered_dataset, split_percentages)
190 return filtered_splits
192 dataset_result = loaded_dataset
193 if split_expressions is not None:
194 # Example: ("train", None, "test") keeps the empty valid slot.
195 loaded_dataset_iterator = iter(loaded_dataset)
196 restored_splits = []
197 for split_expression in split_expressions:
198 if split_expression is None:
199 restored_splits.append(None)
200 continue
201 loaded_split = next(loaded_dataset_iterator)
202 restored_splits.append(loaded_split)
203 dataset_result = (restored_splits[0], restored_splits[1], restored_splits[2])
205 if sample_filter is None:
206 return dataset_result
208 # Example: (train, None, test) filters available splits and preserves the empty valid slot.
209 if isinstance(dataset_result, tuple):
210 filtered_splits = []
211 for split_dataset in dataset_result:
212 if split_dataset is None:
213 filtered_splits.append(None)
214 continue
215 filtered_split = split_dataset.filter(sample_filter)
216 filtered_splits.append(filtered_split)
217 filtered_dataset = (filtered_splits[0], filtered_splits[1], filtered_splits[2])
218 return filtered_dataset
220 filtered_dataset = dataset_result.filter(sample_filter)
221 return filtered_dataset
223 def _load_local_files(self, file_format: str, data_files: Any, selected_split: str | list[str]) -> Any:
224 """Preserve per-file media roots through filtering and split selection."""
225 # datasets remains optional for Indexed-only installations.
226 from datasets import ReadInstruction, concatenate_datasets, load_dataset # pylint: disable=C0415
228 files_by_split = data_files if isinstance(data_files, Mapping) else {"train": data_files}
229 if any("://" in path for paths in files_by_split.values() for path in paths):
230 return load_dataset(
231 file_format, data_files=data_files, split=selected_split,
232 cache_dir=self.data_config.get("cache_dir"),
233 )
234 datasets_by_split = {}
235 for split_name, paths in files_by_split.items():
236 source_datasets = []
237 for path_pattern in paths:
238 matching_paths = sorted(glob.glob(path_pattern)) or [path_pattern]
239 for source_path in matching_paths:
240 source = load_dataset(
241 file_format, data_files=source_path, split="train",
242 cache_dir=self.data_config.get("cache_dir"),
243 )
244 if ONLINE_SOURCE_PATH_KEY in source.column_names:
245 source = source.remove_columns(ONLINE_SOURCE_PATH_KEY)
246 source = source.add_column(
247 ONLINE_SOURCE_PATH_KEY, [os.path.abspath(source_path)] * len(source),
248 )
249 source_datasets.append(source)
250 datasets_by_split[split_name] = concatenate_datasets(source_datasets)
252 split_lengths = {name: len(dataset) for name, dataset in datasets_by_split.items()}
253 selections = [selected_split] if isinstance(selected_split, str) else selected_split
254 selected_datasets = []
255 for selection in selections:
256 instructions = ReadInstruction.from_spec(selection).to_absolute(split_lengths)
257 parts = []
258 for instruction in instructions:
259 source = datasets_by_split[instruction.splitname]
260 start = instruction.from_ or 0
261 end = len(source) if instruction.to is None else instruction.to
262 parts.append(source.select(range(start, end)))
263 selected_datasets.append(concatenate_datasets(parts))
264 return selected_datasets[0] if isinstance(selected_split, str) else selected_datasets
266 def _build_load_request(
267 self,
268 *,
269 streaming: bool,
270 ) -> tuple[str, dict[str, Any], tuple[str | None, str | None, str | None] | None]:
271 """Build the path, keyword arguments, and split selections for ``load_dataset``."""
272 split_config = self.data_config.get("split")
273 split_expressions = None
274 if split_config is not None and streaming:
275 raise ValueError("Online ratio split does not support streaming")
277 if split_config is not None and isinstance(self.data_path, Mapping):
278 raise ValueError("Online ratio split cannot be combined with pre-split paths")
280 if split_config is not None:
281 split_expressions = self._parse_split_ratios(split_config)
283 load_options: dict[str, Any] = {}
284 # Example: {"train": "/data/train.jsonl", "valid": "/data/valid.jsonl"}.
285 if isinstance(self.data_path, Mapping):
286 data_files, load_dataset_path, split_expressions = self._resolve_split_data_files(self.data_path)
287 load_options["data_files"] = data_files
288 resolved_request = (load_dataset_path, load_options, split_expressions)
289 return resolved_request
291 is_local_path = not isinstance(self.data_path, str)
292 if isinstance(self.data_path, str):
293 file_extension = os.path.splitext(self.data_path)[1].lower()
294 is_local_path = os.path.exists(self.data_path) or file_extension in ONLINE_FILE_FORMATS
296 # Examples: ["/data/a.jsonl", "/data/b.jsonl"], "/data/*.parquet".
297 if is_local_path:
298 data_files, load_dataset_path = self._collect_local_data_files(self.data_path)
299 load_options["data_files"] = data_files
300 resolved_request = (load_dataset_path, load_options, split_expressions)
301 return resolved_request
303 # Example: "Salesforce/wikitext" is a Hub dataset ID.
304 load_dataset_path = self.data_path
305 config_name = self.data_config.get("config_name")
306 if config_name is not None:
307 load_options["name"] = config_name
309 resolved_request = (load_dataset_path, load_options, split_expressions)
310 return resolved_request
312 @staticmethod
313 def _split_filtered_dataset(dataset: Any, percentages: Sequence[int]) -> tuple[Any, Any, Any]:
314 """Split one filtered Mapping Dataset using deterministic contiguous ranges."""
315 dataset_length = len(dataset)
316 train_end = int(round(percentages[0] * dataset_length / 100.0))
317 valid_end = int(round((percentages[0] + percentages[1]) * dataset_length / 100.0))
318 split_ranges = ((0, train_end), (train_end, valid_end), (valid_end, dataset_length))
320 split_datasets = []
321 for percentage, (split_start, split_end) in zip(percentages, split_ranges):
322 split_dataset = None
323 if percentage > 0:
324 split_dataset = dataset.select(range(split_start, split_end))
325 split_datasets.append(split_dataset)
326 dataset_splits = (split_datasets[0], split_datasets[1], split_datasets[2])
327 return dataset_splits
329 def _resolve_split_data_files(
330 self,
331 data_path: Mapping[str, OnlinePathSpec],
332 ) -> tuple[dict[str, list[str]], str, tuple[str | None, str | None, str | None]]:
333 """Resolve pre-split local paths into ``load_dataset`` arguments."""
334 if not data_path:
335 raise ValueError("Online split data_path mapping must not be empty")
337 split_data_files: dict[str, list[str]] = {}
338 loader_formats = set()
339 for split_name, split_path in data_path.items():
340 if split_name not in ONLINE_SPLIT_NAMES:
341 raise ValueError(
342 f"Unsupported Online data_path split {split_name!r}; "
343 f"expected one of {ONLINE_SPLIT_NAMES!r}"
344 )
346 data_files, loader_format = self._collect_local_data_files(split_path)
347 split_data_files[split_name] = data_files
348 loader_formats.add(loader_format)
350 if len(loader_formats) != 1:
351 raise ValueError("Online train/valid/test paths must use one common file format")
353 train_split = None
354 if "train" in split_data_files:
355 train_split = "train"
357 valid_split = None
358 if "valid" in split_data_files:
359 valid_split = "valid"
361 test_split = None
362 if "test" in split_data_files:
363 test_split = "test"
365 selected_splits = (train_split, valid_split, test_split)
366 loader_format = loader_formats.pop()
367 return split_data_files, loader_format, selected_splits
369 @staticmethod
370 def _collect_local_data_files(data_path: OnlinePathSpec) -> tuple[list[str], str]:
371 """Collect local files and select their ``load_dataset`` format."""
372 if isinstance(data_path, str):
373 configured_paths = [data_path]
374 else:
375 configured_paths = list(data_path)
377 data_files = []
378 for configured_path in configured_paths:
379 if not os.path.isdir(configured_path):
380 data_files.append(configured_path)
381 continue
383 for filename in sorted(os.listdir(configured_path)):
384 file_extension = os.path.splitext(filename)[1].lower()
385 if file_extension in ONLINE_FILE_FORMATS:
386 data_files.append(os.path.join(configured_path, filename))
388 if not data_files:
389 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
391 first_extension = os.path.splitext(data_files[0])[1].lower()
392 loader_format = ONLINE_FILE_FORMATS.get(first_extension)
393 for data_file in data_files[1:]:
394 file_extension = os.path.splitext(data_file)[1].lower()
395 if ONLINE_FILE_FORMATS.get(file_extension) != loader_format:
396 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
398 if loader_format is None:
399 raise ValueError("Online files must use one format: JSON/JSONL/Parquet/CSV/Arrow")
401 return data_files, loader_format
403 @staticmethod
404 def parse_split_percentages(split: Any) -> tuple[int, int, int]:
405 """Parse and validate train, validation, and test percentages."""
406 try:
407 percentages = []
408 for split_part in split.split(","):
409 percentages.append(int(split_part.strip()))
410 except (AttributeError, ValueError) as error:
411 raise ValueError("Online split must contain three integers such as '98, 1, 1'") from error
413 if len(percentages) != 3:
414 raise ValueError("Online split must contain three integers such as '98, 1, 1'")
416 if sum(percentages) != 100:
417 raise ValueError("Online split percentages must sum to 100")
418 if any(percentage < 0 for percentage in percentages):
419 raise ValueError("Online split percentages must be non-negative")
421 split_percentages = (percentages[0], percentages[1], percentages[2])
422 return split_percentages
424 @classmethod
425 def _parse_split_ratios(cls, split: Any) -> tuple[str | None, str | None, str | None]:
426 """Translate train/valid/test percentages into ``load_dataset`` split expressions."""
427 train_percentage, valid_percentage, _ = cls.parse_split_percentages(split)
429 train_end = train_percentage
430 valid_end = train_percentage + valid_percentage
431 split_ranges = ((0, train_end), (train_end, valid_end), (valid_end, 100))
433 split_expressions: list[str | None] = []
434 for split_begin, split_end in split_ranges:
435 if split_begin == split_end:
436 split_expression = None
437 elif split_begin == 0 and split_end == 100:
438 split_expression = "train"
439 elif split_begin == 0:
440 split_expression = f"train[:{split_end}%]"
441 elif split_end == 100:
442 split_expression = f"train[{split_begin}%:]"
443 else:
444 split_expression = f"train[{split_begin}%:{split_end}%]"
445 split_expressions.append(split_expression)
447 online_split_expressions = (split_expressions[0], split_expressions[1], split_expressions[2])
448 return online_split_expressions
451__all__ = [
452 "OnlineDataPath",
453 "OnlinePathSpec",
454]