Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_helpers.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"""Local index-building helpers for indexed LLM Datasets.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Sequence 

20import importlib 

21import os 

22from pathlib import Path 

23import subprocess 

24import sys 

25import sysconfig 

26import time 

27 

28import numpy as np 

29 

30 

31try: 

32 from hyper_parallel.data.indexed._indexed_helpers_cpp import ( 

33 build_blending_indices as _build_blending_indices_cpp, 

34 build_sample_index_int32, 

35 build_sample_index_int64, 

36 ) 

37except ImportError as error: 

38 source_dir = Path(__file__).resolve().parent / "csrc" 

39 library_path = source_dir.parent / f"_indexed_helpers_cpp{sysconfig.get_config_var('EXT_SUFFIX')}" 

40 start_time = time.monotonic() 

41 

42 # Config resolution imports this module before the process group is initialized. 

43 if int(os.environ.get("LOCAL_RANK", "0")) == 0: 

44 print("> compiling dataset index builder ...", flush=True) 

45 subprocess.run( 

46 ["make", "-B", "-C", str(source_dir), f"PYTHON={sys.executable}", f"OUTPUT={library_path}.tmp"], 

47 check=True, 

48 ) 

49 # Other ranks must only see the library once compilation has completed. 

50 os.replace(f"{library_path}.tmp", library_path) 

51 print( 

52 f">>> done with dataset index builder. Compilation time: {time.monotonic() - start_time:.3f} seconds", 

53 flush=True, 

54 ) 

55 else: 

56 while not library_path.is_file(): 

57 if time.monotonic() - start_time > 300: 

58 raise ImportError("Timed out waiting for rank 0 to compile the Dataset C++ helpers.") from error 

59 time.sleep(0.1) 

60 

61 importlib.invalidate_caches() 

62 from hyper_parallel.data.indexed._indexed_helpers_cpp import ( 

63 build_blending_indices as _build_blending_indices_cpp, 

64 build_sample_index_int32, 

65 build_sample_index_int64, 

66 ) 

67 

68 

69def build_sample_index( 

70 sequence_lengths: np.ndarray, 

71 document_index: np.ndarray, 

72 sequence_length: int, 

73 num_epochs: int, 

74 num_tokens_per_epoch: int, 

75 drop_last_partial_sequence: bool, 

76 add_extra_token_to_sequence: bool, 

77) -> np.ndarray: 

78 """Build the 2-D sample index using the properly typed templated C++ helper. 

79 

80 Args: 

81 sequence_lengths (np.ndarray): The 1-D array of document lengths. 

82 

83 document_index (np.ndarray): The 1-D array of document indices. 

84 

85 sequence_length (int): The sequence length. 

86 

87 num_epochs (int): The number of epochs. 

88 

89 num_tokens_per_epoch (int): The number of tokens per epoch. 

90 

91 drop_last_partial_sequence (bool): Whether to omit the last partial sequence in the sample 

92 index should it exist. 

93 

94 add_extra_token_to_sequence (bool): Whether to build samples with sequence length 

95 ``sequence_length + 1``. 

96 

97 Returns: 

98 np.ndarray: The 2-D sample index. 

99 """ 

100 sample_index_max = max(document_index.shape[0], int(sequence_lengths.max())) 

101 sample_index_builder = ( 

102 build_sample_index_int32 if sample_index_max <= np.iinfo(np.int32).max else build_sample_index_int64 

103 ) 

104 sample_index = sample_index_builder( 

105 sequence_lengths, 

106 document_index, 

107 sequence_length, 

108 num_epochs, 

109 num_tokens_per_epoch, 

110 drop_last_partial_sequence, 

111 add_extra_token_to_sequence, 

112 ) 

113 return sample_index 

114 

115 

116def build_blending_indices( 

117 dataset_index: np.ndarray, 

118 dataset_sample_index: np.ndarray, 

119 weights: Sequence[float], 

120) -> None: 

121 """Build deterministic weighted blend indices. 

122 

123 Args: 

124 dataset_index: Output source-Dataset ID for each blended sample. 

125 dataset_sample_index: Output per-source sample ID for each blended sample. 

126 weights: Normalized source-Dataset weights. 

127 """ 

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

129 _build_blending_indices_cpp( 

130 dataset_index=dataset_index, 

131 dataset_sample_index=dataset_sample_index, 

132 weights=weight_array, 

133 )