File size: 3,706 Bytes
0cf8154 | 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 | """Deterministic sampling of training windows from memory-mapped token shards."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import torch
def effective_mix(target: dict[str, float], group_tokens: dict[str, int], total_tokens: int,
max_epochs: dict[str, float]) -> dict[str, float]:
"""Cap groups at `max_epochs` passes over their data; give the freed share to the other groups."""
capped = {}
for g, w in target.items():
if g in max_epochs:
cap = max_epochs[g] * group_tokens.get(g, 0) / total_tokens
if cap < w:
capped[g] = cap
free = 1.0 - sum(capped.values())
rest = {g: w for g, w in target.items() if g not in capped}
rest_total = sum(rest.values())
out = dict(capped)
for g, w in rest.items():
out[g] = free * w / rest_total
return {g: out[g] for g in target}
class MixtureLoader:
"""Samples (x, y) windows; batch contents depend only on (seed, step, micro)."""
def __init__(self, data_dir: Path, split: str, weights: dict[str, float], context_len: int, micro_batch: int,
seed: int):
self.data_dir = Path(data_dir)
self.manifest = json.loads((self.data_dir / "manifest.json").read_text())
self.context_len, self.micro_batch, self.seed = context_len, micro_batch, seed
self.groups: dict[str, tuple[list[np.memmap], np.ndarray]] = {}
for group, splits in sorted(self.manifest["groups"].items()):
if split not in splits:
continue
arrays, sizes = [], []
for shard in splits[split]["shards"]:
if shard["tokens"] < context_len + 1:
continue # too short to hold one window
# Shard paths are resolved under data_dir so the corpus can be moved (e.g. to Kaggle).
path = self.data_dir / group / Path(shard["path"]).name
arrays.append(np.memmap(path, dtype="<u2", mode="r"))
sizes.append(shard["tokens"] - context_len)
if arrays:
self.groups[group] = (arrays, np.asarray(sizes, dtype=np.float64))
usable = {g: w for g, w in weights.items() if w > 0 and g in self.groups}
if not usable:
raise ValueError(f"no usable shards for weights {weights} in split {split!r}")
total = sum(usable.values())
self.group_names = sorted(usable)
self.group_probs = np.asarray([usable[g] / total for g in self.group_names])
def _window(self, rng: np.random.Generator, group: str) -> np.ndarray:
arrays, sizes = self.groups[group]
k = rng.choice(len(arrays), p=sizes / sizes.sum())
off = int(rng.integers(0, sizes[k]))
return np.asarray(arrays[k][off : off + self.context_len + 1], dtype=np.int64)
def _stack(self, rows: list[np.ndarray]) -> tuple[torch.Tensor, torch.Tensor]:
batch = torch.from_numpy(np.stack(rows))
return batch[:, :-1].contiguous(), batch[:, 1:].contiguous()
def get_batch(self, step: int, micro: int) -> tuple[torch.Tensor, torch.Tensor]:
rng = np.random.default_rng([self.seed, step, micro])
groups = rng.choice(len(self.group_names), size=self.micro_batch, p=self.group_probs)
return self._stack([self._window(rng, self.group_names[g]) for g in groups])
def get_group_batch(self, group: str, index: int) -> tuple[torch.Tensor, torch.Tensor]:
gid = sorted(self.groups).index(group)
rng = np.random.default_rng([self.seed, 1_000_000_007, gid, index])
return self._stack([self._window(rng, group) for _ in range(self.micro_batch)])
|