swift15: corpus generator (10M calib + 1.8M KLD, Froggeric-rendered, disjoint pools)

scripts/make-calib-corpus.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+"""Build the two Swift1.5 campaign corpora from the repo checkout.
+
+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).
+
+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
+reference run and every candidate run consume the SAME file (TOOLS.md
+--kld-record / --kld-ref).
+
+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]
+"""
+import argparse
+import json
+import os
+import random
+import re
+
+from jinja2 import Environment, FileSystemLoader
+from tokenizers import Tokenizer
+
+CODE_EXT = (".sh", ".py", ".jinja", ".recipe", ".env", ".example",
+            ".yaml", ".yml", ".bats", ".service", ".container", ".toml")
+JSONL_EXTS = (".jsonl", ".json")
+DOC_DIR = "serverAppConfigs-docs"
+
+FLEET_CMDS = [
+    "systemctl status frank-primary --no-pager",
+    "journalctl -u frank-primary -n 100 --no-pager",
+    "podman ps --format '{{.Names}} {{.Status}}'",
+    "podman logs --tail 50 frank-primary",
+    "curl -s http://192.168.0.235:7770/health",
+    "curl -s http://192.168.0.235:7770/metrics | grep radiance_",
+    "df -h /home/dank | tail -1",
+    "free -g",
+    "sudo systemctl restart frank-primary",
+    "cd /home/dank/serverAppConfigs && git pull --ff-only",
+    "./scripts/deploy.sh dankserver --service frank-primary",
+    "./tests/run.sh 2>&1 | tail -5",
+    "jq . /home/dank/models/rad/index.json | head -40",
+    "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'",
+    "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')
+
+
+def fnv(s: str) -> int:
+    h = 0xcbf29ce484222325
+    for b in s.encode():
+        h = ((h ^ b) * 0x100000001b3) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
+def harvest_cmds(repo: str, rng: random.Random) -> list[str]:
+    cmds = list(FLEET_CMDS)
+    pat = re.compile(r"^\s*(?:sudo\s+)?([a-z][\w./-]*(?:\s+\S+){1,6})\s*$")
+    for root, dirs, files in os.walk(repo):
+        dirs[:] = [d for d in dirs if not d.startswith(".")]
+        for f in files:
+            if not f.endswith(".sh"):
+                continue
+            try:
+                for line in open(os.path.join(root, f), encoding="utf-8",
+                                 errors="replace"):
+                    m = pat.match(line.rstrip())
+                    if m and len(m.group(1)) > 12:
+                        cmds.append("ssh dank@192.168.0.235 '"
+                                    + m.group(1)[:90] + "'")
+            except OSError:
+                pass
+    rng.shuffle(cmds)
+    return cmds[:4000]
+
+
+def tool_call(rng: random.Random, cmds: list[str]) -> str:
+    return TOOL_JSON % (rng.choice(TOOL_NAMES),
+                        rng.choice(cmds).replace('"', "'"))
+
+
+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 main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--repo", required=True)
+    ap.add_argument("--tokenizer-json", required=True)
+    ap.add_argument("--template", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--calib-tokens", type=int, default=10_000_000)
+    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("--seed", type=int, default=20261006)
+    a = ap.parse_args()
+
+    rng = random.Random(a.seed)
+    tokenizer = Tokenizer.from_file(a.tokenizer_json)
+    env = Environment(loader=FileSystemLoader(os.path.dirname(a.template)),
+                      trim_blocks=True, lstrip_blocks=True)
+    env.filters["tojson"] = lambda v: json.dumps(v, ensure_ascii=False)
+    env.filters["raise_exception"] = lambda m: (_ for _ in ()).throw(
+        ValueError(m))
+    tname = os.path.basename(a.template)
+
+    code, prose, jsonl = [], [], []
+    for root, dirs, files in os.walk(a.repo):
+        dirs[:] = [d for d in dirs if not d.startswith(".")]
+        for f in sorted(files):
+            p = os.path.relpath(os.path.join(root, f), a.repo)
+            if DOC_DIR in p and f.endswith(".md"):
+                prose.append(p)
+            elif p.startswith(DOC_DIR):
+                pass
+            elif f.endswith(CODE_EXT):
+                code.append(p)
+            elif f.endswith(JSONL_EXTS) and f != "package-lock.json":
+                jsonl.append(p)
+    # file-level split: bucket 0 -> KLD pool, 1..4 -> calibration pool
+    pools: dict[str, tuple[list, list]] = {}
+    for cls, lst in (("code", code), ("prose", prose), ("tools", jsonl)):
+        pools[cls] = ([p for p in lst if fnv(p) % 5 == 0],
+                      [p for p in lst if fnv(p) % 5 != 0])
+
+    def text_of(p: str) -> str:
+        path = os.path.join(a.repo, p)
+        raw = open(path, encoding="utf-8", errors="replace").read()
+        if p.endswith(JSONL_EXTS):
+            lines = [json.loads(l)["content"] for l in raw.splitlines()
+                     if l.strip() and "content" in l]
+            raw = "\n".join(lines)
+        toks = tokenizer.encode(raw, add_special_tokens=False).ids
+        return tokenizer.decode(toks[:2400], skip_special_tokens=True)
+
+    # fragment pools per class, corpus-split first
+    frags: dict[str, dict[str, list[str]]] = {}
+    for cls in pools:
+        kld_files, cal_files = pools[cls]
+        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"])]
+        frags[cls] = fr
+
+    def emit_calib() -> tuple[int, list[str]]:
+        targets = {"code": .45, "prose": .30, "tools": .25}
+        rows, used = [], []
+        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:
+                frag = pool[i % len(pool)]
+                n = len(tokenizer.encode(frag, add_special_tokens=False).ids)
+                if nbuf + n > a.chunk_tokens and buf:
+                    rows.append({"prompt": "\n".join(buf)})
+                    used.extend(set())
+                    got += nbuf
+                    buf, nbuf = [], 0
+                buf.append(frag)
+                nbuf += n
+                i += 1
+            if buf:
+                rows.append({"prompt": "\n".join(buf)})
+                got += nbuf
+        rng.shuffle(rows)
+        return len(rows), rows
+
+    def emit_kld() -> tuple[int, list[str]]:
+        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):
+                cls = classes[(len(rows) + t) % 3]
+                pool = frags[cls]["kld"]
+                if not pool:
+                    continue
+                frag = pool[i[cls] % len(pool)]
+                i[cls] += 2
+                msgs.append({"role": "user",
+                             "content": "Summarize and explain this "
+                                        f"engineering material:\n{frag}"})
+                msgs.append({"role": "assistant",
+                             "content": frag[:4000]})
+                srcs.append(cls)
+            rendered, n = render(env, tokenizer, tname, msgs)
+            if n <= a.score_tail:
+                continue
+            rows.append({"prompt": rendered, "score_from": n - a.score_tail,
+                         "source": "+".join(sorted(set(srcs)))})
+            total += n
+        return len(rows), rows
+
+    os.makedirs(a.out_dir, exist_ok=True)
+    nc, rows = emit_calib()
+    with open(os.path.join(a.out_dir, "calib.jsonl"), "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+    nk, krows = emit_kld()
+    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)
+    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,
+        "calib_mix": {"code": .45, "prose": .30, "tools": .25},
+    }
+    with open(os.path.join(a.out_dir, "manifest.json"), "w") as f:
+        json.dump(manifest, f, indent=1)
+    print(json.dumps(manifest, indent=1))
+
+
+if __name__ == "__main__":
+    main()