263 lines
11 KiB
Python
Executable File
263 lines
11 KiB
Python
Executable File
#!/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()
|