diff --git a/scripts/make-calib-corpus.py b/scripts/make-calib-corpus.py index 0f90fb4..86203a3 100755 --- a/scripts/make-calib-corpus.py +++ b/scripts/make-calib-corpus.py @@ -1,34 +1,33 @@ #!/usr/bin/env python3 -"""Build the two Swift1.5 campaign corpora from the repo checkout. +"""Build the two Swift1.5 campaign corpora (calibration + KLD). calib.jsonl — 10M tokens of agent-shaped RAW text (45% code, 30% prose, 25% tool-call JSON), one {"prompt": text} per line, pushed to /v1/completions with max_tokens=1 during the calibration run (the Hessian -tap is of the prompt tokens; docs/MOE-W4.md step 2). +tap is of the prompt tokens; docs/MOE-W4.md step 2). Code and prose mix +the repo checkout with HF dataset parquets (Frank's call 2026-10-06: +local material is only ~1.6M unique tokens; MOE-W4 wants 10M): +nampdn-ai/tiny-codes (code) and Salesforce/wikitext wikitext-103-raw +(prose). Tool calls are synthesized from commands harvested from the +repo's own scripts plus a fleet-ops set, with argument-shaped variants. kld.jsonl — 1.8M tokens, one {"prompt", "score_from", "source"} per -line; each prompt is a two-turn conversation rendered through the -Froggeric v22.5 template, and the last 256 tokens (the tail of the -assistant answer) are the scored span: score_from = n - 256. The KLD +line; each prompt is a multi-turn conversation rendered through the +Froggeric v22.5 template with a >= --kld-min-prompt (8192) token target, +and the last 256 tokens (the tail of the assistant answer) are the +scored span: score_from = n - 256. At 8k-token prompts the 1.8M corpus +is ~220 prompts = ~56k scored positions ~= 25 GB of record, matching +the published Flash-Next KLD scale (83k positions, MOE-W4.md). The KLD reference run and every candidate run consume the SAME file (TOOLS.md ---kld-record / --kld-ref). +--kld-record / --kld-ref). The KLD corpus is LOCAL material only, so the +file-level disjointness rule (fnv1a64(path) % 5 == 0 -> KLD pool) is +untouched by the HF additions. -Disjointness is at file granularity: fnv1a64(relative path) % 5 == 0 -goes to the KLD pool, the rest to the calibration pool; synthesized tool -calls partition by template id. The manifest records the token counts -per source class and the empty intersection. - -Content sources: shell/python/jinja/recipe/config files (code), -serverAppConfigs-docs markdown (prose), and synthesized agent tool -calls — JSON envelopes whose commands are harvested from the repo's own -scripts (ssh-prefixed) plus a vendored fleet-ops command set. - -Deps: tokenizers, jinja2 (see /home/dank/venvs/radcorpus on dankserver). -The tokenizer loads from the checkpoint's tokenizer.json; the template -is the deployed Froggeric v22.5 file. Usage: - - make-calib-corpus.py --repo ROOT --tokenizer-json PATH --template PATH - --out-dir DIR [--calib-tokens 10000000] [--kld-tokens 1800000] +Deps: tokenizers, jinja2, pyarrow (venv /home/dank/venvs/radcorpus). +Usage: make-calib-corpus.py --repo ROOT --tokenizer-json PATH + --template PATH --out-dir DIR [--hf-code-parquet F,F] + [--hf-prose-parquet F,F] [--calib-tokens 10000000] + [--kld-tokens 1800000] """ import argparse import json @@ -61,14 +60,18 @@ FLEET_CMDS = [ "python3 -c \"import json; print(json.load(open('config.json'))['num_hidden_layers'])\"", "ip neigh show 192.168.0.235", "amdgpu_top --device-info 2>/dev/null | head -20", - "cat /proc/meminfo | grep -E 'MemAvailable|MemAvailable'", + "grep -E 'MemAvailable' /proc/meminfo", "sha256sum model-00001-of-00131.safetensors", "git log --oneline -10", "awk '{s+=$1} END {print s}' shard_sizes.txt", ] -TOOL_NAMES = ["bash", "read", "grep", "web_search"] -TOOL_JSON = ('{"name": "%s", "arguments": {"command": "%s"}}' - '<|im_end|>\n') +TOOL_NAMES = ["bash", "read", "grep", "web_search", "edit", "write"] +WORKDIRS = ["/home/dank", "/home/dank/serverAppConfigs", "/home/dank/models", + "/etc/systemd/system", "/var/log"] +DESCS = ["check service state", "tail the engine ledger", "list containers", + "query health", "read metrics", "check disk", "check memory", + "restart the endpoint", "sync the checkout", "run deployer", + "run the test seams", "inspect the index", "probe the peer"] def fnv(s: str) -> int: @@ -100,15 +103,44 @@ def harvest_cmds(repo: str, rng: random.Random) -> list[str]: def tool_call(rng: random.Random, cmds: list[str]) -> str: - return TOOL_JSON % (rng.choice(TOOL_NAMES), - rng.choice(cmds).replace('"', "'")) + """One agent turn: 2-4 tool calls with argument-shaped JSON.""" + parts = [] + for _ in range(rng.randint(2, 4)): + args: dict = {"command": rng.choice(cmds)} + if rng.random() < 0.4: + args["cwd"] = rng.choice(WORKDIRS) + if rng.random() < 0.3: + args["timeout"] = rng.choice([10, 30, 60, 300, 600]) + if rng.random() < 0.5: + args["description"] = rng.choice(DESCS) + parts.append('{"name": "%s", "arguments": %s}<|im_end|>' + % (rng.choice(TOOL_NAMES), + json.dumps(args, ensure_ascii=False))) + return "\n".join(parts) -def render(env: Environment, tokenizer: Tokenizer, template_name: str, - messages: list) -> tuple[str, int]: - out = env.get_template(template_name).render( - messages=messages, add_generation_prompt=False) - return out, len(tokenizer.encode(out, add_special_tokens=False).ids) +def load_parquet_texts(path: str, kind: str, rng: random.Random, + cap: int) -> list[str]: + import pyarrow.parquet as pq + out = [] + pf = pq.ParquetFile(path) + for rg in pf.iter_row_groups(): + d = rg.to_pydict() + n = len(next(iter(d.values()))) + idx = range(n) if n <= cap else rng.sample(range(n), cap) + for i in idx: + if kind == "code": + frag = ("// %s\n### Instruction:\n%s\n### Response:\n%s" + % (d.get("file_name", [""])[i], d["prompt"][i], + d["code"][i] or d["response"][i])) + else: + frag = d["text"][i] + if frag and frag.strip(): + out.append(frag[:8000]) + if len(out) >= cap: + break + rng.shuffle(out) + return out[:cap] def main() -> None: @@ -121,6 +153,9 @@ def main() -> None: ap.add_argument("--kld-tokens", type=int, default=1_800_000) ap.add_argument("--chunk-tokens", type=int, default=8192) ap.add_argument("--score-tail", type=int, default=256) + ap.add_argument("--kld-min-prompt", type=int, default=8192) + ap.add_argument("--hf-code-parquet", default="") + ap.add_argument("--hf-prose-parquet", default="") ap.add_argument("--seed", type=int, default=20261006) a = ap.parse_args() @@ -162,6 +197,7 @@ def main() -> None: toks = tokenizer.encode(raw, add_special_tokens=False).ids return tokenizer.decode(toks[:2400], skip_special_tokens=True) + cmds = harvest_cmds(a.repo, rng) # fragment pools per class, corpus-split first frags: dict[str, dict[str, list[str]]] = {} for cls in pools: @@ -169,49 +205,63 @@ def main() -> None: fr = {} src = {"kld": kld_files, "cal": cal_files} if cls == "tools": - cmds = harvest_cmds(a.repo, rng) - src["kld"] = [f"synth:{i}" for i in range(0, len(cmds), 2)] - src["cal"] = [f"synth:{i}" for i in range(1, len(cmds), 2)] - fr["kld"] = [tool_call(rng, src["kld"][i::1]) if cls == "tools" - else text_of(p) for i, p in enumerate(src["kld"])] - fr["cal"] = [tool_call(rng, src["cal"][i::1]) if cls == "tools" - else text_of(p) for i, p in enumerate(src["cal"])] + fr["kld"] = [tool_call(rng, cmds) for _ in range(600)] + fr["cal"] = [tool_call(rng, cmds) for _ in range(40000)] + else: + fr["kld"] = [text_of(p) for p in src["kld"]] + fr["cal"] = [text_of(p) for p in src["cal"]] frags[cls] = fr + for cls, spec in (("code", a.hf_code_parquet), + ("prose", a.hf_prose_parquet)): + for path in [p for p in spec.split(",") if p]: + frags[cls]["cal"].extend( + load_parquet_texts(path, cls, rng, cap=40000)) + for cls in frags: + rng.shuffle(frags[cls]["cal"]) - def emit_calib() -> tuple[int, list[str]]: + def toks(s: str) -> int: + return len(tokenizer.encode(s, add_special_tokens=False).ids) + + def emit_calib() -> tuple[int, list[dict]]: targets = {"code": .45, "prose": .30, "tools": .25} - rows, used = [], [] + rows: list[dict] = [] for cls, frac in targets.items(): want = int(a.calib_tokens * frac) pool = frags[cls]["cal"] - rng.shuffle(pool) - got, i, buf = 0, 0, [] - nbuf = 0 - while got < want and i < len(pool) * 3: + if not pool: + continue + got, i, buf, nbuf = 0, 0, [], 0 + while got < want: frag = pool[i % len(pool)] - n = len(tokenizer.encode(frag, add_special_tokens=False).ids) + i += 1 + n = toks(frag) if nbuf + n > a.chunk_tokens and buf: - rows.append({"prompt": "\n".join(buf)}) - used.extend(set()) + rows.append({"prompt": "\n".join(buf), "cls": cls}) got += nbuf buf, nbuf = [], 0 + if got >= want: + break buf.append(frag) nbuf += n - i += 1 - if buf: - rows.append({"prompt": "\n".join(buf)}) + if i > len(pool) * 40: + break # pool exhausted with repeats; stop + if buf and got < want: + rows.append({"prompt": "\n".join(buf), "cls": cls}) got += nbuf + print(f"calib {cls}: {got} tokens, {i} fragments") rng.shuffle(rows) - return len(rows), rows + return len(rows), [{k: v for k, v in r.items() if k == "prompt"} + for r in rows] - def emit_kld() -> tuple[int, list[str]]: + def emit_kld() -> tuple[int, list[dict]]: rows, total = [], 0 classes = ["code", "prose", "tools"] i = {c: 0 for c in classes} - while total < a.kld_tokens: - msgs, srcs, ntok = [], [], 0 - turns = rng.choice([2, 3]) - for t in range(turns): + guard = 0 + while total < a.kld_tokens and guard < 5000: + guard += 1 + msgs, srcs = [], [] + for t in range(rng.choice([2, 3, 4])): cls = classes[(len(rows) + t) % 3] pool = frags[cls]["kld"] if not pool: @@ -224,10 +274,13 @@ def main() -> None: msgs.append({"role": "assistant", "content": frag[:4000]}) srcs.append(cls) - rendered, n = render(env, tokenizer, tname, msgs) - if n <= a.score_tail: + rendered = env.get_template(tname).render( + messages=msgs, add_generation_prompt=False) + n = toks(rendered) + if n < a.kld_min_prompt: continue - rows.append({"prompt": rendered, "score_from": n - a.score_tail, + rows.append({"prompt": rendered, + "score_from": n - a.score_tail, "source": "+".join(sorted(set(srcs)))}) total += n return len(rows), rows @@ -241,17 +294,17 @@ def main() -> None: with open(os.path.join(a.out_dir, "kld.jsonl"), "w") as f: for r in krows: f.write(json.dumps(r) + "\n") - scored = sum(len(tokenizer.encode( - r["prompt"][max(0, 0):], add_special_tokens=False).ids) - - r["score_from"] for r in krows) + scored = sum(len(tokenizer.encode(r["prompt"], add_special_tokens=False).ids) + - r["score_from"] for r in krows) manifest = { "seed": a.seed, "template": tname, "tokenizer": os.path.basename(a.tokenizer_json), "calib_prompts": nc, "kld_prompts": nk, "kld_scored_positions": scored, "kld_score_tail": a.score_tail, - "disjoint_file_pools": True, + "disjoint_file_pools": "local KLD material only; HF code feeds calib", "calib_mix": {"code": .45, "prose": .30, "tools": .25}, + "hf_code": a.hf_code_parquet, "hf_prose": a.hf_prose_parquet, } with open(os.path.join(a.out_dir, "manifest.json"), "w") as f: json.dump(manifest, f, indent=1)