Sync InvoiceGuard code for GRPO training job

9a88af0 verified 7 months ago

8.88 kB

"""
Round 2 evaluation harness for InvoiceGuard.
Runs a chosen task slice (canonical / hard / all) against any
OpenAI-compatible model and writes a structured JSON report containing,
for every task: grader score, six grader sub-component scores, decision,
exception type, steps used, shortcut-penalty flag, and the per-step
reward_components log emitted by the environment.
Designed so baseline_*.json (Stage F) and trained_*.json (Stage G) share
the SAME schema and can be diffed with `--compare A.json B.json`.
Usage examples (PowerShell):
    # Baseline run on hard slice with default model from env vars
    uv run python eval_round2.py --slice hard --model-tag baseline
    # Run both slices, write to outputs/round2/
    uv run python eval_round2.py --slice all --model-tag baseline
    # Compare baseline vs trained
    uv run python eval_round2.py --compare outputs/round2/hard__baseline.json outputs/round2/hard__trained.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import List

from dotenv import load_dotenv
from openai import OpenAI

# Reuse the existing inference helpers so prompts/parsing stay identical.
from inference import (  # type: ignore
    API_BASE_URL,
    API_KEY,
    MODEL_NAME,
    run_episode_local,
)
from models import TaskID  # type: ignore
from server.invoice_guard_environment import InvoiceGuardEnvironment  # type: ignore
from tasks import TASK_LIST, HARD_TASK_LIST  # type: ignore


load_dotenv()

OUT_DIR_DEFAULT = Path(__file__).parent / "outputs" / "round2"

# Grader sub-component keys we always extract for the report.
COMPONENT_KEYS = [
    "decision_score",
    "exception_type_score",
    "evidence_score",
    "investigation_score",
    "explanation_score",
    "efficiency_score",
]


def _slice_tasks(slice_name: str) -> List[TaskID]:
    if slice_name == "canonical":
        return list(TASK_LIST)
    if slice_name == "hard":
        return list(HARD_TASK_LIST)
    if slice_name == "all":
        return list(TASK_LIST) + list(HARD_TASK_LIST)
    raise SystemExit(f"Unknown --slice: {slice_name!r}")


def _run_slice(
    slice_name: str,
    model_tag: str,
    out_dir: Path,
) -> Path:
    """Run one slice end-to-end and write the JSON report. Returns path."""
    out_dir.mkdir(parents=True, exist_ok=True)
    task_ids = _slice_tasks(slice_name)

    print("=" * 70, flush=True)
    print(f"InvoiceGuard Round 2 eval | slice={slice_name} | tasks={len(task_ids)}", flush=True)
    print(f"Model: {MODEL_NAME} | Base URL: {API_BASE_URL}", flush=True)
    print(f"Tag:   {model_tag} | Out dir: {out_dir}", flush=True)
    print("=" * 70, flush=True)

    llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    env = InvoiceGuardEnvironment()

    per_task: list[dict] = []
    started_at = datetime.now(timezone.utc).isoformat()
    t0 = time.time()

    for task_id in task_ids:
        start = time.time()
        # `run_episode_local` already handles the conversation loop and stdout
        # in the hackathon-mandated [START]/[STEP]/[END] format.
        result = run_episode_local(env, llm, task_id)
        elapsed = time.time() - start

        grader_breakdown = result.get("grader_breakdown") or {}
        components = {k: float(grader_breakdown.get(k, 0.0)) for k in COMPONENT_KEYS}

        # Pull richer metadata from the env directly (last terminal obs went
        # back through the inference loop but we still hold env.state).
        s = env.state
        reward_components = list(getattr(s, "reward_components", []))

        per_task.append(
            {
                "task_id": result["task_id"],
                "decision": result.get("decision"),
                "exception_type": result.get("exception_type"),
                "steps": result["steps"],
                "grader_score": result["grader_score"],
                "components": components,
                "shortcut_penalty_applied": any(
                    rc.get("penalties", {}).get("shortcut") for rc in reward_components
                ),
                "documents_revealed": list(s.documents_revealed),
                "actions_taken": list(s.actions_taken),
                "reward_components": reward_components,
                "wall_clock_s": round(elapsed, 2),
            }
        )
        print(
            f"  >> {result['task_id']:38s} "
            f"score={result['grader_score']:.4f}  "
            f"steps={result['steps']:>2}  "
            f"decision={result.get('decision')}  "
            f"({elapsed:.1f}s)",
            flush=True,
        )

    total_elapsed = time.time() - t0
    avg_score = (
        sum(t["grader_score"] for t in per_task) / len(per_task) if per_task else 0.0
    )
    component_avgs = {
        k: round(sum(t["components"][k] for t in per_task) / max(len(per_task), 1), 4)
        for k in COMPONENT_KEYS
    }
    decision_correct = sum(
        1 for t in per_task if t["components"]["decision_score"] >= 0.99
    )

    report = {
        "schema_version": 1,
        "slice": slice_name,
        "model_tag": model_tag,
        "model_name": MODEL_NAME,
        "api_base_url": API_BASE_URL,
        "started_at": started_at,
        "wall_clock_s": round(total_elapsed, 2),
        "n_tasks": len(per_task),
        "summary": {
            "avg_score": round(avg_score, 4),
            "decision_correct": decision_correct,
            "decision_correct_rate": round(decision_correct / max(len(per_task), 1), 4),
            "component_avgs": component_avgs,
            "shortcut_episodes": sum(
                1 for t in per_task if t["shortcut_penalty_applied"]
            ),
        },
        "tasks": per_task,
    }

    out_path = out_dir / f"{slice_name}__{model_tag}.json"
    out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print("-" * 70, flush=True)
    print(f"Avg score:           {avg_score:.4f}", flush=True)
    print(f"Decision correct:    {decision_correct}/{len(per_task)}", flush=True)
    print(f"Component averages:  {component_avgs}", flush=True)
    print(f"Wrote report:        {out_path}", flush=True)
    print("=" * 70, flush=True)

    return out_path


def _compare(a_path: Path, b_path: Path) -> None:
    a = json.loads(a_path.read_text(encoding="utf-8"))
    b = json.loads(b_path.read_text(encoding="utf-8"))

    print("=" * 78, flush=True)
    print(f"COMPARE  A: {a_path.name}  ({a['model_tag']})", flush=True)
    print(f"         B: {b_path.name}  ({b['model_tag']})", flush=True)
    print("=" * 78, flush=True)

    a_by = {t["task_id"]: t for t in a["tasks"]}
    b_by = {t["task_id"]: t for t in b["tasks"]}
    keys = list(a_by.keys())

    print(f"{'task':40s} {'A':>7s} {'B':>7s} {'delta':>8s}", flush=True)
    print("-" * 78, flush=True)
    for k in keys:
        if k not in b_by:
            continue
        sa = a_by[k]["grader_score"]
        sb = b_by[k]["grader_score"]
        d = sb - sa
        marker = "  +" if d > 0.01 else ("  -" if d < -0.01 else "   ")
        print(f"{k:40s} {sa:7.4f} {sb:7.4f} {d:+8.4f}{marker}", flush=True)

    print("-" * 78, flush=True)
    print(
        f"{'AVERAGE':40s} {a['summary']['avg_score']:7.4f} "
        f"{b['summary']['avg_score']:7.4f} "
        f"{b['summary']['avg_score'] - a['summary']['avg_score']:+8.4f}",
        flush=True,
    )
    for ck in COMPONENT_KEYS:
        ac = a["summary"]["component_avgs"][ck]
        bc = b["summary"]["component_avgs"][ck]
        print(f"  {ck:38s} {ac:7.4f} {bc:7.4f} {bc - ac:+8.4f}", flush=True)
    print("=" * 78, flush=True)


def main() -> None:
    p = argparse.ArgumentParser(description="InvoiceGuard Round 2 eval harness.")
    p.add_argument("--slice", choices=["canonical", "hard", "all"], default="hard")
    p.add_argument("--model-tag", default="baseline",
                   help="Tag used in the output filename (e.g. baseline, trained, qwen3b-grpo).")
    p.add_argument("--out-dir", type=Path, default=OUT_DIR_DEFAULT)
    p.add_argument("--compare", nargs=2, metavar=("A", "B"), type=Path, default=None,
                   help="Compare two report JSONs and print a delta table.")
    args = p.parse_args()

    if args.compare:
        _compare(args.compare[0], args.compare[1])
        return

    if not API_KEY:
        print(
            "WARNING: no API key found in env (HF_TOKEN / API_KEY / OPENAI_API_KEY). "
            "LLM calls will fail; this run will only verify the harness wiring.",
            file=sys.stderr,
            flush=True,
        )

    if args.slice == "all":
        _run_slice("canonical", args.model_tag, args.out_dir)
        _run_slice("hard", args.model_tag, args.out_dir)
    else:
        _run_slice(args.slice, args.model_tag, args.out_dir)


if __name__ == "__main__":
    main()

	"""
	Round 2 evaluation harness for InvoiceGuard.

	Runs a chosen task slice (canonical / hard / all) against any
	OpenAI-compatible model and writes a structured JSON report containing,
	for every task: grader score, six grader sub-component scores, decision,
	exception type, steps used, shortcut-penalty flag, and the per-step
	reward_components log emitted by the environment.

	Designed so baseline_.json (Stage F) and trained_.json (Stage G) share
	the SAME schema and can be diffed with `--compare A.json B.json`.

	Usage examples (PowerShell):
	# Baseline run on hard slice with default model from env vars
	uv run python eval_round2.py --slice hard --model-tag baseline

	# Run both slices, write to outputs/round2/
	uv run python eval_round2.py --slice all --model-tag baseline

	# Compare baseline vs trained
	uv run python eval_round2.py --compare outputs/round2/hard__baseline.json outputs/round2/hard__trained.json
	"""

	from __future__ import annotations

	import argparse
	import json
	import os
	import sys
	import time
	from datetime import datetime, timezone
	from pathlib import Path
	from typing import List

	from dotenv import load_dotenv
	from openai import OpenAI

	# Reuse the existing inference helpers so prompts/parsing stay identical.
	from inference import ( # type: ignore
	API_BASE_URL,
	API_KEY,
	MODEL_NAME,
	run_episode_local,
	)
	from models import TaskID # type: ignore
	from server.invoice_guard_environment import InvoiceGuardEnvironment # type: ignore
	from tasks import TASK_LIST, HARD_TASK_LIST # type: ignore


	load_dotenv()

	OUT_DIR_DEFAULT = Path(__file__).parent / "outputs" / "round2"

	# Grader sub-component keys we always extract for the report.
	COMPONENT_KEYS = [
	"decision_score",
	"exception_type_score",
	"evidence_score",
	"investigation_score",
	"explanation_score",
	"efficiency_score",
	]


	def _slice_tasks(slice_name: str) -> List[TaskID]:
	if slice_name == "canonical":
	return list(TASK_LIST)
	if slice_name == "hard":
	return list(HARD_TASK_LIST)
	if slice_name == "all":
	return list(TASK_LIST) + list(HARD_TASK_LIST)
	raise SystemExit(f"Unknown --slice: {slice_name!r}")


	def _run_slice(
	slice_name: str,
	model_tag: str,
	out_dir: Path,
	) -> Path:
	"""Run one slice end-to-end and write the JSON report. Returns path."""
	out_dir.mkdir(parents=True, exist_ok=True)
	task_ids = _slice_tasks(slice_name)

	print("=" * 70, flush=True)
	print(f"InvoiceGuard Round 2 eval \| slice={slice_name} \| tasks={len(task_ids)}", flush=True)
	print(f"Model: {MODEL_NAME} \| Base URL: {API_BASE_URL}", flush=True)
	print(f"Tag: {model_tag} \| Out dir: {out_dir}", flush=True)
	print("=" * 70, flush=True)

	llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
	env = InvoiceGuardEnvironment()

	per_task: list[dict] = []
	started_at = datetime.now(timezone.utc).isoformat()
	t0 = time.time()

	for task_id in task_ids:
	start = time.time()
	# `run_episode_local` already handles the conversation loop and stdout
	# in the hackathon-mandated [START]/[STEP]/[END] format.
	result = run_episode_local(env, llm, task_id)
	elapsed = time.time() - start

	grader_breakdown = result.get("grader_breakdown") or {}
	components = {k: float(grader_breakdown.get(k, 0.0)) for k in COMPONENT_KEYS}

	# Pull richer metadata from the env directly (last terminal obs went
	# back through the inference loop but we still hold env.state).
	s = env.state
	reward_components = list(getattr(s, "reward_components", []))

	per_task.append(
	{
	"task_id": result["task_id"],
	"decision": result.get("decision"),
	"exception_type": result.get("exception_type"),
	"steps": result["steps"],
	"grader_score": result["grader_score"],
	"components": components,
	"shortcut_penalty_applied": any(
	rc.get("penalties", {}).get("shortcut") for rc in reward_components
	),
	"documents_revealed": list(s.documents_revealed),
	"actions_taken": list(s.actions_taken),
	"reward_components": reward_components,
	"wall_clock_s": round(elapsed, 2),
	}
	)
	print(
	f" >> {result['task_id']:38s} "
	f"score={result['grader_score']:.4f} "
	f"steps={result['steps']:>2} "
	f"decision={result.get('decision')} "
	f"({elapsed:.1f}s)",
	flush=True,
	)

	total_elapsed = time.time() - t0
	avg_score = (
	sum(t["grader_score"] for t in per_task) / len(per_task) if per_task else 0.0
	)
	component_avgs = {
	k: round(sum(t["components"][k] for t in per_task) / max(len(per_task), 1), 4)
	for k in COMPONENT_KEYS
	}
	decision_correct = sum(
	1 for t in per_task if t["components"]["decision_score"] >= 0.99
	)

	report = {
	"schema_version": 1,
	"slice": slice_name,
	"model_tag": model_tag,
	"model_name": MODEL_NAME,
	"api_base_url": API_BASE_URL,
	"started_at": started_at,
	"wall_clock_s": round(total_elapsed, 2),
	"n_tasks": len(per_task),
	"summary": {
	"avg_score": round(avg_score, 4),
	"decision_correct": decision_correct,
	"decision_correct_rate": round(decision_correct / max(len(per_task), 1), 4),
	"component_avgs": component_avgs,
	"shortcut_episodes": sum(
	1 for t in per_task if t["shortcut_penalty_applied"]
	),
	},
	"tasks": per_task,
	}

	out_path = out_dir / f"{slice_name}__{model_tag}.json"
	out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

	print("-" * 70, flush=True)
	print(f"Avg score: {avg_score:.4f}", flush=True)
	print(f"Decision correct: {decision_correct}/{len(per_task)}", flush=True)
	print(f"Component averages: {component_avgs}", flush=True)
	print(f"Wrote report: {out_path}", flush=True)
	print("=" * 70, flush=True)

	return out_path


	def _compare(a_path: Path, b_path: Path) -> None:
	a = json.loads(a_path.read_text(encoding="utf-8"))
	b = json.loads(b_path.read_text(encoding="utf-8"))

	print("=" * 78, flush=True)
	print(f"COMPARE A: {a_path.name} ({a['model_tag']})", flush=True)
	print(f" B: {b_path.name} ({b['model_tag']})", flush=True)
	print("=" * 78, flush=True)

	a_by = {t["task_id"]: t for t in a["tasks"]}
	b_by = {t["task_id"]: t for t in b["tasks"]}
	keys = list(a_by.keys())

	print(f"{'task':40s} {'A':>7s} {'B':>7s} {'delta':>8s}", flush=True)
	print("-" * 78, flush=True)
	for k in keys:
	if k not in b_by:
	continue
	sa = a_by[k]["grader_score"]
	sb = b_by[k]["grader_score"]
	d = sb - sa
	marker = " +" if d > 0.01 else (" -" if d < -0.01 else " ")
	print(f"{k:40s} {sa:7.4f} {sb:7.4f} {d:+8.4f}{marker}", flush=True)

	print("-" * 78, flush=True)
	print(
	f"{'AVERAGE':40s} {a['summary']['avg_score']:7.4f} "
	f"{b['summary']['avg_score']:7.4f} "
	f"{b['summary']['avg_score'] - a['summary']['avg_score']:+8.4f}",
	flush=True,
	)
	for ck in COMPONENT_KEYS:
	ac = a["summary"]["component_avgs"][ck]
	bc = b["summary"]["component_avgs"][ck]
	print(f" {ck:38s} {ac:7.4f} {bc:7.4f} {bc - ac:+8.4f}", flush=True)
	print("=" * 78, flush=True)


	def main() -> None:
	p = argparse.ArgumentParser(description="InvoiceGuard Round 2 eval harness.")
	p.add_argument("--slice", choices=["canonical", "hard", "all"], default="hard")
	p.add_argument("--model-tag", default="baseline",
	help="Tag used in the output filename (e.g. baseline, trained, qwen3b-grpo).")
	p.add_argument("--out-dir", type=Path, default=OUT_DIR_DEFAULT)
	p.add_argument("--compare", nargs=2, metavar=("A", "B"), type=Path, default=None,
	help="Compare two report JSONs and print a delta table.")
	args = p.parse_args()

	if args.compare:
	_compare(args.compare[0], args.compare[1])
	return

	if not API_KEY:
	print(
	"WARNING: no API key found in env (HF_TOKEN / API_KEY / OPENAI_API_KEY). "
	"LLM calls will fail; this run will only verify the harness wiring.",
	file=sys.stderr,
	flush=True,
	)

	if args.slice == "all":
	_run_slice("canonical", args.model_tag, args.out_dir)
	_run_slice("hard", args.model_tag, args.out_dir)
	else:
	_run_slice(args.slice, args.model_tag, args.out_dir)


	if __name__ == "__main__":
	main()