onf-c1/scripts/train_task_belief.py

#!/usr/bin/env python
"""Fit the task belief and write g_belief_head.npz beside a graph.

Two phases, both in onf.graph.build.belief_fit: extract the frozen retriever's two channels once
into belief_evidence.npz (--reuse-evidence replays it), then train the 163-parameter fusion on that
cache. The whole second phase is seconds; the first is one retriever forward per check.

The pre-registered gate this prints: with the instruction DROPPED, the accumulated belief must
clearly exceed the text encoder's own 64.2% task identification by ~5 checks -- that is the claim
that the deployed system survives a rewording it cannot parse.

Usage:

    OMP_NUM_THREADS=4 PYTHONPATH=src CUDA_VISIBLE_DEVICES=2 python scripts/train_task_belief.py \\
        --artifacts outputs/long/stage2_alpha_v6b/artifacts
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Sequence

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(REPO_ROOT / "src"))

from onf.config import GraphConfig                                          # noqa: E402
from onf.graph.build.belief_fit import (                                    # noqa: E402
    CHUNK_K, EVIDENCE_NPZ, N_CHECKS, REGIMES, BeliefFitter, EvidenceCache, FitReport,
    TaskBeliefProbe, extract_evidence,
)
from onf.graph.lang import TaskEmbeddings                                   # noqa: E402
from onf.graph.net.language import LanguageHead                             # noqa: E402

CHECKPOINTS = (1, 2, 4, 8)         # accumulated checks the accuracy table is read at
ENCODER_TOP1 = 64.2                # the bare instruction encoder, on 500 real rewordings
BLEND_SEM_SPREAD = 1.40            # the soft prior blend_sem deployed; it lost Objects_Layout 6.1pp


def accuracy_table(report: FitReport, label: str) -> str:
    """The accuracy-vs-checks table, one block per text regime.

    Args:
        report: The measurement.
        label: What the rows are, for the header.

    Returns:
        The rendered table.
    """
    checks = sorted(next(iter(report.accuracy.values())))
    head = f"{'channel':<12}" + "".join(f"{f'{m} chk':>10}" for m in checks)
    out = [f"{label}, n={report.n_val_rows} held-out rows", ""]
    for regime in REGIMES:
        out += [f"  [{regime}]", "  " + head, "  " + "-" * (12 + 10 * len(checks))]
        for channel in ("s_text", "s_ret", "s_beh", "fusion"):
            row = report.accuracy[(regime, channel)]
            out.append(
                f"  {channel:<12}" + "".join(f"{100 * row[m]:>9.1f}%" for m in checks)
            )
        out.append("")
    return "\n".join(out)


def gate_lines(report: FitReport) -> list[str]:
    """The pre-registered gate, and the seed spread Objects_Layout turns on.

    Args:
        report: The measurement.

    Returns:
        The verdict lines.
    """
    drop = report.accuracy[("drop_text", "fusion")]
    at = max(m for m in drop if m <= 4)
    passed = 100 * drop[at] > ENCODER_TOP1
    return [
        f"GATE  drop_text fusion @ {at} checks = {100 * drop[at]:.1f}%  vs the bare encoder's "
        f"{ENCODER_TOP1}%  -> {'PASS' if passed else 'FAIL'}",
        "seed factor after 4 checks (correct task / best wrong task), median: " + "  ".join(
            f"{regime} {report.spread[regime]:.4g}x" for regime in REGIMES
        ) + f"   [blend_sem's soft prior managed {BLEND_SEM_SPREAD:.2f}x on Objects_Layout]",
        "  ... over the rows it gets WRONG, where below 1 means the CORRECT lane is suppressed: "
        + "  ".join(f"{regime} {report.spread_wrong[regime]:.4g}x" for regime in REGIMES),
    ]


def head_lines(report: FitReport) -> list[str]:
    """The fitted scalars and which channel the fusion leans on.

    Args:
        report: The measurement.

    Returns:
        The lines.
    """
    return [
        f"fitted: gain={report.gain:.4f} nats/check  decay={report.decay:.4f}  "
        f"saturation={report.saturation:.3f} nats",
        "first-layer weight share: " + "  ".join(
            f"{name} {100 * share:.1f}%" for name, share in report.channel_weight.items()
        ),
        "s_ret one-check accuracy by split: " + "  ".join(
            f"{name} {100 * value:.1f}%" for name, value in report.retrieval_split.items()
        ),
    ]


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Parse the command line.

    Args:
        argv: Argument vector; None reads sys.argv.

    Returns:
        The namespace.
    """
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--artifacts", default="outputs/long/stage2_alpha_v6b/artifacts")
    ap.add_argument("--suite", default="long", help="suite whose ONFField weights the windows")
    ap.add_argument("--rows-per-demo", type=int, default=4, help="check rows drawn per demo")
    ap.add_argument("--n-checks", type=int, default=N_CHECKS, help="checks per row")
    ap.add_argument("--epochs", type=int, default=200)
    ap.add_argument("--batch", type=int, default=32, help="retriever forward-pass batch size")
    ap.add_argument("--device", default="cuda", help="torch device")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--reuse-evidence", action="store_true",
        help=f"replay {EVIDENCE_NPZ} instead of re-running the retriever",
    )
    ap.add_argument("--no-save", action="store_true", help="measure without writing the head")
    return ap.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    """Extract (or reload) the evidence, fit the head, print the table and save.

    Args:
        argv: Argument vector; None reads sys.argv.

    Returns:
        Process exit status.
    """
    args = parse_args(argv)
    artifacts = Path(args.artifacts)

    if args.reuse_evidence:
        cache = EvidenceCache.load(artifacts)
    else:
        # from_env, not a bare GraphConfig: a bare one would silently reset every other GR_* knob
        # away from whatever the head was trained and is deployed under.
        cfg = GraphConfig.from_env()
        cfg.hist, cfg.device = CHUNK_K, args.device
        probe = TaskBeliefProbe(artifacts, cfg, args.suite, args.device)
        language = TaskEmbeddings.load(artifacts).check_against(probe.nodes.task_names)
        cache = extract_evidence(
            probe, language, args.rows_per_demo, args.n_checks,
            np.random.RandomState(args.seed), batch=args.batch,
        )
        print(f"[belief.fit] wrote {cache.save(artifacts)}", file=sys.stderr)
    print(f"[belief.fit] {cache.summary()}", file=sys.stderr)

    fitter = BeliefFitter(
        cache, LanguageHead.load(artifacts, device=args.device), device=args.device,
        seed=args.seed,
    )
    fitter.run(args.epochs)

    for rank_zero, label in ((True, "EPISODE-START rows"), (False, "ALL start offsets")):
        report = fitter.report(CHECKPOINTS, rank_zero=rank_zero)
        print(f"\n=== {label} ===")
        print(accuracy_table(report, label))
        # The fitted scalars describe the head, not the row set, so they are printed once.
        print("\n".join((head_lines(report) if rank_zero else []) + gate_lines(report)))

    if not args.no_save:
        print(f"\n[belief.fit] wrote {fitter.head.save(artifacts)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())