benchmark/dataset.py

"""
Private packed test tensors (clean / noisy / reverberant) and WER aggregation.

Inference is delegated to a `transcribe(audio_np, sampling_rate) -> str` callable
provided by each model-family backend.
"""

import os

from storage import download_bucket_file

# packed/clean.pt, etc.: {"sample_id": {"waveform": Tensor, "sample_rate": int, "transcript": str}, ...}
PACKED_FILES = {
    "clean": "packed/clean.pt",
    "measured": "packed/measured.pt",
    "sim": "packed/sim.pt",
    "low": "packed/low.pt",
    "mid": "packed/mid.pt",
    "high": "packed/high.pt",
    "moving_low": "packed/moving_low.pt",
    "moving_mid": "packed/moving_mid.pt",
    "moving_high": "packed/moving_high.pt",
}

# Bucket path key -> ``run_evaluation`` result field (``wer_<suffix>``).
# Mapped to canonical CSV columns in init.leaderboard_row_from_eval_result.
CONDITION_WER_SUFFIX: dict[str, str] = {
    "clean": "clean",
    "measured": "measured",
    "sim": "sim",
    "low": "low",
    "mid": "mid",
    "high": "high",
    "moving_low": "moving_low",
    "moving_mid": "moving_mid",
    "moving_high": "moving_high",
}


def wer_result_key(condition_key: str) -> str:
    """Result dict key for WER on a packed condition (e.g. ``wer_high_snr``)."""
    suffix = CONDITION_WER_SUFFIX.get(condition_key, condition_key)
    return f"wer_{suffix}"


# Gradio checkbox labels (display, condition_key).
CONDITION_UI_CHOICES: tuple[tuple[str, str], ...] = (
    ("Near Field Speech", "clean"),
    ("Lab Measured", "measured"),
    ("Lab Simulated", "sim"),
    ("High SNR", "high"),
    ("Mid SNR", "mid"),
    ("Low SNR", "low"),
    ("Moving Low SNR", "moving_low"),
    ("Moving Mid SNR", "moving_mid"),
    ("Moving High SNR", "moving_high"),
)

DEFAULT_CONDITION_KEYS: tuple[str, ...] = tuple(PACKED_FILES.keys())


def resolve_condition_keys(keys: list[str] | None) -> tuple[str, ...]:
    """Return validated condition keys; ``None`` or empty means all packed splits."""
    if not keys:
        return DEFAULT_CONDITION_KEYS
    valid = [k for k in keys if k in PACKED_FILES]
    return tuple(valid) if valid else DEFAULT_CONDITION_KEYS


def is_full_condition_set(keys: list[str] | None) -> bool:
    return set(resolve_condition_keys(keys)) == set(PACKED_FILES.keys())


_dataset_cache: dict[str, dict] = {}
_wer_normalizer = None
_wer_normalizer_init_failed = False


def _normalize_for_wer(text: str) -> str:
    """Normalize text for WER with Whisper English normalizer, fallback to lowercase.

    Applied identically to references and predictions so WER is always measured on
    Whisper-normalized text (lowercased, punctuation stripped, numbers/contractions
    expanded), regardless of which backend or custom script produced the hypothesis.
    """
    global _wer_normalizer
    global _wer_normalizer_init_failed
    if _wer_normalizer is None and not _wer_normalizer_init_failed:
        try:
            from transformers.models.whisper.english_normalizer import EnglishTextNormalizer

            # The British->American spelling map is a minor refinement; fetch it when
            # the tokenizer is available, but don't require a network download just to
            # normalize (an empty map still yields full Whisper normalization).
            spelling_map: dict[str, str] = {}
            try:
                from transformers import WhisperTokenizer

                tok = WhisperTokenizer.from_pretrained("openai/whisper-tiny")
                spelling_map = tok.english_spelling_normalizer
            except Exception:
                spelling_map = {}
            _wer_normalizer = EnglishTextNormalizer(spelling_map)
        except Exception:
            _wer_normalizer_init_failed = True
            _wer_normalizer = None
    s = text or ""
    if _wer_normalizer is not None:
        try:
            return _wer_normalizer(s)
        except Exception:
            pass
    # Last-resort fallback (transformers unavailable): mimic the essentials so refs and
    # predictions stay symmetric -- lowercase, drop punctuation, collapse whitespace.
    import re

    s = re.sub(r"[^\w\s]", " ", s.lower())
    return re.sub(r"\s+", " ", s).strip()


def load_packed_condition(condition_key: str) -> dict:
    """Download and load a packed .pt file. Cached after first load."""
    import torch

    if condition_key in _dataset_cache:
        return _dataset_cache[condition_key]

    pt_path = PACKED_FILES[condition_key]
    local_path = download_bucket_file(pt_path)
    data = torch.load(local_path, weights_only=False)
    os.unlink(local_path)

    _dataset_cache[condition_key] = data
    return data


def evaluate_condition_wer(
    transcribe,
    condition_key: str,
) -> tuple[float, int]:
    """
    Run `transcribe(audio_1d_numpy, sampling_rate)` on every sample and return (wer, n).

    References and predictions are normalized for WER.
    """
    wer, n, _a, _w = evaluate_condition_wer_timed(transcribe, condition_key)
    return wer, n


def accumulate_predictions_for_slice(
    transcribe,
    condition_key: str,
    start: int,
    end: int,
    progress_cb=None,
) -> tuple[list[str], list[str], float, float, int]:
    """
    Run ``transcribe`` on packed samples ``[start, end)`` for ``condition_key``.

    Returns ``(predictions, references, audio_seconds, inference_seconds, n_samples)``.
    ``progress_cb`` is invoked as ``progress_cb(condition_key, done_in_slice, slice_len)`` per sample.
    """
    import time

    data = load_packed_condition(condition_key)
    if not data:
        if progress_cb is not None:
            try:
                progress_cb(condition_key, 0, 0)
            except Exception:
                pass
        return [], [], 0.0, 0.0, 0

    items = list(data.items())
    n = len(items)
    start = max(0, min(start, n))
    end = max(start, min(end, n))
    slice_items = items[start:end]
    slice_len = len(slice_items)

    predictions: list[str] = []
    references: list[str] = []
    audio_seconds = 0.0
    inference_seconds = 0.0

    for j, (_sample_id, sample) in enumerate(slice_items):
        audio_np = sample["waveform"].numpy()
        sr = int(sample.get("sample_rate", 16000))
        audio_seconds += float(len(audio_np)) / float(sr)
        t0 = time.perf_counter()
        text = transcribe(audio_np, sr)
        inference_seconds += time.perf_counter() - t0
        predictions.append(_normalize_for_wer(str(text).strip().lower()))
        references.append(_normalize_for_wer(str(sample["transcript"]).strip().lower()))
        if progress_cb is not None:
            try:
                progress_cb(condition_key, j + 1, slice_len)
            except Exception:
                pass

    return predictions, references, audio_seconds, inference_seconds, slice_len


def evaluate_condition_wer_timed(
    transcribe,
    condition_key: str,
    progress_cb=None,
) -> tuple[float, int, float, float]:
    """
    Same as WER evaluation but also returns (audio_seconds, inference_wall_seconds) for RTF.

    If `progress_cb` is provided it is invoked as
    `progress_cb(condition_key, samples_done_in_condition, samples_total_in_condition)`
    once per sample (after inference completes) so callers can expose progress.
    """
    from jiwer import wer as compute_wer

    data = load_packed_condition(condition_key)
    n = len(data)
    preds, refs, audio_seconds, inference_seconds, _cnt = accumulate_predictions_for_slice(
        transcribe, condition_key, 0, n, progress_cb=progress_cb
    )
    if not preds:
        return 0.0, 0, 0.0, 0.0

    wer = round(compute_wer(refs, preds), 4)
    return wer, len(refs), audio_seconds, inference_seconds