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