Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """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) | |