mARTin-B78
4d59853da6
- BenchRow gains `best_for` field: "Short clips" / "Short / medium" / "Long / batch" / "Streaming" — derived from engine type and model name - Device now shows "CUDA" instead of "GPU" for clarity - Benchmark table gains a "Best for" column between Device and Time(s) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
124 lines
4.0 KiB
Python
124 lines
4.0 KiB
Python
"""Benchmark STT engines against a reference clip.
|
||
|
||
Given a WAV and its reference transcript, transcribe with each engine, measure
|
||
the time, and score accuracy as 1 − WER (word error rate). Used by the Settings
|
||
Benchmark tab to find the fastest and most accurate engine/model.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
from . import stt
|
||
from .routing import normalize
|
||
|
||
|
||
@dataclass
|
||
class BenchRow:
|
||
engine: str
|
||
model: str
|
||
device: str # "CPU" | "CUDA" | "remote"
|
||
best_for: str # "Short clips" | "Short / medium" | "Long / batch" | "Streaming"
|
||
ok: bool
|
||
seconds: float
|
||
wer: float
|
||
accuracy: float # percent, max(0, 1-wer)*100
|
||
text: str
|
||
error: str = ""
|
||
|
||
|
||
def _tokens(text: str, case_sensitive: bool) -> list[str]:
|
||
"""Word tokens. case_sensitive keeps capitalisation (and accents); both
|
||
drop punctuation so only the words/spelling are compared."""
|
||
if not case_sensitive:
|
||
return normalize(text)
|
||
return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE).split()
|
||
|
||
|
||
def _edit_distance(a: list[str], b: list[str]) -> int:
|
||
m, n = len(a), len(b)
|
||
dp = list(range(n + 1))
|
||
for i in range(1, m + 1):
|
||
prev = dp[0]
|
||
dp[0] = i
|
||
for j in range(1, n + 1):
|
||
cur = dp[j]
|
||
dp[j] = min(dp[j] + 1, dp[j - 1] + 1, prev + (a[i - 1] != b[j - 1]))
|
||
prev = cur
|
||
return dp[n]
|
||
|
||
|
||
def wer(reference: str, hypothesis: str, *, case_sensitive: bool = False) -> float:
|
||
"""Word error rate (0 = perfect). case_sensitive makes capitalisation count."""
|
||
ref = _tokens(reference, case_sensitive)
|
||
hyp = _tokens(hypothesis, case_sensitive)
|
||
if not ref:
|
||
return 0.0 if not hyp else 1.0
|
||
return _edit_distance(ref, hyp) / len(ref)
|
||
|
||
|
||
def _engine_device(engine, transcriber) -> str:
|
||
if engine.is_local:
|
||
return "CUDA" if getattr(transcriber, "device", "cpu") == "cuda" else "CPU"
|
||
return "remote"
|
||
|
||
|
||
def _engine_best_for(engine) -> str:
|
||
if engine.type == "riva_realtime":
|
||
return "Streaming"
|
||
model = (engine.model or "").lower()
|
||
name = engine.name.lower()
|
||
if any(x in name for x in ("stream", "realtime", "real-time", "live")):
|
||
return "Streaming"
|
||
if engine.is_local:
|
||
if any(x in model for x in ("tiny", "base")):
|
||
return "Short clips"
|
||
if any(x in model for x in ("large",)):
|
||
return "Long / batch"
|
||
return "Short / medium"
|
||
# Remote endpoint
|
||
if any(x in name for x in ("large", "batch")):
|
||
return "Long / batch"
|
||
return "Short / medium"
|
||
|
||
|
||
def run(engines, wav_path: Path, reference: str, *, language: str = "",
|
||
case_sensitive: bool = True, get_local_transcriber=None, progress=None) -> list[BenchRow]:
|
||
"""Benchmark each engine; calls progress(row) as each finishes.
|
||
|
||
Accuracy is case-sensitive by default so wrong capitalisation counts.
|
||
"""
|
||
rows: list[BenchRow] = []
|
||
for e in engines:
|
||
tr = get_local_transcriber(e) if (e.is_local and get_local_transcriber) else None
|
||
res = stt.benchmark(e, wav_path, language=language, local_transcriber=tr)
|
||
w = wer(reference, res.text, case_sensitive=case_sensitive) if res.ok else 1.0
|
||
row = BenchRow(
|
||
engine=e.name,
|
||
model=e.model or ("local" if e.is_local else "(default)"),
|
||
device=_engine_device(e, tr),
|
||
best_for=_engine_best_for(e),
|
||
ok=res.ok,
|
||
seconds=res.seconds,
|
||
wer=w,
|
||
accuracy=max(0.0, 1.0 - w) * 100.0,
|
||
text=res.text,
|
||
error=res.error,
|
||
)
|
||
rows.append(row)
|
||
if progress:
|
||
progress(row)
|
||
return rows
|
||
|
||
|
||
def best(rows: list[BenchRow]) -> tuple[BenchRow | None, BenchRow | None]:
|
||
"""Return (fastest, most_accurate) among successful rows."""
|
||
ok = [r for r in rows if r.ok]
|
||
if not ok:
|
||
return None, None
|
||
fastest = min(ok, key=lambda r: r.seconds)
|
||
most_accurate = max(ok, key=lambda r: (r.accuracy, -r.seconds))
|
||
return fastest, most_accurate
|