scripts/build_compare.py

#!/usr/bin/env python3
"""Build monaco_dev `compare/` shards — side-by-side per-qid across 9 configs.

For each of the 100 monaco_dev test qids, joins the raw model ``answer`` +
the AML LLM-judge output (with ``--deterministic-extract``, the canonical
MoNaCo scoring protocol) from **9 configurations**:

  1. closed-book        — cell 1 · single-shot, no context
  2. with-docs          — cell 2 · single-shot, gold docs
  3. with-structures    — cell 3 · single-shot, eval structures
  4. structure per q    — cell 4 · agentic Baseline A (per-qid scaffolds)
  5. structure per ds   — cell 5 · agentic Baseline B (flat structure corpus)
  6. DCI                — cell 6 · agentic Baseline C v2 rerun (rawtext corpus)
  7. e2e v0             — e2e_pipeline run v0 (on-the-fly per-shape extraction)
  8. e2e v2             — e2e_pipeline run v2
  9. e2e v3             — e2e_pipeline run v3

Scoring column is the LLM judge's ``judge_score`` (0.0–1.0);
``correct = 1 iff judge_score >= 0.99``.

Output layout:
    compare/index.json
    compare/records/<qid>.json

Usage (defaults point at the canonical paths on this workstation):
    python scripts/build_compare.py
    python scripts/build_compare.py --scaffolds-root /path/to/information-scaffolds
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Tuple

DATASET_STAMPS = {"monaco", "monaco_dev"}
HERE = Path(__file__).resolve().parent
REPO = HERE.parent

DEFAULT_SCAFFOLDS = Path("/home/azureuser/projects/information-scaffolds")
DEFAULT_E2E_ROOT = Path("/home/azureuser/run_logs/e2e-monaco_dev-comparison")
DEFAULT_GOLD = REPO / "unified" / "records"
DEFAULT_OUT = REPO / "compare"


@dataclass(frozen=True)
class ConfigSpec:
    label: str
    pred_rel: str        # relative to --scaffolds-root or --e2e-root
    judged_rel: str      # relative to --scaffolds-root or --e2e-root
    source_root: str     # "scaffolds" | "e2e"
    shape: str           # "single" | "agentic"


CONFIGS: List[ConfigSpec] = [
    ConfigSpec("closed-book",
               "outputs/monaco_dev_smoke/sincere_pepper_d54pl76cjj/named-outputs/response/response",
               "outputs/monaco_dev/judges/c1_closedbook/named-outputs/judged/judged",
               "scaffolds", "single"),
    ConfigSpec("with-docs",
               "outputs/monaco_dev/baselines/cell2_withdocs/named-outputs/response/response",
               "outputs/monaco_dev/judges/c2_withdocs/named-outputs/judged/judged",
               "scaffolds", "single"),
    ConfigSpec("with-structures",
               "outputs/monaco_dev/baselines/cell3_withstructures/named-outputs/response/response",
               "outputs/monaco_dev/judges/c3_withstructures/named-outputs/judged/judged",
               "scaffolds", "single"),
    ConfigSpec("structure per q",
               "outputs/monaco_dev/baselines/cell4_agentic_a/named-outputs/response/response",
               "outputs/monaco_dev/judges/c4_agentic_a/named-outputs/judged/judged",
               "scaffolds", "agentic"),
    ConfigSpec("structure per ds",
               "outputs/monaco_dev/baselines/cell5_agentic_b/named-outputs/response/response",
               "outputs/monaco_dev/judges/c5_agentic_b/named-outputs/judged/judged",
               "scaffolds", "agentic"),
    ConfigSpec("DCI",
               "outputs/cell6_v2_reruns/monaco_dev_cell6_v2/named-outputs/response/response",
               "outputs/monaco_dev/judges/c6_agentic_c_v2/named-outputs/judged/judged",
               "scaffolds", "agentic"),
    ConfigSpec("e2e v0",
               "outputs/v0/named-outputs/predictions/predictions",
               "judge_outputs/v0/named-outputs/judged/judged",
               "e2e", "agentic"),
    ConfigSpec("e2e v2",
               "outputs/v2/named-outputs/predictions/predictions",
               "judge_outputs/v2/named-outputs/judged/judged",
               "e2e", "agentic"),
    ConfigSpec("e2e v3",
               "outputs/v3/named-outputs/predictions/predictions",
               "judge_outputs/v3/named-outputs/judged/judged",
               "e2e", "agentic"),
]


# ─── I/O ──────────────────────────────────────────────────────────────────────

def load_jsonl(path: Path) -> Dict[str, Dict[str, Any]]:
    """Load monaco JSONL keyed by qid; accepts stamps 'monaco' / 'monaco_dev'."""
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            ds = d.get("dataset")
            if ds is not None and ds not in DATASET_STAMPS:
                continue
            out[str(d["qid"])] = d
    return out


def load_gold(path: Path) -> Dict[str, Dict[str, Any]]:
    """Load gold from per-qid shard dir or single JSONL."""
    def _norm(d: Dict[str, Any]) -> Dict[str, Any]:
        d = dict(d)
        if "answer_list" not in d and "answers" in d:
            d["answer_list"] = list(d.get("answers") or [])
        if "question_text" not in d and "question" in d:
            d.setdefault("question_text", d["question"])
        return d

    out: Dict[str, Dict[str, Any]] = {}
    if path.is_dir():
        for shard in sorted(path.glob("*.json")):
            with shard.open() as f:
                d = json.load(f)
            qid = str(d.get("qid") or shard.stem)
            out[qid] = _norm(d)
        return out
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            out[str(d["qid"])] = _norm(d)
    return out


# ─── Per-config projection ────────────────────────────────────────────────────

def project_config(pred: Dict[str, Any], judged: Dict[str, Any] | None, shape: str) -> Dict[str, Any]:
    """Build compare-cell payload: light-weight; no per-event dump."""
    j = (judged or {}).get("parsed") or {}
    det = (judged or {}).get("deterministic_exact_answer") or {}
    judge_score = float(j.get("judge_score") or 0.0)
    metrics = {
        "judge_score": round(judge_score, 4),
        "correct": 1 if judge_score >= 0.99 else 0,
        "precision": round(float(j.get("precision") or 0.0), 4),
        "recall": round(float(j.get("recall") or 0.0), 4),
        "extracted_final_answer": j.get("extracted_final_answer"),
    }
    base: Dict[str, Any] = {
        "answer": pred.get("answer") or "",
        "pred_items": list(det.get("items") or []),
        "metrics": metrics,
        "model": pred.get("model"),
        "mode": pred.get("mode"),
        "system_prompt_file": pred.get("system_prompt_file"),
        "max_completion_tokens": pred.get("max_completion_tokens"),
        "latency_ms": pred.get("latency_ms"),
        "judged_missing": judged is None,
    }
    if shape == "single":
        base["finish_reason"] = pred.get("finish_reason")
        base["usage"] = pred.get("usage") or pred.get("tokens")
    else:
        base["stop_reason"] = pred.get("stop_reason")
        base["n_turns"] = pred.get("turns") or pred.get("n_turns")
        base["max_turns"] = pred.get("max_turns")
        base["tool_call_counts"] = pred.get("tool_call_counts") or {}
        base["tokens"] = pred.get("tokens") or {}
        base["finish_reasons"] = pred.get("finish_reasons")
        base["timeout_retries"] = pred.get("timeout_retries") or 0
    return base


# ─── Main ─────────────────────────────────────────────────────────────────────

def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--scaffolds-root", type=Path, default=DEFAULT_SCAFFOLDS,
                    help="Root of `information-scaffolds` repo (for cells 1-6).")
    ap.add_argument("--e2e-root", type=Path, default=DEFAULT_E2E_ROOT,
                    help="Root of the e2e-monaco_dev-comparison run_logs dir "
                         "(has outputs/v0|v2|v3/ + judge_outputs/v0|v2|v3/).")
    ap.add_argument("--gold", type=Path, default=DEFAULT_GOLD)
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    args = ap.parse_args()

    gold = load_gold(args.gold)
    print(f"gold: {len(gold)} qids from {args.gold}", file=sys.stderr)

    # Load all 9 configs.
    cells: List[Tuple[str, ConfigSpec, Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]] = []
    for cfg in CONFIGS:
        root = args.scaffolds_root if cfg.source_root == "scaffolds" else args.e2e_root
        pred_path = root / cfg.pred_rel
        judged_path = root / cfg.judged_rel
        if not pred_path.exists():
            print(f"warning: pred missing, skipping {cfg.label!r}: {pred_path}", file=sys.stderr)
            continue
        preds = load_jsonl(pred_path)
        judged = load_jsonl(judged_path) if judged_path.exists() else {}
        cells.append((cfg.label, cfg, preds, judged))
        n_j_missing = sum(1 for q in preds if q not in judged)
        print(f"  {cfg.label:18s} preds={len(preds):3d}  judged={len(judged):3d}  "
              f"(missing_judge={n_j_missing})", file=sys.stderr)

    # Intersection of predictions × gold across every config (each qid must
    # be in every cell to render a complete side-by-side row).
    qids = set(gold)
    for _label, _cfg, preds, _judged in cells:
        qids &= set(preds)
    qids_sorted = sorted(qids, key=lambda q: int(q) if q.isdigit() else q)
    print(f"common qids: {len(qids_sorted)}", file=sys.stderr)
    if not qids_sorted:
        print("ERROR: no overlap across cells × gold", file=sys.stderr)
        return 1

    # Write per-qid shards + index.
    out_dir = args.out
    out_dir.mkdir(parents=True, exist_ok=True)
    rec_dir = out_dir / "records"
    if rec_dir.exists():
        shutil.rmtree(rec_dir)
    rec_dir.mkdir(parents=True)

    index_rows: List[Dict[str, Any]] = []
    sum_metrics: Dict[str, Dict[str, float]] = {
        label: {"judge_score": 0.0, "correct": 0.0}
        for label, _cfg, _p, _j in cells
    }

    for qid in qids_sorted:
        gold_row = gold[qid]
        gold_items = list(gold_row.get("answer_list") or [])
        configs_payload: Dict[str, Dict[str, Any]] = {}
        per_cell_brief: List[Dict[str, Any]] = []
        first_question = gold_row.get("question_text")
        for label, cfg, preds, judged in cells:
            pred = preds[qid]
            j = judged.get(qid)
            payload = project_config(pred, j, cfg.shape)
            configs_payload[label] = payload
            per_cell_brief.append({
                "label": label,
                "judge_score": payload["metrics"]["judge_score"],
                "correct": payload["metrics"]["correct"],
                "judged_missing": payload["judged_missing"],
            })
            sum_metrics[label]["judge_score"] += payload["metrics"]["judge_score"]
            sum_metrics[label]["correct"] += payload["metrics"]["correct"]
            if not first_question and pred.get("question"):
                first_question = pred.get("question")

        record = {
            "qid": qid,
            "question": first_question or "",
            "gold_answers": gold_items,
            "configs_order": [label for label, *_ in cells],
            "configs": configs_payload,
        }
        (rec_dir / f"{qid}.json").write_text(json.dumps(record, ensure_ascii=False))
        index_rows.append({
            "qid": qid,
            "question": record["question"],
            "gold_answers_length": len(gold_items),
            "cells": per_cell_brief,
        })

    n = len(qids_sorted)
    summary = {
        "n": n,
        "configs": [label for label, *_ in cells],
        "mean_metrics": {
            label: {
                "mean_judge_score": round(s["judge_score"] / n, 4),
                "mean_correct": round(s["correct"] / n, 4),
            }
            for label, s in sum_metrics.items()
        },
    }
    (out_dir / "index.json").write_text(json.dumps({"meta": summary, "rows": index_rows}, ensure_ascii=False))

    print(f"\n✓ Wrote {out_dir}/index.json + {n} record shards", file=sys.stderr)
    for label in summary["mean_metrics"]:
        mm = summary["mean_metrics"][label]
        print(f"   {label:18s} mean judge_score = {mm['mean_judge_score']*100:6.2f}   "
              f"mean correct = {mm['mean_correct']*100:6.2f}", file=sys.stderr)
    return 0


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