scripts/run_eval_and_publish.py

#!/usr/bin/env python3
"""
Run evaluation locally against locally-staged packed conditions, and write the
two CSVs the Space normally publishes to the bucket as plain files on disk:

  - <out-dir>/leaderboard.csv  (same schema as results/leaderboard.csv)
  - <out-dir>/jobs_state.csv   (same schema as results/jobs_state.csv)

Nothing is pushed to the Hub. The packed test tensors are read from
`--packed-dir` instead of being downloaded from the bucket — point it at
`data/new_local_generated/packed_small/`, which mirrors `data/packed/` in the
bucket. Files expected: clean.pt, noisy.pt, reverberant.pt (others are ignored).

Apple Silicon / GPU acceleration:
  * `--device auto` (default) picks MPS (Apple GPU) on M-series Macs, then CUDA, then CPU.
  * `--batch-size N` enables HF-pipeline batched inference (used by the `auto`
    cascade for Whisper / Wav2Vec2-pipeline models). Non-pipeline backends
    (Granite chat, SpeechBrain, custom CTC) ignore batching and run sample-by-sample.

Usage:
  python scripts/run_eval_and_publish.py --model openai/whisper-tiny
  python scripts/run_eval_and_publish.py --model openai/whisper-small \\
      --device mps --batch-size 16
  python scripts/run_eval_and_publish.py --model org/model --family transformers_ctc \\
      --packed-dir /path/to/packed_small --out-dir ./local_results
"""

from __future__ import annotations

import argparse
import csv
import io
import json
import os
import sys
import time
import traceback
import uuid
import warnings
warnings.filterwarnings("ignore", category=UserWarning)  # HF tokenizer parallelism
from datetime import datetime, timezone

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

DEFAULT_PACKED_DIR = "/Users/shivamsaini/speech-rir-asr-dataset/data/new_local_generated/packed_small"
DEFAULT_OUT_DIR = os.path.join(ROOT, "local_results")


# ---------------------------------------------------------------------------
# Device selection (CUDA / MPS / CPU)
# ---------------------------------------------------------------------------


def _resolve_device(choice: str) -> tuple[str, int]:
    """Return (device_str, device_int) for orchestrator / pipeline APIs.

    `device_int` is the value HF pipeline expects: -1 for CPU, 0+ for CUDA.
    For MPS we still pass `device_str="mps"` to the pipeline (it accepts strings),
    and report device_int=-1 so any code that branches on it treats it as non-CUDA.
    """
    import torch

    choice = (choice or "auto").lower()
    has_cuda = torch.cuda.is_available()
    has_mps = bool(getattr(torch.backends, "mps", None) and torch.backends.mps.is_available())

    if choice == "auto":
        if has_mps:
            choice = "mps"
        elif has_cuda:
            choice = "cuda"
        else:
            choice = "cpu"

    if choice == "cuda":
        if not has_cuda:
            raise RuntimeError("CUDA requested but torch.cuda.is_available() is False.")
        return "cuda", 0
    if choice == "mps":
        if not has_mps:
            raise RuntimeError(
                "MPS requested but torch.backends.mps.is_available() is False. "
                "Need PyTorch with MPS support on Apple Silicon."
            )
        # Improve coverage: fall back to CPU for ops MPS doesn't implement.
        os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
        return "mps", -1
    if choice == "cpu":
        return "cpu", -1
    raise ValueError(f"--device must be one of auto/cpu/cuda/mps; got {choice!r}")


# ---------------------------------------------------------------------------
# Local packed-file loader (replaces `download_bucket_file` for `.pt` reads)
# ---------------------------------------------------------------------------


def _patch_packed_loader(packed_dir: str) -> None:
    """Redirect benchmark.dataset.load_packed_condition to read local .pt files.

    Must run BEFORE importing orchestrator/run_evaluation so the orchestrator
    sees the patched function.
    """
    import torch

    from benchmark import dataset as bench_ds

    cache: dict[str, dict] = {}

    def _local_loader(condition_key: str) -> dict:
        if condition_key in cache:
            return cache[condition_key]
        rel = bench_ds.PACKED_FILES[condition_key]
        local_path = os.path.join(packed_dir, os.path.basename(rel))
        if not os.path.isfile(local_path):
            raise FileNotFoundError(
                f"Packed file for condition {condition_key!r} not found: {local_path}"
            )
        data = torch.load(local_path, weights_only=False)
        cache[condition_key] = data
        return data

    bench_ds.load_packed_condition = _local_loader


# ---------------------------------------------------------------------------
# Inline orchestrator: explicit device + optional HF-pipeline batching
# ---------------------------------------------------------------------------


def _resolve_dtype(choice: str, device_str: str):
    """Pick torch dtype.

    `auto` policy:
      * CUDA -> fp16 (fast + generally stable)
      * MPS  -> fp32 (several speech models overflow / misbehave in fp16 on Apple GPU)
      * CPU  -> fp32
    """
    import torch

    choice = (choice or "auto").lower()
    if choice == "auto":
        return torch.float16 if device_str == "cuda" else torch.float32
    return {
        "float32": torch.float32,
        "fp32": torch.float32,
        "float16": torch.float16,
        "fp16": torch.float16,
        "half": torch.float16,
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
    }.get(choice, torch.float32)


_SEQ2SEQ_INCOMPATIBLE_MODEL_TYPES = {
    # Chat-prompted speech models — their processors require a `text` arg with the
    # chat template, so the audio-only processor call we use below is wrong.
    # Their dedicated backends know the right prompt; let those handle it.
    "granite_speech",
    "qwen2_audio",
    "qwen_audio",
}

# When `--family auto`, route specific model_type values to their dedicated
# backend instead of letting the auto cascade fall through pipeline → universal,
# which loads but crashes at inference for chat-prompted speech models, or
# crashes at batched-collation time for models with variable-length features
# the HF pipeline padding can't reconcile (Cohere ASR).
_MODEL_TYPE_TO_FAMILY = {
    "granite_speech": "granite_speech",
    "cohere_asr": "universal",
}


def _resolve_family_for_model(family_id: str, model_id: str) -> tuple[str, str]:
    """Return (resolved_family, model_type). If `family_id` is `auto` and the
    config exposes a model_type that has a dedicated backend, switch to it."""
    try:
        from transformers import AutoConfig

        cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
        model_type = getattr(cfg, "model_type", "") or ""
    except Exception:
        return family_id, ""
    if family_id == "auto" and model_type in _MODEL_TYPE_TO_FAMILY:
        return _MODEL_TYPE_TO_FAMILY[model_type], model_type
    return family_id, model_type


def _try_build_batched_seq2seq(model_id: str, device_str: str, dtype):
    """Hand-rolled batched seq2seq transcriber (Whisper-style).

    Bypasses the HF pipeline so encoder/decoder forward passes are actually
    batched, preprocessing pads once per batch, and we can pick fp16. Works for
    any model loadable via `AutoModelForSpeechSeq2Seq`.

    Returns (transcribe_batch, num_params, cleanup, info) or None when this
    path doesn't apply (model needs chat prompts, can't load on this device,
    or the audio-only processor signature isn't supported).
    """
    import torch

    try:
        from transformers import AutoConfig, AutoModelForSpeechSeq2Seq, AutoProcessor
    except Exception:
        return None

    # Cheap pre-flight: read the config first and bail for known chat-prompted
    # speech models so we don't pay the model-load cost just to crash.
    try:
        cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
        cfg_model_type = getattr(cfg, "model_type", "") or ""
    except Exception:
        cfg_model_type = ""
    if cfg_model_type in _SEQ2SEQ_INCOMPATIBLE_MODEL_TYPES:
        return None

    try:
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    except Exception:
        return None

    # Some chat-style processors are missing `feature_extractor` entirely; their
    # `__call__` requires `text=`. Detecting this cleanly is hard, so we treat
    # the warm-up below as a probe: if it fails, we fall back.
    fe = getattr(processor, "feature_extractor", None)
    target_sr = int(getattr(fe, "sampling_rate", 0)) or 16000

    try:
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id,
            torch_dtype=dtype,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
    except Exception:
        return None

    if not hasattr(model, "generate"):
        del model
        return None

    try:
        model.to(device_str)
    except Exception:
        # Some ops not implemented on MPS; bail and let the pipeline path take over.
        del model
        return None
    model.eval()

    # Confirm placement (and fall back if torch silently kept it on CPU).
    try:
        param_device = str(next(model.parameters()).device)
    except StopIteration:
        param_device = "?"

    # Whisper-only kwargs are tried first, then fall back to plain generate().
    is_whisper = getattr(getattr(model, "config", None), "model_type", "") == "whisper"

    def _generate(input_features, attention_mask=None):
        gen_kwargs = {"max_new_tokens": 440}
        if attention_mask is not None:
            gen_kwargs["attention_mask"] = attention_mask
        if cfg_model_type == "cohere_asr":
            # Cohere's remote generation code expects this key to be a tensor,
            # not None. Without it, HF generation crashes when extending masks.
            gen_kwargs["decoder_attention_mask"] = torch.ones(
                (int(input_features.shape[0]), 1),
                device=input_features.device,
                dtype=torch.long,
            )
        if is_whisper:
            try:
                with torch.inference_mode():
                    return model.generate(
                        input_features,
                        language="english",
                        task="transcribe",
                        **gen_kwargs,
                    )
            except (TypeError, ValueError):
                pass
        with torch.inference_mode():
            return model.generate(input_features, **gen_kwargs)

    def transcribe_batch(items: list[dict], batch_size: int) -> list[str]:
        audios = [it["raw"] for it in items]
        # Packed test data is already 16k mono; processor handles padding.
        inputs = processor(audios, sampling_rate=target_sr, return_tensors="pt", padding=True)
        if "input_features" in inputs:
            feats = inputs["input_features"].to(device_str, dtype=dtype)
        elif "input_values" in inputs:
            feats = inputs["input_values"].to(device_str, dtype=dtype)
        else:
            raise RuntimeError(f"Processor returned unexpected keys: {list(inputs.keys())}")
        attn = inputs.get("attention_mask")
        if attn is not None:
            attn = attn.to(device_str)
        ids = _generate(feats, attention_mask=attn)
        return processor.batch_decode(ids, skip_special_tokens=True)

    # Warm-up doubles as a probe: if the processor refuses an audio-only call
    # (e.g. chat-prompted speech models we didn't catch above), fall back so
    # the dedicated backend can take over.
    import numpy as np

    try:
        warm = [{"raw": np.zeros(target_sr, dtype=np.float32), "sampling_rate": target_sr}] * 2
        transcribe_batch(warm, batch_size=2)
    except Exception as e:
        print(
            f"[eval] seq2seq path rejected for {model_id} ({type(e).__name__}: {e}); "
            "falling back to pipeline / per-sample.",
            file=sys.stderr,
        )
        del model, processor
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return None

    from backends._model_utils import count_params

    num_params = count_params(model)

    def cleanup() -> None:
        nonlocal model, processor
        del model, processor

    info = {
        "kind": "seq2seq",
        "param_device": param_device,
        "dtype": str(dtype).replace("torch.", ""),
        "model_type": getattr(getattr(model, "config", None), "model_type", "?"),
        "target_sr": target_sr,
    }
    return transcribe_batch, num_params, cleanup, info


def _try_build_batched_pipeline(model_id: str, device_str: str):
    """Fallback batched path via the HF pipeline. Less aggressive batching than
    the hand-rolled seq2seq path but works for CTC models too."""
    # Cohere ASR pipeline batching currently breaks in HF pipeline collation for
    # variable-length feature tensors. Keep pipeline available for per-sample use,
    # but avoid this batched path so seq2seq/per-sample code handles it.
    try:
        from transformers import AutoConfig

        cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
        if getattr(cfg, "model_type", "") == "cohere_asr":
            return None
    except Exception:
        pass

    try:
        from transformers import pipeline
    except Exception:
        return None

    try:
        pipe = pipeline(
            "automatic-speech-recognition",
            model=model_id,
            device=device_str,
            trust_remote_code=True,
        )
    except Exception:
        return None

    import numpy as np

    try:
        pipe({"raw": np.zeros(16000, dtype=np.float32), "sampling_rate": 16000})
    except Exception:
        pass

    from backends._model_utils import count_params

    model_obj = getattr(pipe, "model", None)
    num_params = count_params(model_obj)
    try:
        param_device = str(next(model_obj.parameters()).device) if model_obj else "?"
    except Exception:
        param_device = "?"

    def transcribe_batch(audios_with_sr: list[dict], batch_size: int) -> list[str]:
        outs = pipe(audios_with_sr, batch_size=batch_size)
        return [str(o["text"]) for o in outs]

    def cleanup() -> None:
        nonlocal pipe
        del pipe

    info = {
        "kind": "pipeline",
        "param_device": param_device,
        "dtype": "float32",
        "model_type": getattr(getattr(model_obj, "config", None), "model_type", "?"),
    }
    return transcribe_batch, num_params, cleanup, info


def _eval_condition_batched(
    transcribe_batch,
    condition_key: str,
    batch_size: int,
    progress_cb,
):
    """Batched WER eval — calls `transcribe_batch(list[dict], batch_size)` once per condition."""
    from jiwer import wer as compute_wer

    from backends._audio_utils import safe_pad_audio
    from benchmark.dataset import _normalize_for_wer, load_packed_condition

    data = load_packed_condition(condition_key)
    if not data:
        if progress_cb is not None:
            progress_cb(condition_key, 0, 0)
        return 0.0, 0, 0.0, 0.0

    items = list(data.items())
    inputs: list[dict] = []
    references: list[str] = []
    audio_seconds = 0.0
    for _sid, sample in items:
        audio_np = safe_pad_audio(sample["waveform"].numpy())
        sr = int(sample.get("sample_rate", 16000))
        audio_seconds += float(len(audio_np)) / float(sr)
        inputs.append({"raw": audio_np, "sampling_rate": sr})
        references.append(_normalize_for_wer(str(sample["transcript"])))

    total = len(items)
    predictions: list[str] = []
    t0 = time.perf_counter()
    # Process in chunks so we can report progress between batches.
    for start in range(0, total, batch_size):
        chunk = inputs[start : start + batch_size]
        texts = transcribe_batch(chunk, batch_size=batch_size)
        predictions.extend(_normalize_for_wer(str(t)) for t in texts)
        if progress_cb is not None:
            try:
                progress_cb(condition_key, min(start + len(chunk), total), total)
            except Exception:
                pass
    inference_seconds = time.perf_counter() - t0

    return round(compute_wer(references, predictions), 4), len(references), audio_seconds, inference_seconds


def _run_evaluation_local(
    model_id: str,
    family_id: str,
    device_str: str,
    device_int: int,
    batch_size: int,
    dtype_choice: str,
    progress_cb,
) -> dict:
    """Inline orchestrator: chosen device + batched fast-path when possible."""
    import torch

    from evaluation.runtime import disable_broken_torchcodec

    disable_broken_torchcodec()

    from backends.registry import build_transcriber, resolve_label
    from benchmark.dataset import (
        PACKED_FILES,
        evaluate_condition_wer_timed,
        wer_result_key,
        load_packed_condition,
    )

    dtype = _resolve_dtype(dtype_choice, device_str)

    resolved_family, model_type = _resolve_family_for_model(family_id, model_id)
    if resolved_family != family_id:
        print(
            f"[eval] family override: {family_id} -> {resolved_family} "
            f"(detected model_type={model_type!r})",
            file=sys.stderr,
        )
    family_id = resolved_family

    # Pre-compute totals so progress is meaningful from the start.
    condition_totals: dict[str, int] = {}
    grand_total = 0
    for ck in PACKED_FILES:
        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

    # Try the batched fast paths in order: hand-rolled seq2seq → HF pipeline.
    batched = None
    info = None
    if batch_size > 1 and family_id in ("auto", "transformers_pipeline", "universal"):
        built = _try_build_batched_seq2seq(model_id, device_str, dtype)
        if built is None:
            built = _try_build_batched_pipeline(model_id, device_str)
        if built is not None:
            transcribe_batch, num_params, cleanup, info = built
            batched = (transcribe_batch, num_params, cleanup)
            print(
                f"[eval] batched path: {info['kind']} | model_type={info.get('model_type','?')} | "
                f"param_device={info['param_device']} | dtype={info['dtype']} | batch_size={batch_size}",
                file=sys.stderr,
            )

    if batched is not None:
        transcribe_batch, num_params, cleanup = batched
        eval_family_label = resolve_label("transformers_pipeline")
        try:
            results: dict = {
                "model_id": model_id,
                "eval_family": eval_family_label,
                "num_params": int(num_params or 0),
            }
            total_samples = 0
            total_audio_s = 0.0
            total_infer_s = 0.0
            done_so_far = 0

            for condition_key in PACKED_FILES:
                base = done_so_far

                def _cb(ck: str, i: int, _n: int, _b: int = base) -> None:
                    if progress_cb is not None:
                        try:
                            progress_cb(_b + i, grand_total, ck)
                        except Exception:
                            pass

                wer, count, audio_s, infer_s = _eval_condition_batched(
                    transcribe_batch, condition_key, batch_size, progress_cb=_cb
                )
                results[wer_result_key(condition_key)] = wer
                total_samples = max(total_samples, count)
                total_audio_s += audio_s
                total_infer_s += infer_s
                done_so_far += count

            results["num_samples"] = total_samples
            results["eval_audio_seconds"] = round(total_audio_s, 3)
            results["eval_wall_time_s"] = round(total_infer_s, 3)
            results["eval_rtf"] = (
                round(total_audio_s / total_infer_s, 4) if total_infer_s > 1e-6 else ""
            )
            return results
        finally:
            cleanup()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    transcribe, cleanup = build_transcriber(family_id, model_id, device_str, device_int)
    try:
        num_params = int(getattr(transcribe, "_num_params", 0) or 0)
        backend_batch = getattr(transcribe, "_batch_transcribe", None)

        if batch_size > 1 and callable(backend_batch):
            print(
                f"[eval] backend-native batched path: family={family_id} device={device_str} "
                f"batch_size={batch_size}",
                file=sys.stderr,
            )
            results = {
                "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 PACKED_FILES:
                base = done_so_far

                def _cb(ck: str, i: int, _n: int, _b: int = base) -> None:
                    if progress_cb is not None:
                        try:
                            progress_cb(_b + i, grand_total, ck)
                        except Exception:
                            pass

                wer, count, audio_s, infer_s = _eval_condition_batched(
                    backend_batch, condition_key, batch_size, progress_cb=_cb
                )
                results[wer_result_key(condition_key)] = wer
                total_samples = max(total_samples, count)
                total_audio_s += audio_s
                total_infer_s += infer_s
                done_so_far += count
            results["num_samples"] = total_samples
            results["eval_audio_seconds"] = round(total_audio_s, 3)
            results["eval_wall_time_s"] = round(total_infer_s, 3)
            results["eval_rtf"] = (
                round(total_audio_s / total_infer_s, 4) if total_infer_s > 1e-6 else ""
            )
            return results

        # Per-sample fallback (any backend; honors device_str).
        print(
            f"[eval] per-sample path (no batching): family={family_id} device={device_str}",
            file=sys.stderr,
        )
        results = {
            "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 PACKED_FILES:
            base = done_so_far

            def _cb(ck: str, i: int, _n: int, _b: int = base) -> None:
                if progress_cb is not None:
                    try:
                        progress_cb(_b + i, grand_total, ck)
                    except Exception:
                        pass

            wer, count, audio_s, infer_s = evaluate_condition_wer_timed(
                transcribe, condition_key, progress_cb=_cb
            )
            results[wer_result_key(condition_key)] = wer
            total_samples = max(total_samples, count)
            total_audio_s += audio_s
            total_infer_s += infer_s
            done_so_far += count
        results["num_samples"] = total_samples
        results["eval_audio_seconds"] = round(total_audio_s, 3)
        results["eval_wall_time_s"] = round(total_infer_s, 3)
        results["eval_rtf"] = (
            round(total_audio_s / total_infer_s, 4) if total_infer_s > 1e-6 else ""
        )
        return results
    finally:
        cleanup()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ---------------------------------------------------------------------------
# CSV helpers (local files; same schema as the Space's bucket CSVs)
# ---------------------------------------------------------------------------


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _progress_print(done: int, total: int, condition: str) -> None:
    if total <= 0:
        sys.stderr.write(f"\r[{condition or 'preparing'}] preparing…")
    else:
        pct = 100.0 * done / total
        sys.stderr.write(
            f"\r[{condition or '—'}] {done}/{total} ({pct:5.1f}%)            "
        )
    sys.stderr.flush()


def _read_csv(path: str) -> list[dict]:
    if not os.path.isfile(path):
        return []
    with open(path, "r", encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def _write_csv(path: str, fields: list[str], rows: list[dict]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    buf = io.StringIO()
    w = csv.DictWriter(buf, fieldnames=fields, extrasaction="ignore")
    w.writeheader()
    for row in rows:
        w.writerow({k: row.get(k, "") for k in fields})
    with open(path, "w", encoding="utf-8", newline="") as f:
        f.write(buf.getvalue())


def _append_leaderboard_local(csv_path: str, result: dict, notes: str) -> dict:
    from analytics import sort_leaderboard_rows_inplace
    from init import CSV_FIELDS, leaderboard_row_from_eval_result, normalize_legacy_csv_row

    rows = _read_csv(csv_path)
    for row in rows:
        for k in CSV_FIELDS:
            row.setdefault(k, "")
        normalize_legacy_csv_row(row)
    if any(r.get("model_id") == result["model_id"] for r in rows):
        raise RuntimeError(
            f"Model {result['model_id']!r} already in {csv_path}; refusing to duplicate."
        )
    new_row = leaderboard_row_from_eval_result(result, _now_iso(), submission_notes=notes)
    normalize_legacy_csv_row(new_row)
    rows.append(new_row)

    sort_leaderboard_rows_inplace(rows)
    _write_csv(csv_path, CSV_FIELDS, rows)
    return new_row


_JOBS_CSV_FIELDS = [
    "job_id",
    "model_id",
    "family_id",
    "status",
    "created_at",
    "updated_at",
    "error",
    "submission_notes",
]


def _upsert_job_local(
    csv_path: str,
    *,
    job_id: str,
    model_id: str,
    family_id: str,
    status: str,
    created_at: str,
    error: str | None = None,
    submission_notes: str = "",
) -> None:
    rows = _read_csv(csv_path)
    by_id = {(r.get("job_id") or "").strip(): r for r in rows if r.get("job_id")}
    by_id[job_id] = {
        "job_id": job_id,
        "model_id": model_id,
        "family_id": family_id,
        "status": status,
        "created_at": created_at,
        "updated_at": _now_iso(),
        "error": (error or "").replace("\n", " ")[:2000],
        "submission_notes": (submission_notes or "").replace("\n", " ")[:4000],
    }
    merged = list(by_id.values())
    merged.sort(key=lambda r: r.get("created_at") or "")
    _write_csv(csv_path, _JOBS_CSV_FIELDS, merged)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Run FFASR evaluation locally against staged packed files; write CSVs to disk."
    )
    parser.add_argument("--model", required=True, help="Hugging Face model id (author/name).")
    parser.add_argument(
        "--family",
        default=None,
        help="Inference backend family (default: auto). Validated against backends.registry.",
    )
    parser.add_argument(
        "--device",
        default="auto",
        choices=("auto", "cpu", "cuda", "mps"),
        help="Compute device. `auto` picks MPS on Apple Silicon, then CUDA, then CPU.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=8,
        help="Batch size for HF-pipeline backends (whisper, wav2vec2 pipeline, …). "
        "Non-pipeline backends ignore this and run sample-by-sample. Use 1 to disable batching.",
    )
    parser.add_argument(
        "--dtype",
        default="auto",
        choices=("auto", "float32", "float16", "bfloat16", "fp32", "fp16", "bf16", "half"),
        help="Model dtype. `auto` uses fp16 on CUDA/MPS and fp32 on CPU. "
        "fp16 typically halves Whisper inference time on M-series Macs.",
    )
    parser.add_argument(
        "--packed-dir",
        default=DEFAULT_PACKED_DIR,
        help=f"Directory containing clean.pt / noisy.pt / reverberant.pt (default: {DEFAULT_PACKED_DIR}).",
    )
    parser.add_argument(
        "--out-dir",
        default=DEFAULT_OUT_DIR,
        help=f"Directory to write leaderboard.csv and jobs_state.csv (default: {DEFAULT_OUT_DIR}).",
    )
    parser.add_argument(
        "--leaderboard-csv",
        default=None,
        help="Override path for leaderboard.csv (default: <out-dir>/leaderboard.csv).",
    )
    parser.add_argument(
        "--jobs-csv",
        default=None,
        help="Override path for jobs_state.csv (default: <out-dir>/jobs_state.csv).",
    )
    parser.add_argument(
        "--notes",
        default="",
        help="Optional submission notes recorded with the leaderboard row.",
    )
    parser.add_argument(
        "--job-id",
        default=None,
        help="Reuse an existing job_id. Defaults to a fresh uuid.",
    )
    parser.add_argument("--json", action="store_true", help="Print result as JSON at the end.")
    args = parser.parse_args()

    packed_dir = os.path.abspath(args.packed_dir)
    if not os.path.isdir(packed_dir):
        print(f"ERROR: --packed-dir does not exist: {packed_dir}", file=sys.stderr)
        return 2
    for cond in ("clean.pt", "noisy.pt", "reverberant.pt"):
        if not os.path.isfile(os.path.join(packed_dir, cond)):
            print(f"ERROR: missing {cond} in {packed_dir}", file=sys.stderr)
            return 2

    _patch_packed_loader(packed_dir)

    try:
        device_str, device_int = _resolve_device(args.device)
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    from backends.registry import FAMILY_IDS, default_family_id  # noqa: E402

    family_id = args.family or default_family_id()
    if family_id not in FAMILY_IDS:
        print(
            f"ERROR: invalid --family {family_id!r}; choose from {list(FAMILY_IDS)}",
            file=sys.stderr,
        )
        return 2

    if args.batch_size < 1:
        print("ERROR: --batch-size must be >= 1", file=sys.stderr)
        return 2

    model_id = args.model.strip()
    notes = (args.notes or "").strip()[:4000]

    out_dir = os.path.abspath(args.out_dir)
    leaderboard_csv = args.leaderboard_csv or os.path.join(out_dir, "leaderboard.csv")
    jobs_csv = args.jobs_csv or os.path.join(out_dir, "jobs_state.csv")

    existing = _read_csv(leaderboard_csv)
    if any(r.get("model_id") == model_id for r in existing):
        print(
            f"ERROR: {model_id!r} already in {leaderboard_csv}; nothing to do.",
            file=sys.stderr,
        )
        return 1

    job_id = (args.job_id or str(uuid.uuid4())[:8]).strip()
    created = _now_iso()
    print(
        f"[job {job_id}] running {model_id} (family={family_id}, device={device_str}, "
        f"batch_size={args.batch_size}, dtype={args.dtype}) using packed dir {packed_dir}",
        file=sys.stderr,
    )
    _upsert_job_local(
        jobs_csv,
        job_id=job_id,
        model_id=model_id,
        family_id=family_id,
        status="running",
        created_at=created,
        submission_notes=notes,
    )

    try:
        result = _run_evaluation_local(
            model_id,
            family_id,
            device_str=device_str,
            device_int=device_int,
            batch_size=args.batch_size,
            dtype_choice=args.dtype,
            progress_cb=_progress_print,
        )
        sys.stderr.write("\n")
    except Exception as e:
        sys.stderr.write("\n")
        traceback.print_exc()
        _upsert_job_local(
            jobs_csv,
            job_id=job_id,
            model_id=model_id,
            family_id=family_id,
            status="failed",
            created_at=created,
            error=str(e),
            submission_notes=notes,
        )
        print(f"[job {job_id}] FAILED: {e}", file=sys.stderr)
        return 1

    try:
        new_row = _append_leaderboard_local(leaderboard_csv, result, notes)
    except Exception as e:
        traceback.print_exc()
        _upsert_job_local(
            jobs_csv,
            job_id=job_id,
            model_id=model_id,
            family_id=family_id,
            status="failed",
            created_at=created,
            error=f"publish failed: {e}",
            submission_notes=notes,
        )
        return 1

    _upsert_job_local(
        jobs_csv,
        job_id=job_id,
        model_id=model_id,
        family_id=family_id,
        status="done",
        created_at=created,
        submission_notes=notes,
    )

    print(
        f"[job {job_id}] done — wrote row to {leaderboard_csv}, job to {jobs_csv}",
        file=sys.stderr,
    )

    if args.json:
        print(json.dumps({"job_id": job_id, "result": result, "row": new_row}, indent=2))
    else:
        print(f"job_id:           {job_id}")
        print(f"model_id:         {result['model_id']}")
        print(f"eval_family:      {result['eval_family']}")
        print(f"device:           {device_str}")
        print(f"dtype:            {args.dtype}")
        print(f"batch_size:       {args.batch_size}")
        print(f"wer_clean (→ anechoic CSV):        {result.get('wer_clean')}")
        print(f"wer_noisy (→ lab measured CSV):     {result.get('wer_noisy')}")
        print(f"wer_reverberant (→ lab sim CSV):   {result.get('wer_reverberant')}")
        print(f"wer_high (→ High SNR CSV):          {result.get('wer_high')}")
        print(f"wer_moving (→ Moving Sources CSV):  {result.get('wer_moving')}")
        print(f"wer_mid (→ Mid SNR CSV):            {result.get('wer_mid')}")
        print(f"wer_low (→ Low SNR CSV):            {result.get('wer_low')}")
        print(f"num_samples:      {result.get('num_samples')}")
        print(f"eval_wall_time_s: {result.get('eval_wall_time_s')}")
        print(f"eval_rtf:         {result.get('eval_rtf')}")
        print(f"num_params:       {result.get('num_params')}")
    return 0


if __name__ == "__main__":
    sys.exit(main())