File size: 2,463 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import numpy as np
import torch
class BinDataset:
"""Serves training batches from one or more memory-mapped .bin token files.
Memory-mapping lets us read slices of a huge file without loading it into RAM.
Each sample is a random window; input and target are the window shifted by one.
Also returns segment IDs so attention can stay inside document boundaries.
"""
def __init__(
self,
bin_paths: list[str],
seq_len: int,
weights: list[float] | None = None,
eot_token_id: int = 151643,
):
self.seq_len = seq_len
self.eot_token_id = eot_token_id
# open each source as a read-only memory map of uint32 token ids
self.data = [np.memmap(p, dtype=np.uint32, mode="r") for p in bin_paths]
self.lengths = [len(d) for d in self.data]
# sampling weights = the data mix (e.g. 0.65 english, 0.20 code, 0.15 math)
if weights is None:
weights = [n / sum(self.lengths) for n in self.lengths]
self.weights = np.array(weights) / np.sum(weights)
def get_batch(self, batch_size: int, device: str = "cuda"):
x = np.empty((batch_size, self.seq_len), dtype=np.int64)
y = np.empty((batch_size, self.seq_len), dtype=np.int64)
seg = np.empty((batch_size, self.seq_len), dtype=np.int32)
for i in range(batch_size):
# pick a source according to the mix, then a random window inside it
src = np.random.choice(len(self.data), p=self.weights)
arr, n = self.data[src], self.lengths[src]
start = np.random.randint(0, n - self.seq_len - 1)
chunk = arr[start : start + self.seq_len + 1].astype(np.int64)
x[i] = chunk[:-1] # input: tokens 0..seq_len-1
y[i] = chunk[1:] # target: same window shifted by one
# segment IDs bump AFTER each EOT, so EOT stays with the doc it ends.
# Example x[i]=[a,b,EOT,c,d,EOT,e] -> seg[i]=[0,0,0,1,1,1,2]
is_eot = (x[i] == self.eot_token_id).astype(np.int32)
seg[i] = np.cumsum(is_eot) - is_eot
# pin_memory + non_blocking speeds up the CPU->GPU copy during training
xt = torch.from_numpy(x).pin_memory().to(device, non_blocking=True)
yt = torch.from_numpy(y).pin_memory().to(device, non_blocking=True)
st = torch.from_numpy(seg).pin_memory().to(device, non_blocking=True)
return xt, yt, st |