Upload training_v2/data/curriculum_dataset.py with huggingface_hub

training_v2/data/curriculum_dataset.py

@@ -0,0 +1,128 @@
+"""Curriculum dataset with replay buffer for continual pre-training.
+
+Mixes multiple phase corpora at configurable ratios so later phases keep seeing
+earlier-phase data (replay), preventing catastrophic forgetting of conversation
+ability when training on technical content.
+
+Each shard directory is treated as one big concatenated stream of token IDs.
+We memory-map the .bin files so we can sample windows of `block_size+1` tokens
+without loading anything into RAM upfront.
+"""
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import IterableDataset
+
+
+def _load_shards(phase_dir, dtype):
+    phase_dir = Path(phase_dir)
+    shards = sorted(phase_dir.glob("*.bin"))
+    if not shards:
+        raise FileNotFoundError(f"no shards under {phase_dir}")
+    arrs = [np.memmap(s, dtype=dtype, mode="r") for s in shards]
+    sizes = np.array([a.shape[0] for a in arrs], dtype=np.int64)
+    total = int(sizes.sum())
+    return arrs, sizes, total
+
+
+class PhaseStream:
+    """Random-window sampler over the concatenation of all shards in one phase."""
+
+    def __init__(self, phase_dir, block_size, dtype=np.uint16, seed=0):
+        self.dir = Path(phase_dir)
+        self.block_size = block_size
+        self.arrs, self.sizes, self.total = _load_shards(phase_dir, dtype)
+        self.cum = np.concatenate([[0], np.cumsum(self.sizes)])
+        self.rng = np.random.default_rng(seed)
+        if self.total <= block_size + 1:
+            raise ValueError(f"phase {phase_dir} too small: {self.total} <= {block_size+1}")
+
+    def sample(self):
+        start = int(self.rng.integers(0, self.total - self.block_size - 1))
+        shard_idx = int(np.searchsorted(self.cum, start, side="right") - 1)
+        local_start = start - int(self.cum[shard_idx])
+        end_local = local_start + self.block_size + 1
+        a = self.arrs[shard_idx]
+        if end_local <= a.shape[0]:
+            buf = np.asarray(a[local_start:end_local], dtype=np.int64)
+        else:
+            head = np.asarray(a[local_start:], dtype=np.int64)
+            need = self.block_size + 1 - head.shape[0]
+            nxt = self.arrs[(shard_idx + 1) % len(self.arrs)]
+            tail = np.asarray(nxt[:need], dtype=np.int64)
+            buf = np.concatenate([head, tail])
+        return buf
+
+
+class MixedCurriculumDataset(IterableDataset):
+    """Sample from multiple phases according to mixing weights.
+
+    Args:
+        phase_dirs: dict {name: path}
+        weights:    dict {name: float} (will be normalized)
+        block_size: context length per sample
+        dtype:      numpy dtype of binary shards
+        seed:       base RNG seed (offset by worker id)
+    """
+
+    def __init__(self, phase_dirs, weights, block_size, dtype=np.uint16, seed=0):
+        super().__init__()
+        self.phase_dirs = phase_dirs
+        self.weights = weights
+        self.block_size = block_size
+        self.dtype = dtype
+        self.seed = seed
+
+        names = [n for n in phase_dirs if weights.get(n, 0) > 0]
+        if not names:
+            raise ValueError("at least one phase must have weight > 0")
+        ws = np.array([weights[n] for n in names], dtype=np.float64)
+        ws = ws / ws.sum()
+        self.names = names
+        self.probs = ws
+
+    def __iter__(self):
+        info = torch.utils.data.get_worker_info()
+        worker_id = info.id if info is not None else 0
+        seed = self.seed + worker_id * 1009
+        streams = {n: PhaseStream(self.phase_dirs[n], self.block_size,
+                                  dtype=self.dtype, seed=seed + i)
+                   for i, n in enumerate(self.names)}
+        rng = np.random.default_rng(seed + 7919)
+        while True:
+            name = self.names[int(rng.choice(len(self.names), p=self.probs))]
+            buf = streams[name].sample()
+            x = torch.from_numpy(buf[:-1]).long()
+            y = torch.from_numpy(buf[1:]).long()
+            yield x, y, name
+
+
+def load_phase_summary(out_root):
+    summary_path = Path(out_root) / "summary.json"
+    return json.loads(summary_path.read_text()) if summary_path.exists() else {}
+
+
+def make_phase_mix(phase_idx, replay_conv=None, replay_tech=None):
+    """Return mixing weights for each curriculum phase.
+
+    Phase 1 (conversational base): 100% conv
+    Phase 2 (cybersec, with replay): (1-replay_conv) tech + replay_conv conv
+    Phase 3 (tools, with replay):     0.7 tools + replay_tech tech + replay_conv conv
+
+    Defaults match v2 (25%/10%). For v4 (less subtitle contamination)
+    pass replay_conv=0.10 in phase 2.
+    """
+    if phase_idx == 1:
+        return {"phase1_conv": 1.0, "phase2_tech": 0.0, "phase3_tools": 0.0}
+    if phase_idx == 2:
+        rc = 0.25 if replay_conv is None else replay_conv
+        return {"phase1_conv": rc, "phase2_tech": 1.0 - rc, "phase3_tools": 0.0}
+    if phase_idx == 3:
+        rc = 0.10 if replay_conv is None else replay_conv
+        rt = 0.20 if replay_tech is None else replay_tech
+        return {"phase1_conv": rc, "phase2_tech": rt, "phase3_tools": 1.0 - rc - rt}
+    raise ValueError(f"phase_idx must be 1..3, got {phase_idx}")