File size: 1,558 Bytes
37781fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | 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())
|