Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_blended_dataset.py: 0%

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"""Standard deterministic weighted Dataset blending.""" 

16 

17from __future__ import annotations 

18 

19import hashlib 

20import json 

21import os 

22from collections import OrderedDict 

23from collections.abc import Mapping, Sequence 

24from typing import Any 

25 

26import numpy as np 

27 

28from hyper_parallel.data.dataset_logging import get_dataset_logger 

29from hyper_parallel.data.indexed.indexed_data_config import GPTDatasetConfig 

30from hyper_parallel.data.indexed.indexed_helpers import build_blending_indices 

31 

32logger = get_dataset_logger(__name__) 

33 

34 

35class BlendedDataset: 

36 """Expose several Datasets through deterministic weighted sample indices.""" 

37 

38 def __init__( 

39 self, 

40 datasets: Sequence[Any], 

41 weights: Sequence[float], 

42 size: int, 

43 config: GPTDatasetConfig, 

44 ) -> None: 

45 """Validate the inputs and create weighted sample indices.""" 

46 if not datasets or len(datasets) != len(weights): 

47 raise ValueError("datasets and weights must be non-empty and have the same length") 

48 if len(datasets) >= np.iinfo(np.int16).max: 

49 raise ValueError("number of datasets must be less than 32767") 

50 if not all(isinstance(dataset, type(datasets[0])) for dataset in datasets): 

51 raise ValueError("all datasets must be of the same type") 

52 self.datasets = list(datasets) 

53 self.weights = _normalize_weights(weights) 

54 self.size = size 

55 self.config = config 

56 self.unique_identifiers = OrderedDict() 

57 self.unique_identifiers["class"] = type(self).__name__ 

58 self.unique_identifiers["datasets"] = self._collect_dataset_identifiers() 

59 self.unique_identifiers["weights"] = self.weights 

60 self.unique_identifiers["size"] = self.size 

61 self.unique_description = json.dumps( 

62 self.unique_identifiers, 

63 indent=4, 

64 default=lambda value: value.unique_identifiers, 

65 ) 

66 self.unique_description_hash = hashlib.md5( 

67 self.unique_description.encode("utf-8") 

68 ).hexdigest() 

69 self.dataset_index, self.dataset_sample_index = self._build_indices() 

70 

71 def _collect_dataset_identifiers(self) -> list[Any]: 

72 """Collect component Dataset cache identities.""" 

73 identifiers = [dataset.unique_identifiers for dataset in self.datasets] 

74 return identifiers 

75 

76 def __len__(self) -> int: 

77 """Return the requested blended sample count.""" 

78 return self.size 

79 

80 def __getitem__(self, index: int) -> Mapping[str, Any]: 

81 """Read one sample from the Dataset selected by the blend indices.""" 

82 dataset_id = int(self.dataset_index[index]) 

83 sample_id = int(self.dataset_sample_index[index]) 

84 sample = {"dataset_id": dataset_id, **self.datasets[dataset_id][sample_id]} 

85 return sample 

86 

87 def _build_indices(self) -> tuple[np.ndarray, np.ndarray]: 

88 """Load or build the top-level blend index cache.""" 

89 cache_directory = self.config.path_to_cache 

90 if cache_directory is None: 

91 logger.debug( 

92 "Building blended indices in memory: size=%d, sources=%d, weights are %s", 

93 self.size, len(self.weights), self.weights[:3], 

94 ) 

95 indices = self._build_indices_in_memory() 

96 return indices 

97 

98 cache_prefix = os.path.join( 

99 cache_directory, 

100 f"{self.unique_description_hash}-{type(self).__name__}", 

101 ) 

102 description_path = f"{cache_prefix}-description.txt" 

103 dataset_index_path = f"{cache_prefix}-dataset_index.npy" 

104 sample_index_path = f"{cache_prefix}-dataset_sample_index.npy" 

105 cache_paths = (description_path, dataset_index_path, sample_index_path) 

106 if all(os.path.isfile(path) for path in cache_paths): 

107 logger.debug("Loading blended index cache: prefix=%s", cache_prefix) 

108 dataset_index = np.load(dataset_index_path, allow_pickle=True, mmap_mode="r") 

109 sample_index = np.load(sample_index_path, allow_pickle=True, mmap_mode="r") 

110 return dataset_index, sample_index 

111 

112 dataset_index, sample_index = self._build_indices_in_memory() 

113 os.makedirs(cache_directory, exist_ok=True) 

114 with open(description_path, "w", encoding="utf-8") as description_file: 

115 description_file.write(self.unique_description) 

116 np.save(dataset_index_path, dataset_index, allow_pickle=True) 

117 np.save(sample_index_path, sample_index, allow_pickle=True) 

118 logger.debug("Saved blended index cache: prefix=%s", cache_prefix) 

119 return dataset_index, sample_index 

120 

121 def _build_indices_in_memory(self) -> tuple[np.ndarray, np.ndarray]: 

122 """Build a deterministic weighted source schedule.""" 

123 dataset_index = np.empty(self.size, dtype=np.int16) 

124 sample_index = np.empty(self.size, dtype=np.int64) 

125 build_blending_indices( 

126 dataset_index=dataset_index, 

127 dataset_sample_index=sample_index, 

128 weights=self.weights, 

129 ) 

130 requested_counts = np.bincount(dataset_index, minlength=len(self.datasets)) 

131 for dataset_id, requested_count in enumerate(requested_counts): 

132 dataset_size = len(self.datasets[dataset_id]) 

133 if requested_count > dataset_size: 

134 raise ValueError( 

135 f"Dataset {dataset_id} has only {dataset_size} samples, " 

136 f"but the blend requested {requested_count} samples" 

137 ) 

138 return dataset_index, sample_index 

139 

140 

141def _normalize_weights(weights: Sequence[float]) -> list[float]: 

142 """Validate and normalize positive blend weights.""" 

143 if not weights or any(weight <= 0.0 for weight in weights): 

144 raise ValueError("Dataset blend weights must be positive") 

145 weight_array = np.asarray(weights, dtype=np.float64) 

146 normalized_weights = (weight_array / np.sum(weight_array)).tolist() 

147 return normalized_weights