| import torch |
|
|
| with open("data/input.txt") as f: |
| text = f.read() |
|
|
| chars = sorted(list(set(text))) |
| vocab_size = len(chars) |
|
|
| stoi = {ch: i for i, ch in enumerate(chars)} |
| itos = {i: ch for i, ch in enumerate(chars)} |
|
|
|
|
| def encode(s): |
| return [stoi[c] for c in s] |
|
|
|
|
| def decode(l): |
| return "".join([itos[i] for i in l]) |
|
|
|
|
| data = torch.tensor(encode(text), dtype=torch.long) |
| n = int(0.9 * len(data)) |
| train_data = data[:n] |
| val_data = data[n:] |
|
|
|
|
| def get_batch(split, block_size, batch_size): |
| data = train_data if split == "train" else val_data |
| ix = torch.randint(len(data) - block_size, (batch_size,)) |
| x = torch.stack([data[i : i + block_size] for i in ix]) |
| y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix]) |
| return x, y |
|
|
