File size: 4,419 Bytes
3fc79a4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
#!/usr/bin/env python3
"""Build the static leaderboard data bundle from MIMIC-CDM result repositories."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
from collections import defaultdict
from datetime import datetime, timezone
from decimal import Decimal, ROUND_HALF_EVEN
from pathlib import Path
from typing import Any


TASKS = ("Appendicitis", "Cholecystitis", "Diverticulitis", "Pancreatitis")
VARIANT_SUFFIX = re.compile(r"\s+(?:4|8|16)bit$", re.IGNORECASE)


def read_git_metadata(path: Path) -> dict[str, str]:
    try:
        output = subprocess.run(
            ["git", "-C", str(path), "log", "-1", "--format=%H%n%cI"],
            check=True,
            capture_output=True,
            text=True,
        ).stdout.splitlines()
        return {"revision": output[0], "updated_at": output[1]}
    except (OSError, subprocess.CalledProcessError, IndexError):
        return {
            "revision": "local",
            "updated_at": datetime.now(timezone.utc).isoformat(),
        }


def model_hub_id(model_name: str) -> str:
    return VARIANT_SUFFIX.sub("", model_name).strip()


def load_result_files(results_path: Path) -> list[dict[str, Any]]:
    merged: dict[tuple[str, int], dict[str, Any]] = defaultdict(
        lambda: {"scores": {}}
    )

    for result_file in sorted(results_path.rglob("*.json")):
        try:
            payload = json.loads(result_file.read_text())
            config = payload["config"]
            model_name = str(config["model_name"]).strip()
            bits = int(config.get("model_quantization_bits") or 0)
            key = (model_name, bits)
            run = merged[key]
            run.update(
                {
                    "model": model_name,
                    "hub_model": model_hub_id(model_name),
                    "params_b": config.get("params"),
                    "context_length": config.get("max_sequence_length"),
                    "quantization_bits": bits or None,
                }
            )
            for task in TASKS:
                score = payload.get("results", {}).get(task, {}).get("acc")
                if score is not None:
                    run["scores"][task] = float(
                        (Decimal(str(score)) * 100).quantize(
                            Decimal("0.01"), rounding=ROUND_HALF_EVEN
                        )
                    )
        except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error:
            raise ValueError(f"Invalid result file {result_file}: {error}") from error

    rows = []
    for run in merged.values():
        if not all(task in run["scores"] for task in TASKS):
            continue
        run["average"] = float(
            (
                sum(Decimal(str(run["scores"][task])) for task in TASKS)
                / len(TASKS)
            ).quantize(Decimal("0.01"), rounding=ROUND_HALF_EVEN)
        )
        rows.append(run)

    rows.sort(key=lambda row: (-row["average"], row["model"].lower()))
    for rank, row in enumerate(rows, start=1):
        row["rank"] = rank
    return rows


def build_bundle(cdm_path: Path, cdm_fi_path: Path) -> dict[str, Any]:
    sources = {
        "cdm": {
            "repo": "MIMIC-CDM/results-CDM",
            **read_git_metadata(cdm_path),
        },
        "cdm_fi": {
            "repo": "MIMIC-CDM/results-CDM-FI",
            **read_git_metadata(cdm_fi_path),
        },
    }
    return {
        "schema_version": 1,
        "tasks": list(TASKS),
        "sources": sources,
        "leaderboards": {
            "cdm": load_result_files(cdm_path),
            "cdm_fi": load_result_files(cdm_fi_path),
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cdm-dir", required=True, type=Path)
    parser.add_argument("--cdm-fi-dir", required=True, type=Path)
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("data/leaderboards.json"),
    )
    args = parser.parse_args()

    bundle = build_bundle(args.cdm_dir, args.cdm_fi_dir)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(bundle, indent=2) + "\n")
    print(
        f"Wrote {len(bundle['leaderboards']['cdm'])} CDM and "
        f"{len(bundle['leaderboards']['cdm_fi'])} CDM-FI entries to {args.output}"
    )


if __name__ == "__main__":
    main()