| """rsi/selfprove.py — self-proof + repair for generated output (session v8 proofer). |
| |
| Drives the denoising expert (trained on <<<NOISY ...>>> / <<<CLEAN>>> records from |
| build_proof_u8.py) with an iterative masked-span refinement loop — the discrete |
| diffusion inference half: per round, the worst-CE spans are masked out and |
| re-sampled from the model's reconstruction; repeated over a decaying top-k schedule. |
| |
| Causal ctx=256 trunk ⇒ span-local edits with rolling context (chunking mirrors the |
| bundle brain's continuation_logprob): the prompt is always truncated to the last |
| cfg.ctx bytes, so the NOISY record is the local window being repaired. |
| |
| CLI: |
| python3 rsi/selfprove.py --model experts/cand_v8/proof.pt --in out.html \ |
| --report report.json --repaired fixed.html [--rounds 3] |
| |
| Exit 0 always — this is a measurement tool, not a gate. |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| from html.parser import HTMLParser |
|
|
| import numpy as np |
| import torch |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
|
|
| from gwlm_v6 import GoldWormV6, GoldWormV6Config |
| from gwlm_v7 import GoldWormV7, GoldWormV7Config |
|
|
| LN2 = float(np.log(2.0)) |
| NOISY_H = lambda n: b"<<<NOISY " + str(n).encode() + b">>>\n" |
| CLEAN_H = b"\n<<<CLEAN>>>\n" |
| END_MARK = b"<<<END>>>" |
|
|
| VOID = {"AREA", "BASE", "BR", "COL", "EMBED", "HR", "IMG", "INPUT", "LINK", |
| "META", "PARAM", "SOURCE", "TRACK", "WBR"} |
|
|
|
|
| def load(path, dev="cuda"): |
| """Family dispatch mirrors moe_gate.upcycle_experts (uor_tau -> v7 else v6).""" |
| ck = torch.load(path, map_location=dev, weights_only=False) |
| cfg = ck["cfg"] |
| if "uor_tau" in cfg: |
| m = GoldWormV7(GoldWormV7Config(**cfg)).to(dev) |
| else: |
| m = GoldWormV6(GoldWormV6Config(**cfg)).to(dev) |
| m.load_state_dict(ck["model_state_dict"]) |
| m.eval() |
| return m |
|
|
|
|
| @torch.no_grad() |
| def score(model, text, dev="cuda", W=256, stride=128): |
| """Rolling-window mean next-byte CE (bits). Returns (bpc_mean, spans) with |
| spans = [(offset, bpc)] sorted worst-first. Windows shorter than 8 B skipped.""" |
| if len(text) < 16: |
| return 0.0, [] |
| offs = list(range(0, max(1, len(text) - W), stride)) |
| if not offs or offs[-1] + W < len(text): |
| offs.append(max(0, len(text) - W)) |
| spans = [] |
| tot = n = 0.0 |
| for off in offs: |
| w = text[off:off + W] |
| if len(w) < 8: |
| continue |
| b = np.frombuffer(w, dtype=np.uint8).astype(np.int64) |
| x = torch.from_numpy(b[:-1]).long().unsqueeze(0).to(dev) |
| y = torch.from_numpy(b[1:]).long().unsqueeze(0).to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): |
| _, loss = model(x, targets=y) |
| bpc = float(loss.item()) / LN2 |
| spans.append((off, bpc)) |
| tot += float(loss.item()) * (len(w) - 1) |
| n += len(w) - 1 |
| spans.sort(key=lambda t: -t[1]) |
| return (tot / n / LN2 if n else 0.0), spans |
|
|
|
|
| @torch.no_grad() |
| def repair_span(model, ctx_prefix, span, dev="cuda", ctx=256): |
| """One denoising shot: NOISY record prompt -> greedy reconstruction.""" |
| prompt = ctx_prefix[-128:] + NOISY_H(len(span)) + span + CLEAN_H |
| toks = list(prompt) |
| gen = bytearray() |
| budget = int(len(span) * 1.2) + 16 |
| for _ in range(budget): |
| x = torch.tensor([toks[-ctx:]], dtype=torch.long, device=dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): |
| logits = model(x) |
| nxt = int(logits[:, -1].float().argmax(-1).item()) |
| gen.append(nxt) |
| toks.append(nxt) |
| if gen.endswith(END_MARK): |
| gen = gen[:-len(END_MARK)] |
| break |
| return bytes(gen) |
|
|
|
|
| class TagBalance(HTMLParser): |
| """Minimal strict-ish checker: every non-void open must close in LIFO order.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.stack = [] |
| self.mismatch_off = None |
|
|
| def handle_starttag(self, tag, attrs): |
| if tag.upper() not in VOID: |
| self.stack.append(tag) |
|
|
| def handle_startendtag(self, tag, attrs): |
| pass |
|
|
| def handle_endtag(self, tag): |
| if tag.upper() in VOID: |
| return |
| line, col = self.getpos() |
| if self.stack and self.stack[-1] == tag: |
| self.stack.pop() |
| elif tag in self.stack: |
| if self.mismatch_off is None: |
| self.mismatch_off = (line, col) |
| while self.stack and self.stack[-1] != tag: |
| self.stack.pop() |
| if self.stack: |
| self.stack.pop() |
| else: |
| if self.mismatch_off is None: |
| self.mismatch_off = (line, col) |
|
|
|
|
| def html_validate(text): |
| """(ok, mismatch_offset) — ok = balanced stack and no mismatch seen.""" |
| p = TagBalance() |
| try: |
| p.feed(text.decode("utf-8", "replace")) |
| p.close() |
| except Exception as e: |
| return False, str(e)[:80] |
| ok = (p.mismatch_off is None) and not p.stack |
| return ok, (None if p.mismatch_off is None else str(p.mismatch_off)) |
|
|
|
|
| def looks_html(text): |
| head = text.lstrip()[:64].lower() |
| return head.startswith(b"<!doctype") or head.startswith(b"<html") |
|
|
|
|
| def prove(model, text, dev="cuda", rounds=3, topk=(4, 2, 1), W=256): |
| """Iterative masked-span refinement. Returns (repaired_bytes, report_dict).""" |
| html_mode = looks_html(text) |
| bpc0, _ = score(model, text, dev, W=W) |
| html_ok0, _ = html_validate(text) if html_mode else (None, None) |
| cur = text |
| edits = 0 |
| worst = [] |
| rounds_applied = 0 |
| for r in range(rounds): |
| k = topk[min(r, len(topk) - 1)] |
| bpc_mean, spans = score(model, cur, dev, W=W) |
| if not spans: |
| break |
| worst = spans[:1] |
| floor = bpc_mean * 2.0 |
| chosen = [] |
| for off, bpc in spans: |
| if bpc <= floor: |
| continue |
| if len(chosen) >= k: |
| break |
| if all(abs(off - o) >= W for o, _ in chosen): |
| chosen.append((off, bpc)) |
| if not chosen: |
| break |
| rounds_applied += 1 |
| for off, _bpc in sorted(chosen, key=lambda t: -t[0]): |
| span = cur[off:off + W] |
| prefix = cur[max(0, off - 128):off] |
| new = repair_span(model, prefix, span, dev, ctx=model.cfg.ctx) |
| if new and len(new) <= int(len(span) * 1.2) + 16: |
| cur = cur[:off] + new + cur[off + len(span):] |
| edits += 1 |
| bpc1, _ = score(model, cur, dev, W=W) |
| html_ok1, _ = html_validate(cur) if html_mode else (None, None) |
| report = {"bpc_before": round(bpc0, 4), "bpc_after": round(bpc1, 4), |
| "worst_spans": [{"off": o, "bpc": round(b, 4)} for o, b in worst], |
| "rounds_applied": rounds_applied, "edits": edits, |
| "html_ok_before": html_ok0, "html_ok_after": html_ok1} |
| return cur, report |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--in", dest="inp", required=True) |
| ap.add_argument("--report", default=None) |
| ap.add_argument("--repaired", default=None) |
| ap.add_argument("--rounds", type=int, default=3) |
| a = ap.parse_args() |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| model = load(a.model, dev) |
| with open(a.inp, "rb") as f: |
| text = f.read() |
| repaired, report = prove(model, text, dev, rounds=a.rounds) |
| report = {"file": a.inp, **report} |
| if a.repaired: |
| with open(a.repaired, "wb") as f: |
| f.write(repaired) |
| if a.report: |
| with open(a.report, "w") as f: |
| json.dump(report, f, indent=1) |
| print(json.dumps(report, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
