Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_data_reader.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"""Read token sequences from paired ``.idx/.bin`` files.""" 

16 

17from __future__ import annotations 

18 

19import os 

20import struct 

21import time 

22from functools import lru_cache 

23from typing import ClassVar 

24 

25import numpy as np 

26 

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

28from hyper_parallel.data.dataset_logging import get_dataset_logger 

29 

30logger = get_dataset_logger(__name__) 

31 

32 

33class _IndexReader: 

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

35 

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

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

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

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

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

41 if header != INDEX_HEADER: 

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

43 

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

45 if version != INDEX_VERSION: 

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

47 

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

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

50 self.dtype_size = self.dtype.itemsize 

51 

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

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

54 offset = stream.tell() 

55 

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

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

58 # has document_count entries. 

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

60 self.bin_buffer = memoryview(self.bin_buffer_mmap) 

61 self.sequence_lengths = np.frombuffer( 

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

63 ) 

64 

65 self.sequence_pointers = np.frombuffer( 

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

67 ) 

68 

69 self.document_indices = np.frombuffer( 

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

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

72 ) 

73 

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

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

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

77 

78 def __del__(self) -> None: 

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

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

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

82 if mmap_handle is not None: 

83 mmap_handle.close() 

84 

85 def __len__(self) -> int: 

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

87 return self.sequence_count 

88 

89 @lru_cache(maxsize=8) 

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

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

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

93 

94 

95class IndexedDataReader: 

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

97 

98 Args: 

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

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

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

102 matches corpora whose shards share an identical index layout. 

103 """ 

104 

105 cached_index_reader: ClassVar[_IndexReader | None] = None 

106 

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

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

109 self.path_prefix = "" 

110 self.mmap = False 

111 self.reuse_index = reuse_index 

112 self.index: _IndexReader | None = None 

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

114 self.bin_buffer: memoryview | None = None 

115 self.initialize(path_prefix, mmap) 

116 

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

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

119 

120 Args: 

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

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

123 """ 

124 start_time = time.time() 

125 index_path = path_prefix + ".idx" 

126 data_path = path_prefix + ".bin" 

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

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

129 

130 self.path_prefix = path_prefix 

131 self.mmap = mmap 

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

133 if self.index is None: 

134 if reused_cached_index: 

135 self.index = IndexedDataReader.cached_index_reader 

136 else: 

137 self.index = _IndexReader(index_path) 

138 if self.reuse_index: 

139 IndexedDataReader.cached_index_reader = self.index 

140 

141 self.bin_buffer_mmap = None 

142 self.bin_buffer = None 

143 if mmap: 

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

145 self.bin_buffer = memoryview(self.bin_buffer_mmap) 

146 logger.debug( 

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

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

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

150 reused_cached_index, time.time() - start_time, 

151 ) 

152 

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

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

155 return self.path_prefix, self.mmap, self.reuse_index 

156 

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

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

159 path_prefix, mmap, reuse_index = state 

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

161 

162 def __del__(self) -> None: 

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

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

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

166 if mmap_handle is not None: 

167 mmap_handle.close() 

168 

169 def __len__(self) -> int: 

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

171 return len(self._require_index()) 

172 

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

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

175 index_reader = self._require_index() 

176 

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

178 sequence_pointer, sequence_length = index_reader[index] 

179 sequence = np.frombuffer( 

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

181 ) 

182 return sequence 

183 

184 if isinstance(index, slice): 

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

186 if step != 1: 

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

188 if start == stop: 

189 return [] 

190 

191 sequence_lengths = index_reader.sequence_lengths[index] 

192 sequence_offsets = np.cumsum(sequence_lengths) 

193 sequences = np.frombuffer( 

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

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

196 ) 

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

198 return sequences 

199 

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

201 

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

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

204 if isinstance(index, slice): 

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

206 

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

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

209 

210 return self.get(int(index)) 

211 

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

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

214 if self.bin_buffer is not None: 

215 return self._getitem_mmap(index) 

216 

217 return self._getitem_file(index) 

218 

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

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

221 

222 Args: 

223 index: Sequence index. 

224 offset: First token offset in the sequence. 

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

226 

227 Returns: 

228 A NumPy view of the requested tokens. 

229 """ 

230 index_reader = self._require_index() 

231 sequence_pointer, sequence_length = index_reader[index] 

232 sequence_length = int(sequence_length) 

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

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

235 raise ValueError( 

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

237 ) 

238 

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

240 if self.bin_buffer: 

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

242 else: 

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

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

245 bin_buffer_file.seek(byte_offset) 

246 bytes_read = bin_buffer_file.readinto(sequence) 

247 

248 byte_length = read_length * index_reader.dtype.itemsize 

249 if bytes_read != byte_length: 

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

251 

252 return sequence 

253 

254 @property 

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

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

257 return self._require_index().sequence_lengths 

258 

259 @property 

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

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

262 return self._require_index().document_indices 

263 

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

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

266 return self.document_indices 

267 

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

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

270 self._require_index().document_indices = document_indices 

271 

272 def _require_index(self) -> _IndexReader: 

273 """Return the initialized index reader.""" 

274 if self.index is None: 

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

276 return self.index