File size: 1,617 Bytes
27d37fd | 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """
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
|