import numpy as np
import torch
from torch.utils.data import Dataset


class TokenDataset(Dataset):
    def __init__(
        self,
        path,
        seq_len=2048,
    ):
        self.seq_len = seq_len

        self.tokens = np.memmap(
            path,
            dtype=np.uint32,
            mode="r",
        )

        self.length = (
            len(self.tokens) // seq_len
        ) - 1

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        start = idx * self.seq_len
        end = start + self.seq_len + 1

        chunk = self.tokens[start:end]

        x = torch.from_numpy(
            chunk[:-1].astype(np.int64)
        )

        y = torch.from_numpy(
            chunk[1:].astype(np.int64)
        )

        return x, y