File size: 7,794 Bytes
c83c617 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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()
|