| import torch | |
| from torch.utils.data import Dataset | |
| class TextDataset(Dataset): | |
| def __init__(self, tokens, block_size): | |
| self.tokens = tokens | |
| self.block_size = block_size | |
| def __len__(self): | |
| return len(self.tokens) - self.block_size - 1 | |
| def __getitem__(self, idx): | |
| x = self.tokens[idx:idx+self.block_size] | |
| y = self.tokens[idx+1:idx+self.block_size+1] | |
| return x, y | |