python-basic-pitch/shadow_benchmark.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from temporal_chord_metrics import ChordInterval, evaluate_temporal_chords


def load_timeline(path: Path) -> list[ChordInterval]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    raw_segments = payload.get("segments") if isinstance(payload, dict) else payload
    if not isinstance(raw_segments, list):
        raise ValueError(f"{path}: esperado array ou objeto com 'segments'")
    return [
        ChordInterval(
            start=float(item["start"]),
            end=float(item["end"]),
            label=str(item["label"]),
        )
        for item in raw_segments
    ]


def compare_shadow_outputs(
    reference: list[ChordInterval],
    predictions: dict[str, list[ChordInterval]],
) -> dict[str, Any]:
    providers = {
        name: evaluate_temporal_chords(reference, timeline)
        for name, timeline in predictions.items()
    }
    ranking = sorted(
        providers,
        key=lambda name: (
            float(providers[name]["exact_wcsr"]),
            float(providers[name]["root_wcsr"]),
            float(providers[name]["boundaries"]["f1"]),
        ),
        reverse=True,
    )
    return {
        "schema_version": "shadow-chord-benchmark-v1",
        "winner": ranking[0] if ranking else None,
        "ranking": ranking,
        "providers": providers,
    }


def parse_prediction(value: str) -> tuple[str, Path]:
    name, separator, path = value.partition("=")
    if not separator or not name.strip() or not path.strip():
        raise argparse.ArgumentTypeError("Use NOME=CAMINHO_JSON")
    return name.strip(), Path(path.strip())


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Compara saidas shadow sem alterar respostas dos endpoints.",
    )
    parser.add_argument("--reference", type=Path, required=True)
    parser.add_argument(
        "--prediction",
        action="append",
        type=parse_prediction,
        required=True,
        help="Saida de um motor no formato NOME=CAMINHO_JSON.",
    )
    args = parser.parse_args()
    predictions = {name: load_timeline(path) for name, path in args.prediction}
    result = compare_shadow_outputs(load_timeline(args.reference), predictions)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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