arborist/bench/teacher_judge.py

"""Teacher-model judge — cross-judge model answers from a bench run.

Reads `<ts>.judge_input.jsonl` (emitted by cross_model_selfplay.py
`--teacher-model-judge`), grades each (question, answerer, answer) by
calling a *different* model to judge it, writes
`<ts>.verdicts.jsonl`, prints a correctness summary.

Cross-judging avoids self-defense bias: a model's natural tendency to
rationalize its own answer. Default with 2 models: hermes-answers get
judged by qwen, qwen-answers by hermes. Override with `--self-judge`
(every model judges its own answers — useful as a sanity check on the
cross-judge: if self-judge agrees with cross-judge, the verdict is
robust).

Verdict tokens: CORRECT | WRONG | PARTIAL | UNCERTAIN

Output (per row):

    {"question", "answerer", "audit_mode", "answer_text",
     "judge", "verdict", "rationale", "raw_reply", "_elapsed_s"}

Usage:

    python3 bench/teacher_judge.py \\
        --input bench/cross_model_results/20260531T150122Z.judge_input.jsonl
    # → bench/cross_model_results/20260531T150122Z.verdicts.jsonl

Cost: <1 cent for 150 rows on owned-hardware infrastructure (hermes
on 3090, qwen on 4090). For external paid APIs, audit the math first.
"""
from __future__ import annotations

import argparse
import datetime as _dt
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from collections import defaultdict
from pathlib import Path

# Reuse the prompt builder from cross_model_selfplay so iterating on the
# prompt template doesn't require re-running benches — teacher_judge.py
# regenerates a fresh prompt from (question, answer) at judge time.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from cross_model_selfplay import _judge_prompt  # noqa: E402

DEFAULT_MODELS = [
    ("hermes", "https://hermes.ai.unturf.com/v1",
     "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic"),
    ("qwen", "https://qwen.ai.unturf.com/v1",
     "Qwen3.6-27B-UD-Q4_K_XL.gguf"),
]

VERDICT_TOKENS = {"CORRECT", "WRONG", "PARTIAL", "UNCERTAIN"}
# Match VERDICT at the start of a line, allowing some leading whitespace and
# optional leading `**` or quote markers. Capture verdict and rationale.
_VERDICT_RE = re.compile(
    r"^\s*[*_\"`]*\s*(CORRECT|WRONG|PARTIAL|UNCERTAIN)\b[\s:_*\"`]*[:\-]?\s*(.*)$",
    re.IGNORECASE | re.MULTILINE,
)


def _parse_models(spec: str) -> list[tuple[str, str, str]]:
    """Same format as cross_model_selfplay.py: name=endpoint:model_id,..."""
    out = []
    for chunk in spec.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        if "=" not in chunk:
            raise SystemExit(
                f"--models entry must be name=endpoint:model_id, got: {chunk!r}"
            )
        name, rest = chunk.split("=", 1)
        endpoint, model_id = rest.rsplit(":", 1)
        out.append((name.strip(), endpoint.strip(), model_id.strip()))
    return out


def _post_completion(
    endpoint: str, model_id: str, prompt: str, *,
    timeout_s: int, max_tokens: int = 200,
) -> tuple[str, float]:
    """POST a chat-completion request, return (assistant_text, elapsed_s).

    Uses stdlib urllib — no third-party deps. Raises on transport error.

    Sends `chat_template_kwargs: {enable_thinking: false}` to disable
    reasoning mode on Qwen3-style endpoints (verified against llama.cpp's
    qwen.ai.unturf.com 2026-05-31 — content lands in `message.content`
    cleanly instead of being eaten by `reasoning_content`). Hermes / vLLM
    endpoints silently ignore the unknown kwarg.
    """
    url = endpoint.rstrip("/") + "/chat/completions"
    payload = json.dumps({
        "model": model_id,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.0,
        "chat_template_kwargs": {"enable_thinking": False},
    }).encode("utf-8")
    req = urllib.request.Request(
        url, data=payload,
        headers={"Content-Type": "application/json"},
    )
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout_s) as resp:
        body = resp.read()
    elapsed = round(time.time() - t0, 2)
    j = json.loads(body)
    choices = j.get("choices") or []
    if not choices:
        return "", elapsed
    msg = choices[0].get("message") or {}
    return (msg.get("content") or "").strip(), elapsed


def _parse_verdict(reply: str) -> tuple[str, str]:
    """Extract (verdict, rationale) from a teacher's reply.

    Looks for the first line containing one of the verdict tokens. Falls
    back to scanning the whole reply if no clean format match. Returns
    ('PARSE_ERROR', raw_reply) if no token found at all.
    """
    if not reply:
        return ("PARSE_ERROR", "(empty reply)")
    m = _VERDICT_RE.search(reply)
    if m:
        verdict = m.group(1).upper()
        rationale = (m.group(2) or "").strip().rstrip(".")
        if not rationale:
            # rationale might be on the next line
            lines = reply.splitlines()
            for i, L in enumerate(lines):
                if verdict in L.upper():
                    if i + 1 < len(lines):
                        rationale = lines[i + 1].strip()
                    break
        return (verdict, rationale[:300])
    # Fallback: just look for any verdict token anywhere
    for token in VERDICT_TOKENS:
        if re.search(rf"\b{token}\b", reply, re.IGNORECASE):
            return (token, reply[:300].strip())
    return ("PARSE_ERROR", reply[:300])


def _pick_judge(
    answerer: str, models: list[tuple[str, str, str]], self_judge: bool,
) -> tuple[str, str, str] | None:
    """Return (judge_name, endpoint, model_id) or None if no valid judge."""
    if self_judge:
        for m in models:
            if m[0] == answerer:
                return m
        return None
    # Cross-judge: first model whose name != answerer
    for m in models:
        if m[0] != answerer:
            return m
    return None


def _format_summary(verdict_rows: list[dict]) -> str:
    """Markdown summary keyed by answerer: verdict counts per model."""
    by_answerer: dict[str, list[dict]] = defaultdict(list)
    for v in verdict_rows:
        by_answerer[v["answerer"]].append(v)

    out = []
    out.append(f"# Teacher-model judge — {_dt.datetime.now(_dt.timezone.utc).isoformat()}")
    out.append("")
    out.append(f"Rows judged: {len(verdict_rows)}")
    out.append("")
    out.append("## Verdicts by answerer (cross-judged)")
    out.append("")
    out.append("| answerer | judge | CORRECT | WRONG | PARTIAL | UNCERTAIN | PARSE_ERROR | accuracy* |")
    out.append("|---|---|---:|---:|---:|---:|---:|---:|")
    for answerer in sorted(by_answerer):
        rows = by_answerer[answerer]
        judges = sorted({r["judge"] for r in rows})
        judge_label = "+".join(judges) if judges else "—"
        counts = defaultdict(int)
        for r in rows:
            counts[r["verdict"]] += 1
        n = len(rows)
        n_judged = n - counts["UNCERTAIN"] - counts["PARSE_ERROR"]
        acc = (counts["CORRECT"] / n_judged) if n_judged else 0.0
        out.append(
            f"| {answerer} | {judge_label} "
            f"| {counts['CORRECT']} | {counts['WRONG']} | {counts['PARTIAL']} "
            f"| {counts['UNCERTAIN']} | {counts['PARSE_ERROR']} "
            f"| {acc:.0%} |"
        )
    out.append("")
    out.append("\\*accuracy = CORRECT / (CORRECT + WRONG + PARTIAL); UNCERTAIN and PARSE_ERROR excluded from denominator.")
    out.append("")

    # False-STRICT highlight: WRONG verdicts on STRICT audits
    false_stricts = [
        r for r in verdict_rows
        if r["verdict"] == "WRONG" and r["audit_mode"] == "STRICT"
    ]
    if false_stricts:
        out.append("## False STRICTs detected (model audited STRICT, judge said WRONG)")
        out.append("")
        for r in false_stricts:
            ans_short = r["answer_text"][:140].replace("\n", " ")
            out.append(f"* **{r['answerer']} → Q: {r['question'][:80]}**")
            out.append(f"  * Judge ({r['judge']}): {r['rationale'][:200]}")
            out.append(f"  * Answer: {ans_short}…")
        out.append("")
    else:
        out.append("## False STRICTs detected")
        out.append("")
        out.append("None.")
        out.append("")

    return "\n".join(out)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--input", type=Path, required=True,
        help="path to <ts>.judge_input.jsonl from cross_model_selfplay.py",
    )
    ap.add_argument(
        "--out", type=Path, default=None,
        help="output verdicts.jsonl path (default: derived from --input)",
    )
    ap.add_argument(
        "--models", default=None,
        help="comma-separated name=endpoint:model_id; defaults to hermes + qwen",
    )
    ap.add_argument("--timeout", type=int, default=120)
    ap.add_argument(
        "--self-judge", action="store_true",
        help="every model judges its own answers (sanity check; default cross-judge)",
    )
    ap.add_argument(
        "--max-tokens", type=int, default=200,
        help="max tokens per judge reply (default 200 — judges should be terse; "
             "reasoning is disabled via chat_template_kwargs so this caps the "
             "actual verdict, not internal thinking)",
    )
    ap.add_argument(
        "--limit", type=int, default=0,
        help="cap rows judged (0 = all). Useful for smoke-testing the loop.",
    )
    args = ap.parse_args()

    if not args.input.exists():
        print(f"input not found: {args.input}", file=sys.stderr)
        return 1

    models = _parse_models(args.models) if args.models else list(DEFAULT_MODELS)
    if len(models) < 2 and not args.self_judge:
        print("cross-judge needs at least 2 models; use --self-judge for 1-model runs",
              file=sys.stderr)
        return 1

    rows_in = []
    with args.input.open() as f:
        for line in f:
            if not line.strip():
                continue
            rows_in.append(json.loads(line))
    if args.limit:
        rows_in = rows_in[: args.limit]
    print(f"# judging {len(rows_in)} rows ({'self-judge' if args.self_judge else 'cross-judge'})",
          file=sys.stderr)

    out_path = args.out
    if out_path is None:
        # Derive: <stem-minus-.judge_input>.verdicts.jsonl
        stem = args.input.stem
        if stem.endswith(".judge_input"):
            stem = stem[: -len(".judge_input")]
        out_path = args.input.parent / f"{stem}.verdicts.jsonl"

    verdict_rows: list[dict] = []
    with out_path.open("w") as f:
        for i, r in enumerate(rows_in, 1):
            answerer = r["model_name"]
            judge = _pick_judge(answerer, models, args.self_judge)
            if judge is None:
                print(f"  [{i}/{len(rows_in)}] {answerer} on {r['question'][:60]}: NO JUDGE AVAILABLE",
                      file=sys.stderr)
                continue
            judge_name, endpoint, model_id = judge
            print(f"  [{i}/{len(rows_in)}] {answerer} judged by {judge_name}: {r['question'][:60]}",
                  file=sys.stderr)
            # Regenerate the prompt fresh — lets us iterate on the prompt
            # template without re-running the underlying bench. The
            # judge_input.jsonl's stored `judge_prompt` is kept as an audit
            # record but not used.
            prompt = _judge_prompt(
                r["question"], answerer, r["audit_mode"], r["answer_text"],
            )
            try:
                reply, elapsed = _post_completion(
                    endpoint, model_id, prompt,
                    timeout_s=args.timeout, max_tokens=args.max_tokens,
                )
            except (urllib.error.URLError, urllib.error.HTTPError,
                    TimeoutError, OSError) as e:
                print(f"    ERR: {e}", file=sys.stderr)
                reply, elapsed = "", float(args.timeout)
            verdict, rationale = _parse_verdict(reply)
            print(f"    → {verdict} ({elapsed}s)", file=sys.stderr)
            row = {
                "question": r["question"],
                "answerer": answerer,
                "audit_mode": r["audit_mode"],
                "answer_text": r["answer_text"],
                "judge": judge_name,
                "verdict": verdict,
                "rationale": rationale,
                "raw_reply": reply,
                "_elapsed_s": elapsed,
            }
            verdict_rows.append(row)
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
            f.flush()

    md_path = out_path.with_suffix(".md")
    md_path.write_text(_format_summary(verdict_rows))
    print(f"\nVerdicts: {out_path}", file=sys.stderr)
    print(f"Summary: {md_path}", file=sys.stderr)
    print()
    print(md_path.read_text())
    return 0


if __name__ == "__main__":
    sys.exit(main())