| 177 | |
| 178 | |
| 179 | class PackedDatasetIterator: |
| 180 | def __init__(self, filenames, n_chunks, block_size, seed, shuffle, wrap): |
| 181 | self._seed = seed |
| 182 | self._shuffle = shuffle |
| 183 | self._rng = np.random.default_rng(seed) if shuffle else None |
| 184 | self._block_idxs = None |
| 185 | |
| 186 | self._wrap = wrap |
| 187 | |
| 188 | # TODO: instead of filenames, we could have a single text stream |
| 189 | # (or text file) with the sequence of all files to be |
| 190 | # fetched/loaded. |
| 191 | self._filenames = filenames |
| 192 | self._file_idx = 0 |
| 193 | |
| 194 | self._n_chunks = n_chunks |
| 195 | |
| 196 | self._dtype = None |
| 197 | self._block_size = block_size |
| 198 | self._n_blocks = None |
| 199 | |
| 200 | self._mmaps = [] |
| 201 | self._buffers = [] |
| 202 | |
| 203 | self._block_idxs = [] |
| 204 | self._curr_idx = 0 |
| 205 | |
| 206 | self._load_n_chunks() |
| 207 | |
| 208 | def _read_header(self, path): |
| 209 | with open(path, 'rb') as f: |
| 210 | magic = f.read(len(HDR_MAGIC)) |
| 211 | assert magic == HDR_MAGIC, "File doesn't match expected format." |
| 212 | version = struct.unpack('<Q', f.read(8)) |
| 213 | assert version == (1,) |
| 214 | (dtype_code,) = struct.unpack('<B', f.read(1)) |
| 215 | dtype = dtypes[dtype_code] |
| 216 | (chunk_size,) = struct.unpack('<Q', f.read(8)) |
| 217 | return dtype, chunk_size |
| 218 | |
| 219 | def _close_mmaps(self): |
| 220 | for mmap in self._mmaps: |
| 221 | mmap._mmap.close() |
| 222 | |
| 223 | def _load_n_chunks(self): |
| 224 | self._close_mmaps() |
| 225 | self._mmaps = [] |
| 226 | self._buffers = [] |
| 227 | |
| 228 | # if n_chunks is larger than the number of files assigned |
| 229 | if self._n_chunks > len(self._filenames[self._file_idx :]): |
| 230 | if not self._wrap: |
| 231 | print( |
| 232 | 'No more chunks, stopping. (Note: If this happens when preparing data, see https://github.com/Lightning-AI/lit-llama/issues/425)' |
| 233 | ) |
| 234 | raise StopIteration |
| 235 | |
| 236 | self._file_idx = 0 |