evaluation/orchestrator.py

"""Load a model family backend and run WER across packed conditions."""

from __future__ import annotations

import os
from collections.abc import Callable
from functools import lru_cache
from typing import Any

import torch

from benchmark.dataset import (
    accumulate_predictions_for_slice,
    evaluate_condition_wer_timed,
    is_full_condition_set,
    load_packed_condition,
    resolve_condition_keys,
    wer_result_key,
)
from backends.registry import build_transcriber, resolve_label

from evaluation.runtime import (
    apply_cpu_thread_settings_once,
    disable_broken_torchcodec,
    resolve_eval_devices,
    use_spaces_gpu_decorator,
)


def _hub_max_duration_s() -> int:
    """Upper bound the Hub accepts for ``spaces.GPU(duration=...)`` (override if HF raises)."""
    try:
        return max(60, int(os.environ.get("FFASR_ZEROGPU_HUB_MAX_DURATION_S", "600")))
    except ValueError:
        return 600


def _requested_duration_s() -> int:
    try:
        return max(60, int(os.environ.get("FFASR_ZEROGPU_MAX_DURATION_S", "600")))
    except ValueError:
        return 600


def _effective_gpu_duration_s() -> int:
    """``min(FFASR_ZEROGPU_MAX_DURATION_S, FFASR_ZEROGPU_HUB_MAX_DURATION_S)``."""
    return min(_requested_duration_s(), _hub_max_duration_s())


def _samples_per_gpu_segment() -> int:
    """
    When >0 and ``spaces`` is available, each condition is split into slices of this many
    samples; each slice runs in its own ``spaces.GPU`` call (model reload per slice).
    """
    try:
        v = int(os.environ.get("FFASR_ZEROGPU_SAMPLES_PER_SEGMENT", "0"))
    except ValueError:
        return 0
    return max(0, v)


def _build_transcriber(
    family_id: str,
    model_id: str,
    device_str: str,
    device_int: int,
    custom_script: str = "",
) -> tuple[Callable[..., str], Callable[[], None]]:
    """Family backend transcriber, or user ``evaluate(Path)`` script when provided."""
    if (custom_script or "").strip():
        from backends.custom_eval import build_transcriber_from_custom_script

        return build_transcriber_from_custom_script(custom_script)
    return build_transcriber(family_id, model_id, device_str, device_int)


def _zerogpu_run_segment(payload: dict[str, Any]) -> dict[str, Any]:
    """Single ZeroGPU segment: build model, transcribe ``[start:end)`` for one condition."""
    model_id = str(payload["model_id"])
    family_id = str(payload["family_id"])
    condition_key = str(payload["condition_key"])
    start = int(payload["start"])
    end = int(payload["end"])
    custom_script = str(payload.get("custom_script") or "")

    device_str, device_int = resolve_eval_devices()
    transcribe, cleanup = _build_transcriber(
        family_id, model_id, device_str, device_int, custom_script=custom_script
    )
    try:
        num_params = int(getattr(transcribe, "_num_params", 0) or 0)
        preds, refs, audio_s, infer_s, n_done = accumulate_predictions_for_slice(
            transcribe, condition_key, start, end, progress_cb=None
        )
        return {
            "predictions": preds,
            "references": refs,
            "audio_s": audio_s,
            "infer_s": infer_s,
            "n": n_done,
            "num_params": num_params,
        }
    finally:
        cleanup()
        if device_str == "cuda" and torch.cuda.is_available():
            torch.cuda.empty_cache()


def _run_evaluation_core(
    model_id: str,
    family_id: str = "auto",
    progress_cb: Callable[[int, int, str], None] | None = None,
    custom_script: str = "",
    condition_keys: list[str] | None = None,
) -> dict:
    device_str, device_int = resolve_eval_devices()
    keys = resolve_condition_keys(condition_keys)

    # Precompute totals so the progress bar has a known denominator from the start.
    condition_totals: dict[str, int] = {}
    grand_total = 0
    for ck in keys:
        try:
            condition_totals[ck] = len(load_packed_condition(ck))
        except Exception:
            condition_totals[ck] = 0
        grand_total += condition_totals[ck]

    if progress_cb is not None:
        try:
            progress_cb(0, grand_total, "")
        except Exception:
            pass

    transcribe, cleanup = _build_transcriber(
        family_id, model_id, device_str, device_int, custom_script=custom_script
    )
    try:
        # Backends attach `transcribe._num_params` via `_model_utils.attach_params`; defaults to 0
        # if the backend couldn't introspect a module.
        num_params = int(getattr(transcribe, "_num_params", 0) or 0)
        results: dict = {
            "model_id": model_id,
            "eval_family": resolve_label(family_id),
            "num_params": num_params,
        }
        total_samples = 0
        total_audio_s = 0.0
        total_infer_s = 0.0
        done_so_far = 0

        for condition_key in keys:
            condition_base = done_so_far

            def _sample_cb(ck: str, i: int, _n: int, _base: int = condition_base) -> None:
                if progress_cb is None:
                    return
                try:
                    progress_cb(_base + i, grand_total, ck)
                except Exception:
                    pass

            wer_score, count, audio_s, infer_s = evaluate_condition_wer_timed(
                transcribe, condition_key, progress_cb=_sample_cb
            )
            results[wer_result_key(condition_key)] = wer_score
            total_samples = max(total_samples, count)
            total_audio_s += audio_s
            total_infer_s += infer_s
            done_so_far += count

        results["eval_conditions"] = list(keys)
        results["eval_full_run"] = is_full_condition_set(condition_keys)
        results["num_samples"] = total_samples
        results["eval_audio_seconds"] = round(total_audio_s, 3)
        results["eval_wall_time_s"] = round(total_infer_s, 3)
        # RTF: audio duration / compute time (>1 means faster than real-time for batch totals)
        if total_infer_s > 1e-6:
            results["eval_rtf"] = round(total_audio_s / total_infer_s, 4)
        else:
            results["eval_rtf"] = ""

        if progress_cb is not None:
            try:
                progress_cb(done_so_far, grand_total, "")
            except Exception:
                pass
        return results
    finally:
        cleanup()
        if device_str == "cuda" and torch.cuda.is_available():
            torch.cuda.empty_cache()


def _run_evaluation_core_segmented_local(
    model_id: str,
    family_id: str = "auto",
    progress_cb: Callable[[int, int, str], None] | None = None,
    *,
    segment_samples: int,
    custom_script: str = "",
    condition_keys: list[str] | None = None,
) -> dict:
    """
    Same semantics as ``_run_evaluation_core_segmented`` but without ``spaces.GPU``:
    each slice calls ``_zerogpu_run_segment`` in-process (CPU or CUDA per ``FFASR_DEVICE``).
    """
    from jiwer import wer as compute_wer

    keys = resolve_condition_keys(condition_keys)
    condition_totals: dict[str, int] = {}
    grand_total = 0
    for ck in keys:
        try:
            condition_totals[ck] = len(load_packed_condition(ck))
        except Exception:
            condition_totals[ck] = 0
        grand_total += condition_totals[ck]

    if progress_cb is not None:
        try:
            progress_cb(0, grand_total, "")
        except Exception:
            pass

    results: dict = {
        "model_id": model_id,
        "eval_family": resolve_label(family_id),
        "num_params": 0,
    }
    total_audio_s = 0.0
    total_infer_s = 0.0
    total_samples = 0
    done_so_far = 0
    first_num_params: int | None = None

    for condition_key in keys:
        data = load_packed_condition(condition_key)
        n_items = len(data)
        all_p: list[str] = []
        all_r: list[str] = []
        cond_audio = 0.0
        cond_infer = 0.0

        if n_items == 0:
            results[wer_result_key(condition_key)] = 0.0
            continue

        for start in range(0, n_items, segment_samples):
            end = min(n_items, start + segment_samples)
            payload: dict[str, Any] = {
                "model_id": model_id,
                "family_id": family_id,
                "condition_key": condition_key,
                "start": start,
                "end": end,
                "custom_script": custom_script,
            }
            out = _zerogpu_run_segment(payload)
            if first_num_params is None:
                first_num_params = int(out.get("num_params") or 0)
            all_p.extend(out["predictions"])
            all_r.extend(out["references"])
            cond_audio += float(out["audio_s"])
            cond_infer += float(out["infer_s"])
            done_so_far += int(out["n"])
            if progress_cb is not None:
                try:
                    progress_cb(done_so_far, grand_total, condition_key)
                except Exception:
                    pass

        wer_score = round(compute_wer(all_r, all_p), 4) if all_p else 0.0
        results[wer_result_key(condition_key)] = wer_score
        total_audio_s += cond_audio
        total_infer_s += cond_infer
        total_samples = max(total_samples, len(all_p))

    results["num_params"] = first_num_params or 0
    results["eval_conditions"] = list(keys)
    results["eval_full_run"] = is_full_condition_set(condition_keys)
    results["num_samples"] = total_samples
    results["eval_audio_seconds"] = round(total_audio_s, 3)
    results["eval_wall_time_s"] = round(total_infer_s, 3)
    if total_infer_s > 1e-6:
        results["eval_rtf"] = round(total_audio_s / total_infer_s, 4)
    else:
        results["eval_rtf"] = ""

    if progress_cb is not None:
        try:
            progress_cb(done_so_far, grand_total, "")
        except Exception:
            pass
    return results


def _run_evaluation_core_segmented(
    model_id: str,
    family_id: str = "auto",
    progress_cb: Callable[[int, int, str], None] | None = None,
    *,
    segment_samples: int,
    gpu_duration_s: int,
    gpu_size: str,
    custom_script: str = "",
    condition_keys: list[str] | None = None,
) -> dict:
    """
    Long evaluations: each slice of ``segment_samples`` runs under its own ``spaces.GPU``
    budget, then predictions/references are merged and WER is computed once per condition.
    """
    from jiwer import wer as compute_wer

    import spaces

    keys = resolve_condition_keys(condition_keys)
    condition_totals: dict[str, int] = {}
    grand_total = 0
    for ck in keys:
        try:
            condition_totals[ck] = len(load_packed_condition(ck))
        except Exception:
            condition_totals[ck] = 0
        grand_total += condition_totals[ck]

    if progress_cb is not None:
        try:
            progress_cb(0, grand_total, "")
        except Exception:
            pass

    wrapped_seg = _spaces_gpu_segment_wrapper(gpu_duration_s, gpu_size)

    results: dict = {
        "model_id": model_id,
        "eval_family": resolve_label(family_id),
        "num_params": 0,
    }
    total_audio_s = 0.0
    total_infer_s = 0.0
    total_samples = 0
    done_so_far = 0
    first_num_params: int | None = None

    for condition_key in keys:
        data = load_packed_condition(condition_key)
        n_items = len(data)
        all_p: list[str] = []
        all_r: list[str] = []
        cond_audio = 0.0
        cond_infer = 0.0

        if n_items == 0:
            results[wer_result_key(condition_key)] = 0.0
            continue

        for start in range(0, n_items, segment_samples):
            end = min(n_items, start + segment_samples)
            payload: dict[str, Any] = {
                "model_id": model_id,
                "family_id": family_id,
                "condition_key": condition_key,
                "start": start,
                "end": end,
                "custom_script": custom_script,
            }
            out = wrapped_seg(payload)
            if first_num_params is None:
                first_num_params = int(out.get("num_params") or 0)
            all_p.extend(out["predictions"])
            all_r.extend(out["references"])
            cond_audio += float(out["audio_s"])
            cond_infer += float(out["infer_s"])
            done_so_far += int(out["n"])
            if progress_cb is not None:
                try:
                    progress_cb(done_so_far, grand_total, condition_key)
                except Exception:
                    pass

        wer_score = round(compute_wer(all_r, all_p), 4) if all_p else 0.0
        results[wer_result_key(condition_key)] = wer_score
        total_audio_s += cond_audio
        total_infer_s += cond_infer
        total_samples = max(total_samples, len(all_p))

    results["num_params"] = first_num_params or 0
    results["eval_conditions"] = list(keys)
    results["eval_full_run"] = is_full_condition_set(condition_keys)
    results["num_samples"] = total_samples
    results["eval_audio_seconds"] = round(total_audio_s, 3)
    results["eval_wall_time_s"] = round(total_infer_s, 3)
    if total_infer_s > 1e-6:
        results["eval_rtf"] = round(total_audio_s / total_infer_s, 4)
    else:
        results["eval_rtf"] = ""

    if progress_cb is not None:
        try:
            progress_cb(done_so_far, grand_total, "")
        except Exception:
            pass
    return results


@lru_cache(maxsize=16)
def _spaces_gpu_wrapped(duration_s: int, size: str):
    import spaces

    return spaces.GPU(duration=duration_s, size=size)(_run_evaluation_core)


@lru_cache(maxsize=16)
def _spaces_gpu_segment_wrapper(duration_s: int, size: str):
    import spaces

    return spaces.GPU(duration=duration_s, size=size)(_zerogpu_run_segment)


def run_evaluation(
    model_id: str,
    family_id: str = "auto",
    progress_cb: Callable[[int, int, str], None] | None = None,
    custom_script: str = "",
    condition_keys: list[str] | None = None,
) -> dict:
    """
    Evaluate `model_id` on all packed conditions using the selected family backend.

    On Hugging Face **Spaces ZeroGPU**, evaluation runs under ``spaces.GPU``:

    - ``FFASR_ZEROGPU_MAX_DURATION_S`` — requested max GPU seconds per **decorated** call
      (default **600**). Capped by ``FFASR_ZEROGPU_HUB_MAX_DURATION_S`` (default **600**) because
      the Hub rejects overly large values (e.g. 1800s).
    - ``FFASR_ZEROGPU_GPU_SIZE`` — ``large`` (default) or ``xlarge``.
    - ``FFASR_ZEROGPU_SAMPLES_PER_SEGMENT`` — when **>0**, each condition is split into chunks of
      this many samples; **each chunk** uses its own ``spaces.GPU`` call with the duration cap
      (the model is **reloaded** every chunk). Merge is exact for WER (full-condition preds/refs).

    If the ``spaces`` package is not installed (local dev), evaluation runs without the decorator.

    **CPU / device overrides** (see also ``evaluation/runtime.py``):

    - ``FFASR_DEVICE`` — ``auto`` (default), ``cpu``, or ``cuda``. Forces CPU on GPU hosts when set to ``cpu``.
    - ``FFASR_DISABLE_ZEROGPU`` — when ``1``/``true``, never use ``spaces.GPU`` even if ``spaces`` is installed;
      segmented eval (``FFASR_ZEROGPU_SAMPLES_PER_SEGMENT`` > 0) runs in-process slices instead.
    - ``FFASR_TORCH_NUM_THREADS`` / ``FFASR_TORCH_NUM_INTEROP_THREADS`` — optional torch thread caps (applied once per process).

    If ``progress_cb`` is provided, it is called as
    ``progress_cb(samples_done_across_all_conditions, samples_total_across_all_conditions, current_condition_key)``
    periodically during the run.

    When ``custom_script`` is non-empty, it must define ``evaluate(file: pathlib.Path) -> str``;
    that function is called once per packed sample (via a temp WAV) instead of the family backend.

    Returns wer_clean, wer_noisy, wer_reverberant, wer_real, wer_difficult, num_samples, model_id, eval_family, plus timing.
    """
    apply_cpu_thread_settings_once()
    disable_broken_torchcodec()

    seg = _samples_per_gpu_segment()

    if not use_spaces_gpu_decorator():
        if seg > 0:
            return _run_evaluation_core_segmented_local(
                model_id,
                family_id=family_id,
                progress_cb=progress_cb,
                segment_samples=seg,
                custom_script=custom_script,
                condition_keys=condition_keys,
            )
        return _run_evaluation_core(
            model_id,
            family_id=family_id,
            progress_cb=progress_cb,
            custom_script=custom_script,
            condition_keys=condition_keys,
        )

    eff = _effective_gpu_duration_s()
    size = os.environ.get("FFASR_ZEROGPU_GPU_SIZE", "large").strip().lower()
    if size not in ("large", "xlarge"):
        size = "large"

    if seg <= 0:
        wrapped = _spaces_gpu_wrapped(eff, size)
        return wrapped(
            model_id,
            family_id=family_id,
            progress_cb=progress_cb,
            custom_script=custom_script,
            condition_keys=condition_keys,
        )

    return _run_evaluation_core_segmented(
        model_id,
        family_id=family_id,
        progress_cb=progress_cb,
        segment_samples=seg,
        gpu_duration_s=eff,
        gpu_size=size,
        custom_script=custom_script,
        condition_keys=condition_keys,
    )