#!/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())