File size: 2,527 Bytes
35d483e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | from __future__ import annotations
import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
class ModelComparisonTest(unittest.TestCase):
def test_aligned_candidate_reports_paired_improvement(self) -> None:
labels = [0, 0, 0, 0, 1, 1, 1, 1]
baseline_scores = [0.1, 0.4, 0.3, 0.2, 0.3, 0.6, 0.5, 0.4]
candidate_scores = [0.05, 0.10, 0.15, 0.20, 0.80, 0.85, 0.90, 0.95]
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
baseline = root / "baseline.jsonl"
candidate = root / "candidate.jsonl"
output = root / "comparison.json"
for path, scores in ((baseline, baseline_scores), (candidate, candidate_scores)):
path.write_text(
"".join(
json.dumps(
{
"record_id": f"row-{index}",
"group_id": f"group-{index}",
"label": label,
"probability": score,
}
)
+ "\n"
for index, (label, score) in enumerate(zip(labels, scores, strict=True))
),
encoding="utf-8",
)
completed = subprocess.run(
[
sys.executable,
str(ROOT / "scripts/compare_models.py"),
"--baseline",
str(baseline),
"--candidate",
str(candidate),
"--output",
str(output),
"--fpr-budget",
"0.25",
"--bootstrap-samples",
"50",
],
cwd=ROOT,
check=False,
capture_output=True,
text=True,
)
self.assertEqual(completed.returncode, 0, completed.stderr)
report = json.loads(output.read_text(encoding="utf-8"))
self.assertEqual(report["count"], 8)
self.assertEqual(report["group_count"], 8)
self.assertGreater(report["candidate_minus_baseline_95ci"]["roc_auc"]["estimate"], 0.0)
self.assertEqual(report["candidate"]["metrics"]["roc_auc"], 1.0)
if __name__ == "__main__":
unittest.main()
|