1 files changed, 18 insertions, 0 deletions
diff --git a/src/zbp_scaling/data.py b/src/zbp_scaling/data.py
new file mode 100644
index 0000000..c6946a0
--- /dev/null
+++ b/src/zbp_scaling/data.py
@@ -0,0 +1,18 @@
+"""Token shards: uint16 memmap files train.bin / val.bin in --data dir."""
+import os
+import numpy as np
+import torch
+
+
+class Shards:
+    def __init__(self, path, seq_len, device):
+        self.train = np.memmap(os.path.join(path, "train.bin"), dtype=np.uint16, mode="r")
+        self.val = np.memmap(os.path.join(path, "val.bin"), dtype=np.uint16, mode="r")
+        self.T, self.device = seq_len, device
+
+    def batch(self, split, bs, gen):
+        src = self.train if split == "train" else self.val
+        ix = torch.randint(0, len(src) - self.T - 1, (bs,), generator=gen)
+        x = torch.stack([torch.from_numpy(src[i:i + self.T].astype(np.int64)) for i in ix])
+        y = torch.stack([torch.from_numpy(src[i + 1:i + 1 + self.T].astype(np.int64)) for i in ix])
+        return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)