| """Freeze virtual epochs: all web windows once, chat repeated to cropped-token quota.""" |
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "src")) |
| import numpy as np |
| from vimeml.training.data import SentenceWindowDataset, file_sha, write_json |
| from vimeml.training.data_v2 import PrefixCropWindowDataset |
| from vimeml.tokenizer.encode_parallel import output_lock |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--output", type=Path, default=ROOT / "artifacts/training-data/corpus-v3-chat05-c128") |
| parser.add_argument("--chat-fraction", type=float, default=.05) |
| parser.add_argument("--epochs", type=int, default=2) |
| parser.add_argument("--noise-filter", type=Path) |
| args = parser.parse_args() |
| if not 0 < args.chat_fraction < .5 or args.epochs < 1: |
| parser.error("Positive epochs and chat fraction in (0, .5) required") |
| token_dir = ROOT / "artifacts/token-data/corpus-v3-half/mixed" |
| index_dir = ROOT / "artifacts/token-data/corpus-v3-half/indexes/mixed" |
| output = args.output.resolve() |
| output.mkdir(parents=True, exist_ok=True) |
| request = {"token_manifest_sha256": file_sha(token_dir / "manifest.json"), |
| "window_manifest_sha256": file_sha(index_dir / "manifest.json"), |
| "index_dir": index_dir.relative_to(ROOT).as_posix(), "epochs": args.epochs, |
| "chat_fraction": args.chat_fraction, "seed": 42, "crop_probability": .30, "crop_minimum": 8} |
| noise = None |
| if args.noise_filter: |
| noise = json.loads((args.noise_filter / "manifest.json").read_text(encoding="utf-8")) |
| if noise["status"] != "complete" or noise["token_manifest_sha256"] != request["token_manifest_sha256"]: |
| raise ValueError("Noise filter belongs to different token data") |
| request["noise_filter_manifest_sha256"] = file_sha(args.noise_filter / "manifest.json") |
| with output_lock(output): |
| if (output / "manifest.json").exists(): |
| cached = json.loads((output / "manifest.json").read_text(encoding="utf-8")) |
| cached["request"]["index_dir"] = cached["request"]["index_dir"].replace("\\", "/") |
| if cached["request"] != request or cached["status"] != "complete": |
| raise ValueError("Changed mixture request; use a new output directory") |
| if any((output / r["file"]).stat().st_size != r["bytes"] for r in cached["epochs"]): |
| raise ValueError("Incomplete cached mixture") |
| print("Mixture already complete; reused", flush=True) |
| return |
| dataset = PrefixCropWindowDataset(token_dir, index_dir, seed=42, probability=.30) |
| if len(dataset) >= 2 ** 32: |
| raise ValueError("Mixture mapping exceeds uint32") |
| ids = np.arange(len(dataset), dtype=np.int64) |
| position = np.searchsorted(dataset.end, ids, side="right") |
| previous = np.zeros(len(ids), dtype=np.int64) |
| present = position > 0 |
| previous[present] = dataset.end[position[present] - 1] - dataset.sentences[position[present] - 1] - 1 |
| sentences = ids - previous |
| candidates = np.flatnonzero(position < len(dataset.end)) |
| long = candidates[ids[candidates] >= dataset.first[position[candidates]]] |
| sentences[long] = dataset.sentences[position[long]] |
| tm = json.loads((token_dir / "manifest.json").read_text(encoding="utf-8")) |
| chat_sources = [value for name, value in tm["source_ids"].items() if name in ("real-persona-chat", "mrmp-chat")] |
| sources = np.memmap(token_dir / "train.sources.bin", dtype="u1", mode="r") |
| is_chat = np.isin(sources[sentences], chat_sources) |
| sources._mmap.close() |
| excluded = np.zeros(len(ids), dtype=np.bool_) |
| if noise: |
| path = args.noise_filter / noise["file"] |
| if path.stat().st_size != noise["bytes"]: |
| raise ValueError("Noise filter file size mismatch") |
| with_filter = np.load(path, mmap_mode="r", allow_pickle=False) |
| if with_filter.shape != (tm["splits"]["train"]["sentences"],) or with_filter.dtype != np.bool_: |
| raise ValueError("Invalid noise filter shape/dtype") |
| excluded = with_filter[sentences].copy() |
| with_filter._mmap.close() |
| chat, web = ids[is_chat], ids[~is_chat & ~excluded] |
| if not len(chat) or not len(web): |
| raise ValueError("Both chat and web windows required") |
| del sentences, previous, candidates, long, present, position |
| raw = np.empty(len(ids), dtype=np.int64) |
| for start in range(0, len(ids), 262144): |
| raw[start:start + 262144] = SentenceWindowDataset.window_lengths(dataset, ids[start:start + 262144]) |
| records = [] |
| try: |
| for epoch in range(args.epochs): |
| dataset.set_epoch(epoch) |
| lengths = np.empty(len(ids), dtype=np.int64) |
| for start in range(0, len(ids), 262144): |
| lengths[start:start + 262144] = dataset.window_lengths(ids[start:start + 262144]) |
| web_tokens = int(lengths[web].sum()) |
| target = round(web_tokens * args.chat_fraction / (1 - args.chat_fraction)) |
| rng = np.random.default_rng(42 + epoch) |
| chosen, used = [], 0 |
| while used < target: |
| order = rng.permutation(chat) |
| cumulative = np.cumsum(lengths[order], dtype=np.int64) |
| take = int(np.searchsorted(cumulative, target - used, side="right")) |
| if take: |
| chosen.append(order[:take]) |
| used += int(cumulative[take - 1]) |
| if take < len(order): |
| break |
| if not chosen: |
| raise ValueError("Chat quota is smaller than a usable window") |
| chosen = np.concatenate(chosen) |
| mapping = np.concatenate((web, chosen)).astype("<u4") |
| rng.shuffle(mapping) |
| path = output / f"epoch-{epoch}.npy" |
| temporary = path.with_suffix(".tmp.npy") |
| np.save(temporary, mapping, allow_pickle=False) |
| temporary.replace(path) |
| record = {"epoch": epoch, "file": path.name, "bytes": path.stat().st_size, |
| "windows": len(mapping), "web_windows": len(web), "chat_windows": len(chosen), |
| "web_tokens": web_tokens, "chat_tokens": used, "effective_tokens": web_tokens + used, |
| "uncropped_tokens": int(raw[web].sum()) + int(raw[chosen].sum()), |
| "chat_fraction": used / (web_tokens + used), |
| "chat_effective_passes": used / int(lengths[chat].sum()), |
| "chat_independent_tokens": int(lengths[chat].sum())} |
| records.append(record) |
| print(json.dumps(record, indent=2), flush=True) |
| write_json(output / "progress.json", {"status": "building", "completed_epochs": epoch + 1}) |
| write_json(output / "manifest.json", {"status": "complete", "format": "vimeml_token_mixture_v1", |
| "request": request, "epochs": records, "policy": "Web windows exactly once per epoch; training chat only, shuffled repeated cycles; quota after deterministic prefix crop; whole-window rounding; no token copies.", |
| "verification": "Small manifest identity, mapping sizes/shapes; no full corpus or mapping SHA256"}) |
| write_json(output / "progress.json", {"status": "complete"}) |
| finally: |
| dataset.close() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
