#!/usr/bin/env python3 """ Builds `data/agent_persona_v1/strata_parity_refs.pt` with full cryptographic provenance and multi-scale `strata_window_cli` reference logits across: 1. `short_223_eval_tf_chat_json` (T=223, n_eval=64, QSA Indexer off) 2. `long_2255_qsa_indexer_ftpo_6a` (T=2255, n_eval=64, QSA Indexer active) 3. `long_9533_qsa_indexer_supervised_span` (T=9533, n_eval=64, QSA Indexer active on supervised span) 4. `long_15081_qsa_indexer_supervised_span` (T=15081, n_eval=64, QSA Indexer active on supervised span) 5. `max_30554_qsa_indexer_supervised_span` (T=30554, n_eval=64, max train length, QSA Indexer active) Also quantifies `C5` (`STRATA_METAL_MOCE_BONUS=0.85` vs `0.0` KLD & Top-1 agreement on T=9533). """ from __future__ import annotations import gzip import hashlib import json import os import subprocess import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F ROOT = Path(__file__).resolve().parent.parent DATA_DIR = ROOT / "data" / "agent_persona_v1" PACK_PATH = DATA_DIR / "colab_a100_tokenized_pack.json.gz" OUT_PATH = DATA_DIR / "strata_parity_refs.pt" def file_sha256_prefix(p: Path, max_bytes: int = 16 * 1024 * 1024) -> str: h = hashlib.sha256() with p.open("rb") as f: h.update(f.read(max_bytes)) return h.hexdigest() def run_strata_logits( strata_cli: Path, gguf_path: Path, ple_path: Path, toks: list[int], n_eval: int, ctx: int, moce_bonus: str = "0.0", ) -> tuple[torch.Tensor, str]: tfile = Path("/tmp/strata_ref_toks.txt") lbin = Path("/tmp/strata_ref_logits.bin") tfile.write_text(" ".join(str(x) for x in toks), encoding="utf-8") lbin.unlink(missing_ok=True) n_p = len(toks) - n_eval env = dict(os.environ, STRATA_METAL_MOCE_BONUS=moce_bonus, STRATA_METAL_PROMPT_MOCE="0.0") cmd = [ str(strata_cli), "--metal-gguf", str(gguf_path), "--ple-gguf", str(ple_path), "--tokens-file", str(tfile), "--n-expert", "512", "--ctx", str(ctx), "--n-prompt", str(n_p), "--n-gen", str(n_eval), "--logits-out", str(lbin), ] proc = subprocess.run(cmd, check=True, env=env, capture_output=True, text=True) lg = torch.from_numpy(np.fromfile(lbin, dtype=np.float32).reshape(n_eval, 248320)) lbin.unlink(missing_ok=True) tfile.unlink(missing_ok=True) return lg, (proc.stdout + "\n" + proc.stderr) def main() -> int: strata_cli = ROOT / "vendor" / "strata" / "build-metal" / "strata_window_cli" gguf_path = ROOT / "models" / "K512_iq2_xs" / "K512-iq2_xs-metal.gguf" ple_path = ROOT / "models" / "K256_reap_l2_sum" / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf" pack = json.load(gzip.open(PACK_PATH, "rt", encoding="utf-8")) all_wins = pack["sft_train"] + pack["sft_val_id"] + pack["sft_val_ood_repo"] p223_file = ROOT / "data" / "m5_val_prompts" / "eval_tf_chat_json.txt" toks_223 = [int(x) for x in p223_file.read_text(encoding="utf-8").split()] p2255 = next(p for p in pack["ftpo_train"] if len(p["chosen_ids"]) == 2255) toks_2255 = p2255["chosen_ids"] w9533 = next(w for w in all_wins if w["id"] == "b1_local_no_per_01a1361b_168319") toks_9533 = w9533["input_ids"][:9533] w15081 = next(w for w in all_wins if w["id"] == "b3_astaroth_01a11d7d_94856") toks_15081 = w15081["input_ids"][:15081] w_max = max( pack["sft_train"], key=lambda w: max(e for _, e in w["supervised_spans"]), ) max_end = max(e for _, e in w_max["supervised_spans"]) toks_30554 = w_max["input_ids"][:max_end] specs = [ ("short_223_eval_tf_chat_json", toks_223, 64, 2048), ("long_2255_qsa_indexer_ftpo_6a", toks_2255, 64, 4096), ("long_9533_qsa_indexer_supervised_span", toks_9533, 64, 10240), ("long_15081_qsa_indexer_supervised_span", toks_15081, 64, 16384), ("max_30554_qsa_indexer_supervised_span", toks_30554, 64, 32768), ] refs = [] lg_9533_moce0 = None for name, toks, n_eval, ctx in specs: t0 = time.time() lg, log_txt = run_strata_logits(strata_cli, gguf_path, ple_path, toks, n_eval, ctx, moce_bonus="0.0") tgt = torch.tensor(toks[-n_eval:], dtype=torch.long) ce = float(F.cross_entropy(lg, tgt).item()) print(f"Built {name:42s}: T={len(toks):5d}, n_eval={n_eval}, strata_ce={ce:.4f} ({time.time()-t0:.2f}s)", flush=True) if name == "long_9533_qsa_indexer_supervised_span": lg_9533_moce0 = lg refs.append({ "name": name, "input_ids": toks, "n_eval": n_eval, "ctx": ctx, "strata_ce_nats": round(ce, 4), "strata_logits_f16": lg.to(torch.float16), }) # C5: Measure MoCE=0.85 vs MoCE=0.0 on the 9,533-token supervised window lg_9533_moce085, _ = run_strata_logits(strata_cli, gguf_path, ple_path, toks_9533, 64, 10240, moce_bonus="0.85") tgt_9533 = torch.tensor(toks_9533[-64:], dtype=torch.long) ce_0 = float(F.cross_entropy(lg_9533_moce0, tgt_9533).item()) ce_085 = float(F.cross_entropy(lg_9533_moce085, tgt_9533).item()) p0 = F.softmax(lg_9533_moce0.to(torch.float64), dim=-1) lp0 = F.log_softmax(lg_9533_moce0.to(torch.float64), dim=-1) lp085 = F.log_softmax(lg_9533_moce085.to(torch.float64), dim=-1) kld_moce = float((p0 * (lp0 - lp085)).sum(dim=-1).mean().item()) top1_moce = float((lg_9533_moce0.argmax(dim=-1) == lg_9533_moce085.argmax(dim=-1)).float().mean().item()) moce_comparison = { "window": "long_9533_qsa_indexer_supervised_span", "n_tokens": len(toks_9533), "n_eval_tokens": 64, "ce_moce_0_0": round(ce_0, 4), "ce_moce_0_85": round(ce_085, 4), "kld_moce_0_vs_085": round(kld_moce, 5), "top1_agreement_moce_0_vs_085": round(top1_moce, 4), } print(f"C5 MoCE=0.0 vs MoCE=0.85 comparison on T=9533: {json.dumps(moce_comparison)}", flush=True) bundle = { "provenance": { "generator_script": "scripts/build_strata_parity_refs.py", "gguf_file": gguf_path.name, "gguf_size_bytes": gguf_path.stat().st_size, "gguf_head16mb_sha256": file_sha256_prefix(gguf_path), "ple_gguf_file": ple_path.name, "ple_head16mb_sha256": file_sha256_prefix(ple_path), "strata_cli_sha256": file_sha256_prefix(strata_cli), "env": {"STRATA_METAL_MOCE_BONUS": "0.0", "STRATA_METAL_PROMPT_MOCE": "0.0"}, }, "moce_0_vs_085_audit": moce_comparison, "refs": refs, } torch.save(bundle, OUT_PATH) print(f"Saved {OUT_PATH} ({OUT_PATH.stat().st_size / (1024*1024):.2f} MiB) with {len(refs)} multi-scale refs (223..30554 tokens)!") return 0 if __name__ == "__main__": raise SystemExit(main())