File size: 1,247 Bytes
d997548 | 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 | #!/usr/bin/env python3
import numpy as np, torch
from pathlib import Path
from torch.utils.data import Dataset
DATA = Path("/data/berna-r7-v3/core/data/tokenized")
FILES = ["eng.npy", "math.npy", "code.npy"]
class BernaDataset(Dataset):
def __init__(self, seq_len=4096):
self.seq_len = seq_len
parts = []
for f in FILES:
a = np.load(DATA / f, mmap_mode="r")
parts.append(np.asarray(a))
print(f" {f}: {len(a):,} tokens")
self.data = np.concatenate(parts)
self.total = len(self.data)
self.n_samples = (self.total - 1) // seq_len
print(f" Total: {self.total:,} | Samples: {self.n_samples:,}")
def __len__(self): return self.n_samples
def __getitem__(self, i):
# FIX (2026-10-03): return (x, x). Model does the shift internally.
# Previously returned (x, x.shift), which combined with model's shift
# produced 2-token-ahead prediction — a severe training bug.
s = i * self.seq_len
buf = self.data[s:s+self.seq_len+1]
x = torch.from_numpy(buf[:-1].astype(np.int64))
return x, x
if __name__ == "__main__":
ds = BernaDataset()
x, y = ds[0]
print(f"OK: {x.shape}, {y.shape}")
|