Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / omni / build_dataset.py: 87%

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"""Compose an Online source with the configured Omni transform lifecycle.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

22from hyper_parallel.data.constants import ONLINE_SPLIT_COUNT 

23from hyper_parallel.data.omni.omni_transform import ( 

24 OmniDataTransform, 

25 _OmniTransformStrategy, 

26) 

27from hyper_parallel.data.online import ( 

28 MappingTransformDataset, 

29 OnlineDataPath, 

30 build_online_mapping_source, 

31) 

32 

33 

34def build_online_omni_mapping_dataset( 

35 *, 

36 data_config: Mapping[str, Any], 

37 data_path: OnlineDataPath | None = None, 

38 transform: OmniDataTransform | None = None, 

39 training_config: Any = None, 

40 sampler_type: str = "single", 

41) -> Any: 

42 """Build an Online Mapping source and apply its Omni transform lazily. 

43 

44 Args: 

45 data_config: Online Mapping source options. 

46 data_path: Optional local source path, ordered paths, or pre-split 

47 train/valid/test path mapping. 

48 transform: Omni sample transform selected by the Trainer. 

49 training_config: Training plan providing the random seed. 

50 sampler_type: Mapping sampler policy controlling source-level shuffle. 

51 

52 Returns: 

53 A transformed Online Mapping Dataset or train-valid-test tuple. 

54 

55 Raises: 

56 TypeError: If transform is not an OmniDataTransform. 

57 ValueError: If no Omni transform is configured. 

58 """ 

59 if transform is None: 

60 raise ValueError("Online Omni Dataset requires a data_transform") 

61 if not isinstance(transform, OmniDataTransform): 

62 raise TypeError("Online Omni Dataset transform must be an OmniDataTransform") 

63 

64 dataset_config = dict(data_config) 

65 training_seed = getattr(training_config, "seed", None) 

66 dataset_config["random_seed"] = 42 if training_seed is None else int(training_seed) 

67 

68 sample_filter = transform.is_valid_sample 

69 source_dataset = build_online_mapping_source( 

70 data_path=data_path, 

71 data_config=dataset_config, 

72 sample_filter=sample_filter, 

73 sampler_type=sampler_type, 

74 ) 

75 transformed_dataset = _OmniMappingDataset.apply(source_dataset, transform) 

76 return transformed_dataset 

77 

78 

79class _OmniMappingDataset(MappingTransformDataset): 

80 """Retain the Omni encoding lifecycle around one Mapping source.""" 

81 

82 def __init__(self, source_dataset: Any, transform: OmniDataTransform) -> None: 

83 """Build the pre-selection strategy and retain post-selection hooks.""" 

84 self.data_transform = transform 

85 self._transform_strategy = _OmniTransformStrategy.from_transform(transform) 

86 super().__init__(source_dataset, self._transform_strategy.prepare_samples) 

87 

88 @classmethod 

89 def apply(cls, source_dataset: Any, transform: OmniDataTransform) -> Any: 

90 """Wrap one Mapping source or each available train-valid-test split.""" 

91 if isinstance(source_dataset, tuple) and len(source_dataset) == ONLINE_SPLIT_COUNT: 

92 transformed_splits = [] 

93 for split_dataset in source_dataset: 

94 transformed_dataset = None 

95 if split_dataset is not None: 

96 transformed_dataset = cls(split_dataset, transform) 

97 transformed_splits.append(transformed_dataset) 

98 

99 mapping_splits = ( 

100 transformed_splits[0], 

101 transformed_splits[1], 

102 transformed_splits[2], 

103 ) 

104 return mapping_splits 

105 

106 if source_dataset is None: 

107 return None 

108 

109 mapping_dataset = cls(source_dataset, transform) 

110 return mapping_dataset 

111 

112 def encode_selected_sample(self, sample: Any) -> Any: 

113 """Finish encoding one sample after packing selection.""" 

114 encoded_sample = self._transform_strategy.encode_selected_sample(sample) 

115 return encoded_sample 

116 

117 def encode_batch(self, batch: Any) -> Any: 

118 """Apply the transform's final batch encoding hook.""" 

119 encoded_batch = self.data_transform.encode_batch(batch) 

120 return encoded_batch 

121 

122 

123__all__ = ["build_online_omni_mapping_dataset"]