import json, os, re, subprocess, time
from pathlib import Path
PPL_BIN = "/content/llama.cpp/build/bin/llama-perplexity"
PPL_CPP = Path("/content/llama.cpp/tools/perplexity/perplexity.cpp")
assert "common_tokenize(ctx, params.prompt, true, true)" in PPL_CPP.read_text(), "perplexity.cpp missing parse_special=true patch!"
SHARD1 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf"
SHARD2 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00002-of-00002.gguf"
EVAL_DOMAINS = {
"ops_tools": ("/content/eval_ops_tools.txt", "/content/base_logits_ops_tools.bin"),
"chat_json": ("/content/eval_chat_json.txt", "/content/base_logits_chat_json.bin"),
"ja": ("/content/eval_ja.txt", "/content/base_logits_ja.bin"),
}
def grab(pattern: str, text: str) -> float:
m = re.search(pattern, text)
if not m:
raise RuntimeError(f"Failed to match {pattern} in output:\n{text[-1500:]}")
return float(m.group(1))
results_path = Path("/content/phase1c_kld_results.json")
all_res = json.loads(results_path.read_text()) if results_path.exists() else {}
if "K512_unpruned" not in all_res:
print("=== Step 1: Evaluating K512_unpruned baseline & saving teacher logits ===", flush=True)
base_res = {}
for dom, (fpath, lfile) in EVAL_DOMAINS.items():
t0 = time.monotonic()
cmd = [
PPL_BIN, "-m", SHARD1, "-f", fpath,
"-c", "2048", "-b", "512", "-ngl", "99",
"--chunks", "4",
"--save-all-logits", lfile,
]
p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
if p.returncode != 0:
raise RuntimeError(f"Baseline PPL failed on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}")
dt = time.monotonic() - t0
ppl = grab(r"Final estimate: PPL = ([0-9.]+)", p.stdout)
ppl_se = grab(r"Final estimate: PPL = [0-9.]+\s*\+/-\s*([0-9.]+)", p.stdout)
base_res[dom] = {"ppl": ppl, "ppl_se": ppl_se, "sec": round(dt, 1)}
print(f"[K512_unpruned] {dom}: PPL = {ppl:.4f} +/- {ppl_se:.5f} ({dt:.1f}s)", flush=True)
all_res["K512_unpruned"] = base_res
results_path.write_text(json.dumps(all_res, indent=2))
CANDIDATES = [
("K320_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k320.json"),
("K320_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k320.json"),
("K256_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k256.json"),
("K256_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k256.json"),
("K256_reap_shrink", "/content/keep_maps/keep_reap_shrink_k256.json"),
("K224_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k224.json"),
]
os.makedirs("/dev/shm/pruned", exist_ok=True)
for label, km_path in CANDIDATES:
if label in all_res:
continue
out1 = f"/dev/shm/pruned/{label}-00001-of-00002.gguf"
out2 = f"/dev/shm/pruned/{label}-00002-of-00002.gguf"
t0 = time.monotonic()
print(f"=== Slicing {label} ===", flush=True)
subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD1, "--out-gguf", out1, "--keep-json", km_path], check=True)
subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD2, "--out-gguf", out2, "--keep-json", km_path], check=True)
print(f"Sliced {label} in {time.monotonic()-t0:.1f}s", flush=True)
cand_res = {}
for dom, (fpath, lfile) in EVAL_DOMAINS.items():
t1 = time.monotonic()
cmd = [
PPL_BIN, "-m", out1, "-f", fpath,
"-c", "2048", "-b", "512", "-ngl", "99",
"--chunks", "4",
"--kl-divergence", "--kl-divergence-base", lfile,
]
p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
if p.returncode != 0:
raise RuntimeError(f"KLD eval failed for {label} on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}")
dt = time.monotonic() - t1
ppl_q = grab(r"Mean PPL\(Q\)\s*:\s*([0-9.]+)", p.stdout)
ppl_base = grab(r"Mean PPL\(base\)\s*:\s*([0-9.]+)", p.stdout)
ppl_diff = grab(r"Mean PPL\(Q\)-PPL\(base\)\s*:\s*([-+0-9.]+)", p.stdout)
mean_kld = grab(r"Mean\s+KLD:\s*([-+0-9.]+)", p.stdout)
mean_kld_se = grab(r"Mean\s+KLD:\s*[-+0-9.]+\s*±\s*([0-9.]+)", p.stdout)
med_kld = grab(r"Median\s+KLD:\s*([-+0-9.]+)", p.stdout)
p90_kld = grab(r"90\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout)
p99_kld = grab(r"99\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout)
rms_dp = grab(r"RMS Δp\s*:\s*([0-9.]+)", p.stdout)
same_top = grab(r"Same top p:\s*([0-9.]+)", p.stdout)
base_fp32 = all_res["K512_unpruned"][dom]["ppl"]
cand_res[dom] = {
"ppl": ppl_q,
"ppl_base": ppl_base,
"delta_ppl": ppl_diff,
"exact_delta_ppl": round(ppl_q - base_fp32, 4),
"mean_kld": mean_kld,
"mean_kld_se": mean_kld_se,
"median_kld": med_kld,
"p90_kld": p90_kld,
"p99_kld": p99_kld,
"rms_delta_p_pct": rms_dp,
"same_top1_pct": same_top,
"sec": round(dt, 1),
}
print(f"[{label}] {dom}: PPL={ppl_q:.4f} (Δ={ppl_diff:+.4f}) | KLD={mean_kld:.5f}±{mean_kld_se:.5f} | Top1={same_top:.2f}% ({dt:.1f}s)", flush=True)
all_res[label] = cand_res
results_path.write_text(json.dumps(all_res, indent=2))
os.remove(out1)
os.remove(out2)
print("=== ALL PHASE 1C TIER-1 KLD + TOP-1 + PPL EVALS COMPLETE! ===", flush=True)