| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import IterableDataset, get_worker_info |
|
|
|
|
| class MemmapTokenDataset(IterableDataset): |
| """Minimal contiguous adapter for an existing pretokenized stream. |
| |
| Each item contains ``seq_len`` inputs and their immediately following |
| next-token labels. Consecutive items advance by ``seq_len`` tokens, so the |
| transition at each sequence boundary is preserved instead of skipped. |
| |
| Formal experiments should still use the project's canonical logical-stream |
| loader so exact shard order, cursor and hashes match prior runs. |
| """ |
|
|
| def __init__(self, path: str, seq_len: int, dtype: str = "uint16", seed: int = 1337, start_sequence: int = 0): |
| self.path = path |
| self.seq_len = seq_len |
| self.dtype = np.dtype(dtype) |
| self.seed = seed |
| self.start_sequence = start_sequence |
|
|
| def __iter__(self): |
| arr = np.memmap(self.path, mode="r", dtype=self.dtype) |
| n = (len(arr) - 1) // self.seq_len |
| worker = get_worker_info() |
| worker_id = 0 if worker is None else worker.id |
| workers = 1 if worker is None else worker.num_workers |
| first = self.start_sequence + worker_id |
| for i in range(first, n, workers): |
| start = i * self.seq_len |
| x = np.asarray(arr[start : start + self.seq_len], dtype=np.int64) |
| y = np.asarray(arr[start + 1 : start + self.seq_len + 1], dtype=np.int64) |
| yield torch.from_numpy(x.copy()), torch.from_numpy(y.copy()) |
|
|
