1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
"""Token shards: uint16 memmap files train.bin / val.bin in --data dir."""
import os
import numpy as np
import torch
class Shards:
def __init__(self, path, seq_len, device):
self.train = np.memmap(os.path.join(path, "train.bin"), dtype=np.uint16, mode="r")
self.val = np.memmap(os.path.join(path, "val.bin"), dtype=np.uint16, mode="r")
self.T, self.device = seq_len, device
def batch(self, split, bs, gen):
src = self.train if split == "train" else self.val
ix = torch.randint(0, len(src) - self.T - 1, (bs,), generator=gen)
x = torch.stack([torch.from_numpy(src[i:i + self.T].astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy(src[i + 1:i + 1 + self.T].astype(np.int64)) for i in ix])
return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
|
