scripts/compare_models.py

#!/usr/bin/env python3
"""Compare two aligned prediction files with paired group bootstrap intervals."""

from __future__ import annotations

import argparse
import hashlib
import json
import random
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

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

from turn_detection.training.metrics import (  # noqa: E402
    average_precision,
    binary_classification_metrics,
    roc_auc,
    threshold_at_max_fpr,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--baseline", required=True, help="baseline prediction JSONL")
    parser.add_argument("--candidate", required=True, help="candidate prediction JSONL")
    parser.add_argument("--output", required=True)
    parser.add_argument("--baseline-name", default="baseline")
    parser.add_argument("--candidate-name", default="candidate")
    parser.add_argument("--fpr-budget", type=float, default=0.02)
    parser.add_argument("--bootstrap-samples", type=int, default=2_000)
    parser.add_argument("--seed", type=int, default=17)
    return parser.parse_args()


def _load(path: Path) -> dict[str, dict[str, Any]]:
    records: dict[str, dict[str, Any]] = {}
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: invalid JSON") from exc
            if not isinstance(row, dict):
                raise ValueError(f"{path}:{line_number}: expected an object")
            record_id = str(row.get("record_id", ""))
            if not record_id:
                raise ValueError(f"{path}:{line_number}: record_id is missing")
            if record_id in records:
                raise ValueError(f"{path}:{line_number}: duplicate record_id {record_id!r}")
            label = row.get("label")
            probability = row.get("probability")
            if label not in (0, 1, False, True):
                raise ValueError(f"{path}:{line_number}: label must be binary")
            if not isinstance(probability, int | float) or not 0.0 <= probability <= 1.0:
                raise ValueError(f"{path}:{line_number}: probability must be in [0, 1]")
            records[record_id] = row
    if not records:
        raise ValueError(f"{path}: no predictions")
    return records


def _percentile(values: list[float], quantile: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * quantile
    lower = int(position)
    upper = min(lower + 1, len(ordered) - 1)
    fraction = position - lower
    return ordered[lower] * (1.0 - fraction) + ordered[upper] * fraction


def _interval(values: list[float], point: float) -> dict[str, float]:
    return {
        "estimate": point,
        "lower": _percentile(values, 0.025),
        "upper": _percentile(values, 0.975),
        "bootstrap_fraction_delta_above_zero": sum(value > 0.0 for value in values) / len(values),
    }


def _portable_path(path: Path) -> str:
    try:
        return path.resolve().relative_to(REPOSITORY_ROOT).as_posix()
    except ValueError:
        return path.name


def _file_evidence(path: Path) -> dict[str, str | int]:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return {
        "path": _portable_path(path),
        "bytes": path.stat().st_size,
        "sha256": digest.hexdigest(),
    }


def main() -> int:
    args = parse_args()
    if not 0.0 <= args.fpr_budget <= 1.0:
        raise SystemExit("--fpr-budget must be in [0, 1]")
    if args.bootstrap_samples < 1:
        raise SystemExit("--bootstrap-samples must be positive")
    baseline_path = Path(args.baseline)
    candidate_path = Path(args.candidate)
    if not baseline_path.is_absolute():
        baseline_path = REPOSITORY_ROOT / baseline_path
    if not candidate_path.is_absolute():
        candidate_path = REPOSITORY_ROOT / candidate_path
    baseline = _load(baseline_path)
    candidate = _load(candidate_path)
    if baseline.keys() != candidate.keys():
        only_baseline = len(baseline.keys() - candidate.keys())
        only_candidate = len(candidate.keys() - baseline.keys())
        raise SystemExit(
            "prediction sets are not aligned: "
            f"{only_baseline} only in baseline, {only_candidate} only in candidate"
        )

    record_ids = sorted(baseline)
    labels: list[int] = []
    baseline_scores: list[float] = []
    candidate_scores: list[float] = []
    groups: list[str] = []
    for record_id in record_ids:
        baseline_row = baseline[record_id]
        candidate_row = candidate[record_id]
        if int(baseline_row["label"]) != int(candidate_row["label"]):
            raise SystemExit(f"label mismatch for record {record_id}")
        baseline_group = str(baseline_row.get("group_id") or record_id)
        candidate_group = str(candidate_row.get("group_id") or record_id)
        if baseline_group != candidate_group:
            raise SystemExit(f"group mismatch for record {record_id}")
        labels.append(int(baseline_row["label"]))
        baseline_scores.append(float(baseline_row["probability"]))
        candidate_scores.append(float(candidate_row["probability"]))
        groups.append(baseline_group)

    baseline_point = threshold_at_max_fpr(labels, baseline_scores, args.fpr_budget)
    candidate_point = threshold_at_max_fpr(labels, candidate_scores, args.fpr_budget)
    baseline_threshold = float(baseline_point["threshold"])
    candidate_threshold = float(candidate_point["threshold"])
    by_group: dict[str, list[int]] = defaultdict(list)
    for index, group in enumerate(groups):
        by_group[group].append(index)
    group_names = sorted(by_group)
    rng = random.Random(args.seed)
    bootstrap_deltas: dict[str, list[float]] = defaultdict(list)
    for _ in range(args.bootstrap_samples):
        sampled_groups = [rng.choice(group_names) for _ in group_names]
        indices = [index for group in sampled_groups for index in by_group[group]]
        sample_labels = [labels[index] for index in indices]
        baseline_sample = [baseline_scores[index] for index in indices]
        candidate_sample = [candidate_scores[index] for index in indices]
        baseline_auc = roc_auc(sample_labels, baseline_sample)
        candidate_auc = roc_auc(sample_labels, candidate_sample)
        baseline_ap = average_precision(sample_labels, baseline_sample)
        candidate_ap = average_precision(sample_labels, candidate_sample)
        if baseline_auc is not None and candidate_auc is not None:
            bootstrap_deltas["roc_auc"].append(candidate_auc - baseline_auc)
        if baseline_ap is not None and candidate_ap is not None:
            bootstrap_deltas["average_precision"].append(candidate_ap - baseline_ap)
        baseline_fixed = binary_classification_metrics(
            sample_labels, baseline_sample, baseline_threshold
        )
        candidate_fixed = binary_classification_metrics(
            sample_labels, candidate_sample, candidate_threshold
        )
        for metric in ("recall", "false_positive_rate"):
            baseline_value = baseline_fixed[metric]
            candidate_value = candidate_fixed[metric]
            if baseline_value is not None and candidate_value is not None:
                bootstrap_deltas[metric].append(candidate_value - baseline_value)

    baseline_metrics = binary_classification_metrics(labels, baseline_scores, baseline_threshold)
    candidate_metrics = binary_classification_metrics(labels, candidate_scores, candidate_threshold)
    point_deltas = {
        "roc_auc": float(candidate_metrics["roc_auc"] - baseline_metrics["roc_auc"]),
        "average_precision": float(
            candidate_metrics["average_precision"] - baseline_metrics["average_precision"]
        ),
        "recall": float(candidate_metrics["recall"] - baseline_metrics["recall"]),
        "false_positive_rate": float(
            candidate_metrics["false_positive_rate"] - baseline_metrics["false_positive_rate"]
        ),
    }
    comparison = {
        "scope": "paired development comparison; not independent-test evidence",
        "count": len(labels),
        "group_count": len(group_names),
        "positive_count": sum(labels),
        "negative_count": len(labels) - sum(labels),
        "fpr_budget": args.fpr_budget,
        "threshold_note": (
            "Each threshold was selected on this same development set. Bootstrap samples use "
            "those fixed thresholds; intervals do not remove model-selection or calibration bias."
        ),
        "baseline": {
            "name": args.baseline_name,
            "predictions": _file_evidence(baseline_path),
            "metrics": baseline_metrics,
        },
        "candidate": {
            "name": args.candidate_name,
            "predictions": _file_evidence(candidate_path),
            "metrics": candidate_metrics,
        },
        "candidate_minus_baseline_95ci": {
            metric: _interval(bootstrap_deltas[metric], point)
            for metric, point in point_deltas.items()
        },
        "bootstrap": {
            "unit": "audit leakage group",
            "samples": args.bootstrap_samples,
            "seed": args.seed,
        },
    }
    output = Path(args.output)
    if not output.is_absolute():
        output = REPOSITORY_ROOT / output
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(
        json.dumps(comparison, indent=2, sort_keys=True, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(comparison, indent=2, sort_keys=True, allow_nan=False))
    return 0


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