swift15: corpus generator v2 — HF code/prose ingest, richer tool JSON, 8k KLD prompts

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 = ('<tool_call>{"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('<tool_call>{"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)