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

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) -> Any: 

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

42 

43 Args: 

44 data_config: Online Mapping source options. 

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

46 train/valid/test path mapping. 

47 transform: Omni sample transform selected by the Trainer. 

48 training_config: Training plan providing the random seed. 

49 

50 Returns: 

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

52 

53 Raises: 

54 TypeError: If transform is not an OmniDataTransform. 

55 ValueError: If no Omni transform is configured. 

56 """ 

57 if transform is None: 

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

59 if not isinstance(transform, OmniDataTransform): 

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

61 

62 dataset_config = dict(data_config) 

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

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

65 

66 sample_filter = transform.is_valid_sample 

67 source_dataset = build_online_mapping_source( 

68 data_path=data_path, 

69 data_config=dataset_config, 

70 sample_filter=sample_filter, 

71 ) 

72 transformed_dataset = _OmniMappingDataset.apply(source_dataset, transform) 

73 return transformed_dataset 

74 

75 

76class _OmniMappingDataset(MappingTransformDataset): 

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

78 

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

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

81 self.data_transform = transform 

82 self._transform_strategy = _OmniTransformStrategy.from_transform(transform) 

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

84 

85 @classmethod 

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

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

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

89 transformed_splits = [] 

90 for split_dataset in source_dataset: 

91 transformed_dataset = None 

92 if split_dataset is not None: 

93 transformed_dataset = cls(split_dataset, transform) 

94 transformed_splits.append(transformed_dataset) 

95 

96 mapping_splits = ( 

97 transformed_splits[0], 

98 transformed_splits[1], 

99 transformed_splits[2], 

100 ) 

101 return mapping_splits 

102 

103 if source_dataset is None: 

104 return None 

105 

106 mapping_dataset = cls(source_dataset, transform) 

107 return mapping_dataset 

108 

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

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

111 encoded_sample = self._transform_strategy.encode_selected_sample(sample) 

112 return encoded_sample 

113 

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

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

116 encoded_batch = self.data_transform.encode_batch(batch) 

117 return encoded_batch 

118 

119 

120__all__ = ["build_online_omni_mapping_dataset"]