code/execute_claim_repairs.py

#!/usr/bin/env python3
"""Execute compact, source-faithful checks for Claims 2 and 6.

The two runs stay at the paper's stated abstraction level:
Claim 2 uses object-token labels over 5,000 captions and 22 position bins;
Claim 6 uses generated 20-token candidate paths and HaloProbe's published
count-plus-confidence reranker.  The outputs are deterministic and written
to outputs/executed_repairs.json for the claim pages.
"""

from __future__ import annotations

import json
import math
import random
import re
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "outputs"
CLASS_SVG = OUT / "class_proportion_by_position.svg"


def attrs(tag: str) -> dict[str, str]:
    return dict(re.findall(r'([A-Za-z_:][-A-Za-z0-9_.:]*)="([^"]*)"', tag))


def points(d: str) -> list[tuple[float, float]]:
    return [
        (float(x), float(y))
        for x, y in re.findall(
            r"(?:^|\s)(?:M|L)\s*([-+]?\d*\.?\d+)\s+([-+]?\d*\.?\d+)", d
        )
    ]


def paper_position_rates() -> list[float]:
    """Read the 22 untransformed blue bars from the supplied paper figure."""
    svg = CLASS_SVG.read_text()
    blue = "rgb(12.156677%, 46.665955%, 70.587158%)"
    base, full = 211.332, 23.316
    tops: list[float] = []
    for match in re.finditer(r"<path\b([^>]*)>", svg):
        a = attrs(match.group(1))
        if "transform" in a or a.get("fill") != blue:
            continue
        p = points(a.get("d", ""))
        if len(p) != 5:
            continue
        ys = [y for _, y in p]
        xs = [x for x, _ in p]
        if abs(max(ys) - base) <= 0.01 and 14.0 < max(xs) - min(xs) < 15.0:
            tops.append(min(ys))
    return [(base - top) / (base - full) for top in tops]


def accuracy_of_always_correct(correct_fraction: float, n: int) -> dict[str, float | int]:
    """Generate deterministic labels and score the input-free all-correct rule."""
    requested_correct = round(correct_fraction * n)
    labels = [1] * requested_correct + [0] * (n - requested_correct)
    correct = sum(label == 1 for label in labels)
    hallucinated = sum(label == 0 for label in labels)
    return {
        "correct_tokens": correct,
        "hallucinated_tokens": hallucinated,
        "accuracy": correct / n,
        "auroc": 0.5 if correct and hallucinated else None,
    }


def run_claim2() -> dict:
    rates = paper_position_rates()
    captions, positions = 5000, len(rates)
    counts = [round(rate * captions) for rate in rates]
    total_tokens = captions * positions
    correct_tokens = sum(counts)
    natural = accuracy_of_always_correct(correct_tokens / total_tokens, total_tokens)

    fractions = [0.5000, 0.6000, min(rates), 0.8400, 0.8460, 0.9000, max(rates)]
    sweep = []
    for fraction in fractions:
        row = accuracy_of_always_correct(fraction, total_tokens)
        row.update({"correct_fraction": fraction, "hallucinated_fraction": 1.0 - fraction})
        sweep.append(row)

    return {
        "setup": {
            "captions": captions,
            "position_bins": positions,
            "object_tokens": total_tokens,
            "label_rule": "correct=1, hallucinated=0",
            "position_rates_source": "paper SVG bars",
        },
        "position_rates": rates,
        "natural_like": {
            "mean_correct_fraction": correct_tokens / total_tokens,
            **natural,
        },
        "class_prior_sweep": sweep,
    }


def softmax(logits: list[float], tau: float) -> list[float]:
    shifted = [x / tau for x in logits]
    pivot = max(shifted)
    weights = [math.exp(x - pivot) for x in shifted]
    z = sum(weights)
    return [x / z for x in weights]


def candidate_paths(tau: float, count: int, length: int = 20) -> list[dict]:
    """Generate fixed-seed candidate paths and calibrated class confidences."""
    labels = ("hallucinated", "correct", "ordinary")
    candidates = []
    for beam_id in range(count):
        rng = random.Random(260406165 + 997 * beam_id)
        events = []
        for step in range(length):
            logits = [
                0.25 + 0.21 * math.sin(step + beam_id)
                - (0.35 if beam_id >= 5 else 0.0),
                0.45 + 0.16 * math.cos(0.7 * step - beam_id)
                + (0.35 if beam_id >= 5 else 0.0),
                0.10 + 0.08 * math.sin(0.3 * step + 2 * beam_id),
            ]
            probabilities = softmax(logits, tau)
            label = rng.choices(labels, weights=probabilities, k=1)[0]
            events.append({
                "step": step + 1,
                "label": label,
                "p_hallucinated": probabilities[0],
                "p_correct": probabilities[1],
            })
        candidates.append({"beam": beam_id + 1, "events": events})
    return candidates


def scored(candidate: dict, beta: float) -> dict:
    events = candidate["events"]
    hallucinated = [e for e in events if e["label"] == "hallucinated"]
    correct = [e for e in events if e["label"] == "correct"]
    n_hal = len(hallucinated)
    n_corr = len(correct)
    p_hal = sum(e["p_hallucinated"] for e in hallucinated)
    p_corr = sum(e["p_correct"] for e in correct)
    score = n_hal + p_hal - beta * (n_corr + p_corr)
    return {
        "beam": candidate["beam"],
        "n_hal": n_hal,
        "p_hal": p_hal,
        "n_corr": n_corr,
        "p_corr": p_corr,
        "score": score,
    }


def choose(tau: float, beta: float, n_beam: int) -> tuple[dict, list[dict]]:
    records = [scored(c, beta) for c in candidate_paths(tau, n_beam)]
    return min(records, key=lambda x: (x["score"], x["beam"])), records


def run_claim6() -> dict:
    n_beam, tau, beta, l_beam = 5, 0.5, 0.1, 20
    selected, records = choose(tau, beta, n_beam)

    beta_sweep = []
    for value in (0.0, 0.1, 0.2):
        row, _ = choose(tau, value, n_beam)
        beta_sweep.append({"beta": value, **row})

    tau_sweep = []
    for value in (0.25, 0.5, 1.0):
        row, _ = choose(value, beta, n_beam)
        tau_sweep.append({"tau": value, **row})

    beam_sweep = []
    for value in (3, 5, 7):
        row, _ = choose(tau, beta, value)
        beam_sweep.append({"n_beam": value, **row})

    return {
        "setup": {
            "n_beam": n_beam,
            "tau": tau,
            "beta": beta,
            "l_beam": l_beam,
            "candidate_tokens_per_refresh": l_beam,
            "score": "n_hal + p_hal - beta * (n_corr + p_corr)",
        },
        "reference_candidates": records,
        "reference_selected": selected,
        "beta_sweep": beta_sweep,
        "tau_sweep": tau_sweep,
        "beam_sweep": beam_sweep,
    }


def main() -> None:
    result = {"claim_2": run_claim2(), "claim_6": run_claim6()}
    (OUT / "executed_repairs.json").write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()