Add assembly eval: dataset WER/DER scoring for model selection (#74)

Transcribe an evaluation dataset — a Hugging Face dataset id (via the
datasets-server REST API, no datasets dependency) or a local .csv/.jsonl
manifest — and score word error rate against its reference texts with
jiwer. With --speaker-labels the run also diarizes and scores diarization
error rate against reference speaker turns via pyannote.metrics
(--collar for boundary forgiveness). Per-file rows plus pooled corpus
scores render as a table or --json.

https://claude.ai/code/session_014o9JsqBmqmybikbtHNsbSf

Co-authored-by: Claude <noreply@anthropic.com>

Author: alexkroman

.importlinter

@@ -14,8 +14,10 @@ source_modules =
     aai_cli.config
     aai_cli.config_builder
     aai_cli.context
+    aai_cli.der
     aai_cli.environments
     aai_cli.errors
+    aai_cli.eval_data
     aai_cli.follow
     aai_cli.help_panels
     aai_cli.help_text
@@ -32,6 +34,7 @@ source_modules =
     aai_cli.transcribe_batch
     aai_cli.transcribe_exec
     aai_cli.transcribe_render
+    aai_cli.wer
     aai_cli.ws
     aai_cli.youtube
 forbidden_modules =
@@ -47,6 +50,7 @@ modules =
     aai_cli.commands.deploy
     aai_cli.commands.dev
     aai_cli.commands.doctor
+    aai_cli.commands.evaluate
     aai_cli.commands.init
     aai_cli.commands.keys
     aai_cli.commands.llm
@@ -66,9 +70,12 @@ source_modules =
     aai_cli.client
     aai_cli.config
     aai_cli.config_builder
+    aai_cli.der
     aai_cli.environments
     aai_cli.errors
+    aai_cli.eval_data
     aai_cli.llm
     aai_cli.telemetry
+    aai_cli.wer
 forbidden_modules =
     rich

README.md

@@ -50,6 +50,7 @@ Requires Python 3.12+. On Linux, install PortAudio once for microphone support (
 - **Real-time streaming**: `assembly stream` transcribes the microphone, a file, or a URL live — on macOS it can capture system audio too.
 - **Voice agent**: `assembly agent` runs a full-duplex spoken conversation in your terminal (use headphones).
 - **LLM Gateway**: `assembly llm` prompts an LLM over a transcript, stdin, or a live stream (`assembly stream --llm "summarize as I talk"`).
+- **Model evaluation**: `assembly eval` transcribes a Hugging Face dataset or a local `.csv`/`.jsonl` manifest and scores WER against its references (plus DER with `--speaker-labels`) — handy for picking a speech model.
 - **Starter apps**: `assembly init` scaffolds a self-contained FastAPI + HTML app (`audio-transcription`, `live-captions`, `voice-agent`).
 - **Code generation**: add `--show-code` to `transcribe`/`stream`/`agent` to print the equivalent Python SDK script instead of running.
 - **Account self-service**: `assembly keys` / `balance` / `usage` / `limits` / `sessions` / `audit` via browser login.

aai_cli/commands/evaluate.py

@@ -0,0 +1,212 @@
+"""`assembly eval` — transcribe an evaluation dataset and score it against references.
+
+WER (via jiwer) against the dataset's reference texts; with ``--speaker-labels``
+also DER (via pyannote.metrics) against its reference speaker turns. The module
+is named ``evaluate`` because importing a module named ``eval`` would shadow the
+builtin; the command itself registers as ``eval``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import assemblyai as aai
+import typer
+from rich.console import RenderableType
+
+from aai_cli import client, config, der, eval_data, help_panels, jsonshape, options, output, wer
+from aai_cli.context import AppState, run_command
+from aai_cli.help_text import examples_epilog
+
+app = typer.Typer()
+
+
+def _pct(value: object) -> str:
+    return f"{jsonshape.as_float(value):.2%}"
+
+
+def _hypothesis_turns(transcript: aai.Transcript) -> list[der.Turn]:
+    """The transcript's diarized utterances as DER hypothesis turns (ms → seconds)."""
+    return [
+        der.Turn(
+            speaker=str(getattr(utterance, "speaker", "")),
+            start=jsonshape.as_float(getattr(utterance, "start", None)) / 1000,
+            end=jsonshape.as_float(getattr(utterance, "end", None)) / 1000,
+        )
+        for utterance in jsonshape.object_list(getattr(transcript, "utterances", None))
+    ]
+
+
+@dataclass(frozen=True)
+class _ItemResult:
+    """One scored row: the emitted dict plus the scores kept for pooling."""
+
+    row: dict[str, object]
+    words: wer.Score | None
+    speakers: der.DerScore | None
+
+
+def _score_item(
+    item: eval_data.EvalItem, transcript: aai.Transcript, *, collar: float
+) -> _ItemResult:
+    row: dict[str, object] = {"item": item.item_id}
+    words: wer.Score | None = None
+    speakers: der.DerScore | None = None
+    if item.reference is not None:
+        words = wer.score(item.reference, str(transcript.text or ""))
+        row.update({"words": words.words, "errors": words.errors, "wer": words.wer})
+    if item.turns is not None:
+        speakers = der.score(item.turns, _hypothesis_turns(transcript), collar=collar)
+        row["der"] = speakers.der
+    return _ItemResult(row=row, words=words, speakers=speakers)
+
+
+def _payload(
+    label: str, speech_model: aai.SpeechModel | None, results: list[_ItemResult]
+) -> dict[str, object]:
+    payload: dict[str, object] = {
+        "dataset": label,
+        "speech_model": speech_model.value if speech_model else None,
+        "items": len(results),
+        "rows": [result.row for result in results],
+    }
+    word_scores = [result.words for result in results if result.words is not None]
+    if word_scores:
+        total = wer.pooled(word_scores)
+        payload.update({"words": total.words, "errors": total.errors, "wer": total.wer})
+    der_scores = [result.speakers for result in results if result.speakers is not None]
+    if der_scores:
+        payload["der"] = der.pooled(der_scores).der
+    return payload
+
+
+def _summary(payload: dict[str, object]) -> str:
+    parts: list[str] = []
+    if "wer" in payload:
+        errors = jsonshape.as_int(payload.get("errors"))
+        noun = "error" if errors == 1 else "errors"
+        parts.append(
+            f"WER {_pct(payload.get('wer'))} ({errors} {noun} / {payload.get('words')} words)"
+        )
+    if "der" in payload:
+        parts.append(f"DER {_pct(payload.get('der'))}")
+    return output.heading("   ".join(parts))
+
+
+def _render(payload: dict[str, object]) -> RenderableType:
+    has_wer = "wer" in payload
+    has_der = "der" in payload
+    columns = [
+        "ITEM",
+        *(["WORDS", "ERRORS", "WER"] if has_wer else []),
+        *(["DER"] if has_der else []),
+    ]
+    table = output.data_table(*columns)
+    for row in jsonshape.mapping_list(payload.get("rows")):
+        cells = [str(row.get("item"))]
+        if has_wer:
+            cells += [str(row.get("words")), str(row.get("errors")), _pct(row.get("wer"))]
+        if has_der:
+            cells.append(_pct(row.get("der")))
+        table.add_row(*cells)
+    model = payload.get("speech_model") or "default model"
+    return output.stack(
+        output.muted(f"{payload.get('dataset')} · {model}"), table, _summary(payload)
+    )
+
+
+@app.command(
+    name="eval",
+    rich_help_panel=help_panels.TRANSCRIPTION,
+    epilog=examples_epilog(
+        [
+            ("Score a model on 10 rows of an HF dataset", "assembly eval distil-whisper/meanwhile"),
+            (
+                "Compare models on your own audio",
+                "assembly eval calls.csv --speech-model universal",
+            ),
+            (
+                "Score diarization too (WER + DER)",
+                "assembly eval agent-calls.jsonl --speaker-labels",
+            ),
+            (
+                "Pick a subset/split and more rows",
+                "assembly eval mozilla-foundation/common_voice_17_0 --subset en --limit 50",
+            ),
+        ]
+    ),
+)
+def evaluate(
+    ctx: typer.Context,
+    dataset: str = typer.Argument(
+        ...,
+        help="Hugging Face dataset id, or a local .csv/.jsonl manifest with audio + text columns.",
+    ),
+    split: str | None = typer.Option(
+        None, "--split", help="Hugging Face split to score (default: test)."
+    ),
+    subset: str | None = typer.Option(
+        None, "--subset", help="Hugging Face config/subset name (e.g. a language)."
+    ),
+    limit: int = typer.Option(10, "--limit", min=1, max=100, help="Rows to evaluate (1-100)."),
+    audio_column: str | None = typer.Option(
+        None, "--audio-column", help="Audio column name (default: auto-detect)."
+    ),
+    text_column: str | None = typer.Option(
+        None, "--text-column", help="Reference text column name (default: auto-detect)."
+    ),
+    speech_model: aai.SpeechModel | None = typer.Option(
+        None, "--speech-model", help="Speech model to evaluate."
+    ),
+    language_code: str | None = typer.Option(
+        None, "--language-code", help="Force a language (e.g. en_us)."
+    ),
+    speaker_labels: bool = typer.Option(
+        False,
+        "--speaker-labels",
+        help="Diarize and also score DER against the dataset's reference speaker turns (speakers/timestamps_start/timestamps_end columns, in seconds).",
+    ),
+    collar: float = typer.Option(
+        1.0,
+        "--collar",
+        min=0.0,
+        help="DER forgiveness (seconds) around each reference turn boundary.",
+    ),
+    json_out: bool = options.json_option("Output the rows and summary as one JSON object."),
+) -> None:
+    """Transcribe an evaluation dataset and score WER against its reference texts.
+
+    Handy for picking a model: run once per --speech-model and compare. Datasets
+    come from the Hugging Face Hub (gated ones need HF_TOKEN) or a local .csv/.jsonl
+    manifest. --speaker-labels also scores diarization (DER) against reference
+    speaker turns.
+    """
+
+    def body(state: AppState, json_mode: bool) -> None:
+        data = eval_data.load(
+            dataset,
+            split=split,
+            subset=subset,
+            audio_column=audio_column,
+            text_column=text_column,
+            limit=limit,
+            with_speakers=speaker_labels,
+        )
+        api_key = config.resolve_api_key(profile=state.profile)
+        transcription_config = aai.TranscriptionConfig(
+            speech_model=speech_model,
+            language_code=language_code,
+            speaker_labels=speaker_labels or None,
+        )
+        results: list[_ItemResult] = []
+        for index, item in enumerate(data.items, start=1):
+            with output.status(
+                f"[{index}/{len(data.items)}] Transcribing {item.item_id}…",
+                json_mode=json_mode,
+                quiet=state.quiet,
+            ):
+                transcript = client.transcribe(api_key, item.audio, config=transcription_config)
+            results.append(_score_item(item, transcript, collar=collar))
+        output.emit(_payload(data.label, speech_model, results), _render, json_mode=json_mode)
+
+    run_command(ctx, body, json=json_out)

aai_cli/der.py

@@ -0,0 +1,95 @@
+"""Diarization error rate (DER) scoring for `assembly eval --speaker-labels`.
+
+A thin shim over :mod:`pyannote.metrics` — the de-facto standard DER
+implementation, including its optimal reference↔hypothesis speaker mapping —
+so the alignment math is never re-derived here. pyannote drags numpy/scipy/
+pandas along, so it is imported lazily inside :func:`score`; the dataclasses
+here stay import-cheap for the command layer. No SDK, no Rich.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import TypeAdapter
+
+if TYPE_CHECKING:
+    from pyannote.core import Annotation
+
+# pyannote types the metric call as a bare float; with ``detailed=True`` it
+# actually returns the components dict — validate that shape instead of casting.
+_COMPONENTS: TypeAdapter[dict[str, float]] = TypeAdapter(dict[str, float])
+
+
+@dataclass(frozen=True)
+class Turn:
+    """One speaker turn: who spoke from ``start`` to ``end`` (seconds)."""
+
+    speaker: str
+    start: float
+    end: float
+
+
+@dataclass(frozen=True)
+class DerScore:
+    """DER components in seconds of speech; pooled across files for corpus DER."""
+
+    missed: float
+    false_alarm: float
+    confusion: float
+    total: float
+
+    @property
+    def der(self) -> float:
+        return (self.missed + self.false_alarm + self.confusion) / self.total
+
+
+def _annotation(turns: list[Turn]) -> Annotation:
+    from pyannote.core import Annotation, Segment
+
+    annotation = Annotation()
+    # The list index keys each track, so two overlapping turns by the same
+    # speaker don't collapse into one.
+    for track, turn in enumerate(turns):
+        annotation[Segment(turn.start, turn.end), track] = turn.speaker
+    return annotation
+
+
+def score(reference: list[Turn], hypothesis: list[Turn], *, collar: float = 0.0) -> DerScore:
+    """Score hypothesis speaker turns against the reference diarization.
+
+    The caller guarantees a non-empty reference (the dataset loader rejects
+    rows without speaker turns), so ``DerScore.der`` is always well-defined.
+    ``collar`` forgives that many seconds around each reference turn boundary.
+    """
+    from pyannote.core import Segment, Timeline
+    from pyannote.metrics.diarization import DiarizationErrorRate
+
+    metric = DiarizationErrorRate(collar=collar)
+    # An explicit evaluation extent (audio start through the last turn either
+    # side heard); without it pyannote warns about approximating the UEM.
+    extent = max(turn.end for turn in [*reference, *hypothesis])
+    components: object = metric(
+        _annotation(reference),
+        _annotation(hypothesis),
+        uem=Timeline([Segment(0.0, extent)]),
+        detailed=True,
+    )
+    detail = _COMPONENTS.validate_python(components)
+    return DerScore(
+        missed=detail["missed detection"],
+        false_alarm=detail["false alarm"],
+        confusion=detail["confusion"],
+        total=detail["total"],
+    )
+
+
+def pooled(scores: list[DerScore]) -> DerScore:
+    """Corpus-level score: error seconds over total reference speech seconds."""
+    return DerScore(
+        missed=sum(item.missed for item in scores),
+        false_alarm=sum(item.false_alarm for item in scores),
+        confusion=sum(item.confusion for item in scores),
+        total=sum(item.total for item in scores),
+    )