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

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"""Read indexed token files and prepared Arrow records.""" 

16 

17from __future__ import annotations 

18 

19import os 

20import struct 

21import time 

22from collections.abc import Mapping 

23from functools import lru_cache 

24from pathlib import Path 

25from typing import Any, ClassVar 

26 

27import numpy as np 

28 

29from hyper_parallel.data.constants import INDEX_DTYPE_CODES, INDEX_HEADER, INDEX_VERSION 

30from hyper_parallel.data.dataset_logging import get_dataset_logger 

31 

32logger = get_dataset_logger(__name__) 

33 

34 

35class _IndexReader: 

36 """Memory-map the compact sequence metadata stored in an index file.""" 

37 

38 def __init__(self, idx_path: str) -> None: 

39 """Parse and memory-map one index metadata file.""" 

40 # Read the fixed-size header, which describes how the following index arrays are decoded. 

41 with open(idx_path, "rb") as stream: 

42 header = stream.read(len(INDEX_HEADER)) 

43 if header != INDEX_HEADER: 

44 raise ValueError(f"Invalid indexed Dataset header in {idx_path!r}") 

45 

46 version = struct.unpack("<Q", stream.read(8))[0] 

47 if version != INDEX_VERSION: 

48 raise ValueError(f"Unsupported indexed Dataset version {version}; expected {INDEX_VERSION}") 

49 

50 dtype_code = struct.unpack("<B", stream.read(1))[0] 

51 self.dtype = np.dtype(INDEX_DTYPE_CODES[dtype_code]) 

52 self.dtype_size = self.dtype.itemsize 

53 

54 self.sequence_count = struct.unpack("<Q", stream.read(8))[0] 

55 self.document_count = struct.unpack("<Q", stream.read(8))[0] 

56 offset = stream.tell() 

57 

58 # Map the three arrays stored after the header: token count per sequence, byte offset in .bin per sequence, 

59 # and sequence boundaries per document. The first two arrays have sequence_count entries, while the last 

60 # has document_count entries. 

61 self.bin_buffer_mmap = np.memmap(idx_path, mode="r", order="C") 

62 self.bin_buffer = memoryview(self.bin_buffer_mmap) 

63 self.sequence_lengths = np.frombuffer( 

64 self.bin_buffer, dtype=np.int32, count=self.sequence_count, offset=offset, 

65 ) 

66 

67 self.sequence_pointers = np.frombuffer( 

68 self.bin_buffer, dtype=np.int64, count=self.sequence_count, offset=offset + self.sequence_lengths.nbytes, 

69 ) 

70 

71 self.document_indices = np.frombuffer( 

72 self.bin_buffer, dtype=np.int64, count=self.document_count, 

73 offset=offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes, 

74 ) 

75 

76 # Each document is a half-open sequence range; the final boundary must equal the sequence count. 

77 if self.document_indices.size == 0 or self.document_indices[-1] != self.sequence_count: 

78 raise ValueError("Indexed Dataset document boundaries do not match its sequence count") 

79 

80 def __del__(self) -> None: 

81 """Close the index metadata mmap when it is no longer referenced.""" 

82 bin_buffer_mmap = getattr(self, "bin_buffer_mmap", None) 

83 mmap_handle = getattr(bin_buffer_mmap, "_mmap", None) 

84 if mmap_handle is not None: 

85 mmap_handle.close() 

86 

87 def __len__(self) -> int: 

88 """Return the number of indexed sequences.""" 

89 return self.sequence_count 

90 

91 @lru_cache(maxsize=8) 

92 def __getitem__(self, index: int | np.integer) -> tuple[np.int64, np.int32]: 

93 """Return the byte pointer and token length of one sequence.""" 

94 return self.sequence_pointers[index], self.sequence_lengths[index] 

95 

96 

97class IndexedDataReader: 

98 """Read individual token sequences from a paired ``.idx/.bin`` Dataset. 

99 

100 Args: 

101 path_prefix: Dataset path without the ``.idx`` or ``.bin`` suffix. 

102 mmap: Whether to memory-map the token payload. 

103 reuse_index: Whether to reuse the first loaded index metadata. This 

104 matches corpora whose shards share an identical index layout. 

105 """ 

106 

107 cached_index_reader: ClassVar[_IndexReader | None] = None 

108 

109 def __init__(self, path_prefix: str, mmap: bool = True, reuse_index: bool = False) -> None: 

110 """Open the index metadata and token payload files.""" 

111 self.path_prefix = "" 

112 self.mmap = False 

113 self.reuse_index = reuse_index 

114 self.index: _IndexReader | None = None 

115 self.bin_buffer_mmap: np.memmap | None = None 

116 self.bin_buffer: memoryview | None = None 

117 self.initialize(path_prefix, mmap) 

118 

119 def initialize(self, path_prefix: str, mmap: bool) -> None: 

120 """Open the index metadata and token payload. 

121 

122 Args: 

123 path_prefix: Dataset path without the ``.idx`` or ``.bin`` suffix. 

124 mmap: Whether to memory-map the token payload. 

125 """ 

126 start_time = time.time() 

127 index_path = path_prefix + ".idx" 

128 data_path = path_prefix + ".bin" 

129 if not os.path.isfile(index_path) or not os.path.isfile(data_path): 

130 raise FileNotFoundError(f"Expected indexed Dataset files {index_path!r} and {data_path!r}") 

131 

132 self.path_prefix = path_prefix 

133 self.mmap = mmap 

134 reused_cached_index = self.reuse_index and IndexedDataReader.cached_index_reader is not None 

135 if self.index is None: 

136 if reused_cached_index: 

137 self.index = IndexedDataReader.cached_index_reader 

138 else: 

139 self.index = _IndexReader(index_path) 

140 if self.reuse_index: 

141 IndexedDataReader.cached_index_reader = self.index 

142 

143 self.bin_buffer_mmap = None 

144 self.bin_buffer = None 

145 if mmap: 

146 self.bin_buffer_mmap = np.memmap(data_path, mode="r", order="C") 

147 self.bin_buffer = memoryview(self.bin_buffer_mmap) 

148 logger.debug( 

149 "Opened indexed Dataset: sequences=%d, documents=%d, dtype=%s, mmap=%s, prefix=%s, " 

150 "reused_index=%s, elapsed=%.4f seconds", 

151 len(self.index), self.index.document_indices.size - 1, self.index.dtype, mmap, path_prefix, 

152 reused_cached_index, time.time() - start_time, 

153 ) 

154 

155 def __getstate__(self) -> tuple[str, bool, bool]: 

156 """Serialize construction inputs instead of open mmap resources.""" 

157 return self.path_prefix, self.mmap, self.reuse_index 

158 

159 def __setstate__(self, state: tuple[str, bool, bool]) -> None: 

160 """Reopen index and data resources inside a DataLoader worker.""" 

161 path_prefix, mmap, reuse_index = state 

162 self.__init__(path_prefix, mmap=mmap, reuse_index=reuse_index) 

163 

164 def __del__(self) -> None: 

165 """Close the token mmap when this reader is released.""" 

166 bin_buffer_mmap = getattr(self, "bin_buffer_mmap", None) 

167 mmap_handle = getattr(bin_buffer_mmap, "_mmap", None) 

168 if mmap_handle is not None: 

169 mmap_handle.close() 

170 

171 def __len__(self) -> int: 

172 """Return the number of indexed token sequences.""" 

173 return len(self._require_index()) 

174 

175 def _getitem_mmap(self, index: int | np.integer | slice) -> np.ndarray | list[np.ndarray]: 

176 """Return one sequence or a contiguous sequence slice from the mapped token payload.""" 

177 index_reader = self._require_index() 

178 

179 if isinstance(index, (int, np.integer)): 

180 sequence_pointer, sequence_length = index_reader[index] 

181 sequence = np.frombuffer( 

182 self.bin_buffer, dtype=index_reader.dtype, count=sequence_length, offset=sequence_pointer 

183 ) 

184 return sequence 

185 

186 if isinstance(index, slice): 

187 start, stop, step = index.indices(len(self)) 

188 if step != 1: 

189 raise ValueError("Slices into indexed Dataset must be contiguous") 

190 if start == stop: 

191 return [] 

192 

193 sequence_lengths = index_reader.sequence_lengths[index] 

194 sequence_offsets = np.cumsum(sequence_lengths) 

195 sequences = np.frombuffer( 

196 self.bin_buffer, dtype=index_reader.dtype, count=int(sequence_lengths.sum()), 

197 offset=int(index_reader.sequence_pointers[start]), 

198 ) 

199 sequences = np.split(sequences, sequence_offsets[:-1]) 

200 return sequences 

201 

202 raise TypeError(f"Indexed Dataset indices must be integers or slices, got {type(index).__name__}") 

203 

204 def _getitem_file(self, index: int | np.integer | slice) -> np.ndarray: 

205 """Return one sequence through direct file reading when mmap is disabled.""" 

206 if isinstance(index, slice): 

207 raise NotImplementedError("Slicing is not implemented when mmap is disabled") 

208 

209 if not isinstance(index, (int, np.integer)): 

210 raise TypeError(f"Indexed Dataset indices must be integers or slices, got {type(index).__name__}") 

211 

212 return self.get(int(index)) 

213 

214 def __getitem__(self, index: int | np.integer | slice) -> np.ndarray | list[np.ndarray]: 

215 """Return one sequence, or a contiguous slice when mmap is enabled.""" 

216 if self.bin_buffer is not None: 

217 return self._getitem_mmap(index) 

218 

219 return self._getitem_file(index) 

220 

221 def get(self, index: int, offset: int = 0, length: int | None = None) -> np.ndarray: 

222 """Read a contiguous token range from one sequence. 

223 

224 Args: 

225 index: Sequence index. 

226 offset: First token offset in the sequence. 

227 length: Number of tokens to read, or all remaining tokens. 

228 

229 Returns: 

230 A NumPy view of the requested tokens. 

231 """ 

232 index_reader = self._require_index() 

233 sequence_pointer, sequence_length = index_reader[index] 

234 sequence_length = int(sequence_length) 

235 read_length = sequence_length - offset if length is None else length 

236 if offset < 0 or read_length < 0 or offset + read_length > sequence_length: 

237 raise ValueError( 

238 f"Token range [{offset}, {offset + read_length}) exceeds sequence length {sequence_length}" 

239 ) 

240 

241 byte_offset = int(sequence_pointer) + offset * index_reader.dtype.itemsize 

242 if self.bin_buffer: 

243 sequence = np.frombuffer(self.bin_buffer, dtype=index_reader.dtype, count=read_length, offset=byte_offset) 

244 else: 

245 sequence = np.empty(read_length, dtype=index_reader.dtype) 

246 with open(self.path_prefix + ".bin", mode="rb", buffering=0) as bin_buffer_file: 

247 bin_buffer_file.seek(byte_offset) 

248 bytes_read = bin_buffer_file.readinto(sequence) 

249 

250 byte_length = read_length * index_reader.dtype.itemsize 

251 if bytes_read != byte_length: 

252 raise ValueError(f"Indexed Dataset read returned {bytes_read} bytes; expected {byte_length}") 

253 

254 return sequence 

255 

256 @property 

257 def sequence_lengths(self) -> np.ndarray: 

258 """Return the token length of every indexed sequence.""" 

259 return self._require_index().sequence_lengths 

260 

261 @property 

262 def document_indices(self) -> np.ndarray: 

263 """Return sequence indices that delimit documents.""" 

264 return self._require_index().document_indices 

265 

266 def get_document_indices(self) -> np.ndarray: 

267 """Return sequence indices that delimit documents.""" 

268 return self.document_indices 

269 

270 def set_document_indices(self, document_indices: np.ndarray) -> None: 

271 """Replace sequence indices that delimit documents.""" 

272 self._require_index().document_indices = document_indices 

273 

274 def _require_index(self) -> _IndexReader: 

275 """Return the initialized index reader.""" 

276 if self.index is None: 

277 raise ValueError("Indexed Dataset metadata is not initialized") 

278 return self.index 

279 

280 

281class ArrowDataReader: 

282 """Memory-map one Arrow file, the low-level source of one SFT Dataset.""" 

283 

284 def __init__(self, dataset_path: str) -> None: 

285 """Open one HF Arrow IPC stream; directory discovery belongs to the builder.""" 

286 path = Path(dataset_path) 

287 if path.suffix != ".arrow" or not path.is_file(): 

288 raise ValueError("ArrowDataReader requires one .arrow file, not a directory") 

289 # Keep the HF dependency out of Trainer and bin/idx import paths. 

290 from datasets import Dataset # pylint: disable=C0415 

291 

292 self.dataset = Dataset.from_file(str(path)) 

293 

294 def __len__(self) -> int: 

295 """Return the number of complete prepared records in this file.""" 

296 return len(self.dataset) 

297 

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

299 """Return source columns for one record, including labels and boundaries.""" 

300 return self.dataset[int(index)]