File size: 522 Bytes
ae9e4fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | import numpy as np
def load_ids(path):
return np.load(path)
def get_batch(ids, block_size, batch_size, rng):
hi = len(ids) - block_size - 1
ix = rng.integers(0, hi, size=batch_size)
x = np.stack([ids[i : i + block_size] for i in ix]).astype(np.int64)
y = np.stack([ids[i + 1 : i + 1 + block_size] for i in ix]).astype(np.int64)
return x, y
def iter_batches(ids, block_size, batch_size, rng, n_batches):
for _ in range(n_batches):
yield get_batch(ids, block_size, batch_size, rng)
|