score_benchmark.py

import argparse
import os
import site
import subprocess
import sys
import time
from pathlib import Path


STAGES = [
    {
        "key": "scoreq",
        "label": "SCOREQ",
        "script": "score_scoreq.py",
        "outputs": ["scores.csv", "scores_metadata.json"],
        "aliases": {"scoreq", "scores", "quality"},
    },
    {
        "key": "wer",
        "label": "WER",
        "script": "score_asr_wer.py",
        "outputs": [
            "reference_transcripts.csv",
            "reference_transcripts_metadata.json",
            "transcripts.csv",
            "wer.csv",
            "wer_metadata.json",
        ],
        "aliases": {"wer", "asr"},
    },
    {
        "key": "dnsmos",
        "label": "DNSMOS P.835",
        "script": "score_dnsmos.py",
        "outputs": ["dnsmos.csv", "dnsmos_metadata.json"],
        "aliases": {"dnsmos", "mos", "p835"},
    },
]


def format_seconds(seconds: float) -> str:
    total_seconds = int(round(seconds))
    minutes, remaining_seconds = divmod(total_seconds, 60)
    hours, remaining_minutes = divmod(minutes, 60)
    if hours:
        return f"{hours}h {remaining_minutes}m {remaining_seconds}s"
    if minutes:
        return f"{minutes}m {remaining_seconds}s"
    return f"{remaining_seconds}s"


def benchmark_env() -> dict[str, str]:
    env = os.environ.copy()
    site_packages = Path(site.getusersitepackages())
    runtime_dirs = [
        site_packages / "nvidia" / "cudnn" / "lib",
        site_packages / "nvidia" / "cublas" / "lib",
        site_packages / "nvidia" / "cuda_runtime" / "lib",
    ]
    existing = [str(path) for path in runtime_dirs if path.is_dir()]
    if existing:
        current = env.get("LD_LIBRARY_PATH", "")
        env["LD_LIBRARY_PATH"] = ":".join(existing + ([current] if current else []))
    return env


def run_stage(script_name: str, env: dict[str, str]) -> None:
    subprocess.run([sys.executable, script_name], check=True, env=env)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run Voices for Christ benchmark metrics.")
    parser.add_argument(
        "-m",
        "--metrics",
        nargs="+",
        default=["all"],
        help="Metrics to run: all, scoreq, wer, dnsmos.",
    )
    parser.add_argument(
        "--list-metrics",
        action="store_true",
        help="Print available metrics and exit.",
    )
    return parser.parse_args()


def resolve_stages(metric_names: list[str]) -> list[dict[str, object]]:
    normalized = {metric_name.strip().lower() for metric_name in metric_names if metric_name.strip()}
    if not normalized or "all" in normalized:
        return STAGES

    selected: list[dict[str, object]] = []
    seen_keys: set[str] = set()
    for stage in STAGES:
        aliases = set(stage["aliases"]) | {stage["key"]}
        if normalized & aliases and stage["key"] not in seen_keys:
            selected.append(stage)
            seen_keys.add(stage["key"])

    known = {"all"}
    for stage in STAGES:
        known.update(stage["aliases"])
        known.add(stage["key"])
    unknown = sorted(normalized - known)
    if unknown:
        raise SystemExit(f"Unknown metric(s): {', '.join(unknown)}")
    if not selected:
        raise SystemExit("No metrics selected.")
    return selected


def main() -> None:
    args = parse_args()
    if args.list_metrics:
        print("Available metrics:")
        for stage in STAGES:
            aliases = ", ".join(sorted(stage["aliases"]))
            print(f"- {stage['key']}: {stage['label']} ({aliases})")
        return

    benchmark_root = Path(__file__).resolve().parent
    env = benchmark_env()
    stages = resolve_stages(args.metrics)

    overall_start = time.perf_counter()
    total_stages = len(stages)

    selected_labels = ", ".join(stage["label"] for stage in stages)
    print(f"Running Voices for Christ benchmark: {selected_labels}")
    for index, stage in enumerate(stages, start=1):
        stage_start = time.perf_counter()
        print(f"\n[{index}/{total_stages}] Starting {stage['label']}")
        run_stage(stage["script"], env)
        elapsed = time.perf_counter() - stage_start
        outputs = ", ".join(str(benchmark_root / output_name) for output_name in stage["outputs"])
        print(f"[{index}/{total_stages}] Finished {stage['label']} in {format_seconds(elapsed)}")
        print(f"[{index}/{total_stages}] Outputs: {outputs}")

    total_elapsed = time.perf_counter() - overall_start
    print(f"\nBenchmark complete in {format_seconds(total_elapsed)}")


if __name__ == "__main__":
    main()