File size: 4,092 Bytes
11e2c7c | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Deterministic, resumable batch loader over the packed token corpus.
The corpus is one flat uint16 array in which every source already holds exactly its mixture share.
It is cut into 2048-token blocks and visited once in a fixed random order, so a single pass
reproduces the intended mix with no repetition. Step N always sees the same blocks (resume-safe).
"""
from __future__ import annotations
import json
import queue
import threading
from pathlib import Path
import numpy as np
import torch
BLOCK = 2048
class BlockLoader:
def __init__(self, data_dir: Path, seed: int = 1337, corpus_json: str = "corpus.json", train_file: str = "train.bin"):
self.dir = Path(data_dir)
self.meta = json.loads((self.dir / corpus_json).read_text())
self.mm = np.memmap(self.dir / train_file, dtype=np.uint16, mode="r")
self.n_blocks = (len(self.mm) - 1) // BLOCK
self.perm = np.random.default_rng(seed).permutation(self.n_blocks).astype(np.int64)
def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
"""Returns int64 [blocks_per_step * (BLOCK // seq_len), seq_len + 1]: inputs and shifted targets in one array."""
ids = self.perm[(np.arange(blocks_per_step) + step * blocks_per_step) % self.n_blocks]
out = np.empty((blocks_per_step, BLOCK // seq_len, seq_len + 1), dtype=np.int64)
for i, b in enumerate(ids):
w = self.mm[b * BLOCK : b * BLOCK + BLOCK + 1].astype(np.int64)
for j in range(BLOCK // seq_len):
out[i, j] = w[j * seq_len : j * seq_len + seq_len + 1]
return torch.from_numpy(out.reshape(-1, seq_len + 1))
class StagedLoader:
"""first loader for steps < from_step, then a second corpus from its own beginning (opened when first needed)."""
def __init__(self, first: BlockLoader, from_step: int, second_dir: Path, seed: int):
self.first, self.from_step, self.second_dir, self.seed = first, from_step, Path(second_dir), seed
self.second, self.n_blocks, self.meta = None, first.n_blocks, first.meta
def ready(self) -> bool:
return (self.second_dir / "corpus.json").exists()
def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
if step < self.from_step: return self.first.batch(step, blocks_per_step, seq_len)
if self.second is None: self.second = BlockLoader(self.second_dir, seed=self.seed)
return self.second.batch(step - self.from_step, blocks_per_step, seq_len)
class Prefetcher:
"""Background thread that prepares upcoming steps' batches."""
def __init__(self, loader: BlockLoader, start_step: int, blocks_per_step: int, seq_len_fn, depth: int = 3):
self.q: queue.Queue = queue.Queue(maxsize=depth)
self.loader, self.bps, self.seq_len_fn = loader, blocks_per_step, seq_len_fn
self.step = start_step
self.stop = False
self.t = threading.Thread(target=self._run, daemon=True)
self.t.start()
def _run(self):
s = self.step
while not self.stop:
T = self.seq_len_fn(s)
b = self.loader.batch(s, self.bps, T).pin_memory()
self.q.put((s, T, b))
s += 1
def get(self):
return self.q.get()
class ValSet:
"""Fixed held-out token windows, overall and per source."""
def __init__(self, data_dir: Path, seq_len: int = 2048, per_source_windows: int = 64):
meta = json.loads((Path(data_dir) / "corpus.json").read_text())
mm = np.memmap(Path(data_dir) / "val.bin", dtype=np.uint16, mode="r")
self.sets: dict[str, torch.Tensor] = {}
for name, s in meta["sources"].items():
n = (s["val_tokens"] - 1) // seq_len
if n <= 0:
continue
take = min(n, per_source_windows)
starts = np.linspace(0, n - 1, take).astype(int)
arr = np.stack([mm[s["val_start"] + i * seq_len : s["val_start"] + i * seq_len + seq_len + 1].astype(np.int64) for i in starts])
self.sets[name] = torch.from_numpy(arr)
|