"""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)