utils.py

"""Vendored metrics + manifest utilities for standalone explorer.

Copied from asr_benchmark.utils.manifest, asr_benchmark.utils.metrics, and
asr_benchmark.utils.data so the explorer can run without the parent
benchmark repo installed.
"""

from __future__ import annotations

import json
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np


# ── Manifest (from asr_benchmark.utils.manifest) ──────────────────────────────

def read_manifest(path: str | Path) -> list[dict[str, Any]]:
    """Read a JSONL manifest file and return a list of dicts."""
    records = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


# ── Metrics (from asr_benchmark.utils.metrics) ──────────────────────────────

@dataclass
class EditStats:
    hits: int = 0
    substitutions: int = 0
    insertions: int = 0
    deletions: int = 0

    @property
    def errors(self) -> int:
        return self.substitutions + self.insertions + self.deletions

    @property
    def ref_length(self) -> int:
        return self.hits + self.substitutions + self.deletions

    def __add__(self, other: "EditStats") -> "EditStats":
        return EditStats(
            hits=self.hits + other.hits,
            substitutions=self.substitutions + other.substitutions,
            insertions=self.insertions + other.insertions,
            deletions=self.deletions + other.deletions,
        )


def _align(ref: list[str], hyp: list[str]) -> EditStats:
    r, h = len(ref), len(hyp)
    dp = [[0] * (h + 1) for _ in range(r + 1)]
    for i in range(r + 1):
        dp[i][0] = i
    for j in range(h + 1):
        dp[0][j] = j
    for i in range(1, r + 1):
        for j in range(1, h + 1):
            if ref[i - 1] == hyp[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
    stats = EditStats()
    i, j = r, h
    while i > 0 or j > 0:
        if i > 0 and j > 0 and ref[i - 1] == hyp[j - 1]:
            stats.hits += 1
            i -= 1; j -= 1
        elif i > 0 and j > 0 and dp[i][j] == dp[i - 1][j - 1] + 1:
            stats.substitutions += 1
            i -= 1; j -= 1
        elif j > 0 and dp[i][j] == dp[i][j - 1] + 1:
            stats.insertions += 1
            j -= 1
        else:
            stats.deletions += 1
            i -= 1
    return stats


def _align_words(ref: list[str], hyp: list[str]) -> list[tuple[str, str | None, str | None]]:
    """Word-level alignment; returns (op, ref_word, hyp_word) tuples."""
    r, h = len(ref), len(hyp)
    dp = [[0] * (h + 1) for _ in range(r + 1)]
    for i in range(r + 1):
        dp[i][0] = i
    for j in range(h + 1):
        dp[0][j] = j
    for i in range(1, r + 1):
        for j in range(1, h + 1):
            if ref[i - 1] == hyp[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
    ops: list[tuple[str, str | None, str | None]] = []
    i, j = r, h
    while i > 0 or j > 0:
        if i > 0 and j > 0 and ref[i - 1] == hyp[j - 1]:
            ops.append(("hit", ref[i - 1], hyp[j - 1])); i -= 1; j -= 1
        elif i > 0 and j > 0 and dp[i][j] == dp[i - 1][j - 1] + 1:
            ops.append(("sub", ref[i - 1], hyp[j - 1])); i -= 1; j -= 1
        elif j > 0 and dp[i][j] == dp[i][j - 1] + 1:
            ops.append(("ins", None, hyp[j - 1])); j -= 1
        else:
            ops.append(("del", ref[i - 1], None)); i -= 1
    ops.reverse()
    return ops


def _corpus_word_stats(references: list[str], hypotheses: list[str]) -> EditStats:
    total = EditStats()
    for ref, hyp in zip(references, hypotheses):
        total = total + _align(ref.split(), hyp.split())
    return total


def _corpus_char_stats(references: list[str], hypotheses: list[str]) -> EditStats:
    total = EditStats()
    for ref, hyp in zip(references, hypotheses):
        total = total + _align(list(ref.replace(" ", "")), list(hyp.replace(" ", "")))
    return total


def compute_wer(references: list[str], hypotheses: list[str]) -> float:
    stats = _corpus_word_stats(references, hypotheses)
    if stats.ref_length == 0:
        return 0.0
    return round(100.0 * stats.errors / stats.ref_length, 2)


def compute_cer(references: list[str], hypotheses: list[str]) -> float:
    stats = _corpus_char_stats(references, hypotheses)
    if stats.ref_length == 0:
        return 0.0
    return round(100.0 * stats.errors / stats.ref_length, 2)


# ── Rare-word WER helpers ────────────────────────────────────────────────────

def build_freq_map(references: list[str]) -> Counter:
    freq: Counter = Counter()
    for ref in references:
        freq.update(ref.split())
    return freq


def make_common_words(freq_map: Counter, top_n: int) -> frozenset:
    return frozenset(w for w, _ in freq_map.most_common(top_n))


def compute_rare_wer(
    refs: list[str],
    hyps: list[str],
    common_words: frozenset,
) -> dict:
    rare_ref = rare_hits = rare_subs = rare_dels = 0
    for ref, hyp in zip(refs, hyps):
        for op, rw, _hw in _align_words(ref.split(), hyp.split()):
            if rw is None or rw in common_words:
                continue
            rare_ref += 1
            if op == "hit":
                rare_hits += 1
            elif op == "sub":
                rare_subs += 1
            elif op == "del":
                rare_dels += 1

    def pct(n: int, d: int) -> float:
        return round(100.0 * n / d, 2) if d > 0 else 0.0

    return {
        "rare_wer":           pct(rare_subs + rare_dels, rare_ref),
        "rare_sub_rate":      pct(rare_subs,             rare_ref),
        "rare_del_rate":      pct(rare_dels,             rare_ref),
        "rare_ref_words":     rare_ref,
        "rare_substitutions": rare_subs,
        "rare_deletions":     rare_dels,
    }


# ── Audio decoding (from asr_benchmark.utils.data) ────────────────────────────

def decode_audio(audio_data, target_sr: int) -> tuple[np.ndarray, int]:
    """
    Decode an audio field from a HuggingFace dataset row.

    Handles two formats:
    - Standard dict: {"array": np.ndarray, "sampling_rate": int}
    - torchcodec AudioDecoder: used by newer HF datasets (e.g. Revolab/ASR-Benchmark-Public)
    """
    if isinstance(audio_data, dict):
        array = audio_data["array"].astype(np.float32)
        sr = audio_data["sampling_rate"]
        if sr != target_sr:
            array = _resample(array, sr, target_sr)
        return array, target_sr

    # torchcodec AudioDecoder
    samples = audio_data.get_all_samples()
    data = samples.data
    sr = int(samples.sample_rate)
    try:
        array = data.numpy()
    except Exception:
        array = data.cpu().numpy()
    if array.ndim == 2:
        array = array.mean(axis=0)
    array = array.astype(np.float32)
    if sr != target_sr:
        array = _resample(array, sr, target_sr)
    return array, target_sr


def _resample(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    try:
        import resampy
        return resampy.resample(audio, orig_sr, target_sr)
    except ImportError:
        pass
    try:
        import librosa
        return librosa.resample(audio, orig_sr=orig_sr, target_sr=target_sr)
    except ImportError:
        pass
    n = int(len(audio) * target_sr / orig_sr)
    return np.interp(np.linspace(0, len(audio) - 1, n), np.arange(len(audio)), audio).astype(np.float32)