| """ |
| Binary token dataset. |
| |
| Input: |
| train.bin or val.bin produced by encode_corpus.py |
| |
| Output: |
| x: input token IDs [block_size] |
| y: target token IDs [block_size] |
| """ |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
|
|
| class TokenBinDataset(Dataset): |
| """ |
| Memory-mapped token dataset. |
| |
| This avoids loading the full corpus into RAM. |
| """ |
|
|
| def __init__( |
| self, |
| bin_path: str, |
| meta_path: str, |
| block_size: int, |
| ): |
| self.bin_path = Path(bin_path) |
| self.meta_path = Path(meta_path) |
| self.block_size = block_size |
|
|
| with open(self.meta_path, "r", encoding="utf-8") as f: |
| self.meta = json.load(f) |
|
|
| dtype = np.dtype(self.meta["dtype"]) |
|
|
| self.data = np.memmap( |
| self.bin_path, |
| dtype=dtype, |
| mode="r", |
| ) |
|
|
| if len(self.data) <= block_size + 1: |
| raise ValueError("Dataset too small for selected block_size.") |
|
|
| self.n_samples = (len(self.data) - 1) // block_size |
|
|
| def __len__(self) -> int: |
| return self.n_samples |
|
|
| def __getitem__(self, idx: int): |
| """ |
| Input: |
| idx: sample index. |
| |
| Output: |
| x: input tokens. |
| y: next-token targets. |
| """ |
| start = idx * self.block_size |
| end = start + self.block_size + 1 |
|
|
| chunk = np.asarray(self.data[start:end], dtype=np.int64) |
|
|
| x = torch.from_numpy(chunk[:-1].copy()).long() |
| y = torch.from_numpy(chunk[1:].copy()).long() |
|
|
| return x, y |
|
|
