Update speech model options to universal-3-pro and universal-2

aai_cli/commands/evaluate.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from enum import StrEnum
 
 import assemblyai as aai
 import typer
@@ -21,6 +22,14 @@
 app = typer.Typer()
 
 
+class EvalSpeechModel(StrEnum):
+    """The current-generation models, requested via the SDK's ``speech_models``
+    list parameter (its legacy ``SpeechModel`` enum predates them)."""
+
+    universal_3_pro = "universal-3-pro"
+    universal_2 = "universal-2"
+
+
 def _pct(value: object) -> str:
     return f"{jsonshape.as_float(value):.2%}"
 
@@ -62,7 +71,7 @@ def _score_item(
 
 
 def _payload(
-    label: str, speech_model: aai.SpeechModel | None, results: list[_ItemResult]
+    label: str, speech_model: EvalSpeechModel | None, results: list[_ItemResult]
 ) -> dict[str, object]:
     payload: dict[str, object] = {
         "dataset": label,
@@ -120,10 +129,13 @@ def _render(payload: dict[str, object]) -> RenderableType:
     rich_help_panel=help_panels.TRANSCRIPTION,
     epilog=examples_epilog(
         [
-            ("Score a model on 10 rows of an HF dataset", "assembly eval distil-whisper/meanwhile"),
+            (
+                "Score a model on 10 rows of an HF dataset",
+                "assembly eval sanchit-gandhi/tedlium-data",
+            ),
             (
                 "Compare models on your own audio",
-                "assembly eval calls.csv --speech-model universal",
+                "assembly eval calls.csv --speech-model universal-3-pro",
             ),
             (
                 "Score diarization too (WER + DER)",
@@ -135,11 +147,11 @@ def _render(payload: dict[str, object]) -> RenderableType:
             ),
             (
                 "Evaluate non-English audio",
-                "assembly eval PolyAI/minds14 --subset fr-FR --language-code fr",
+                "assembly eval fixie-ai/common_voice_17_0 --subset fr --language-code fr",
             ),
             (
                 "DER on a Hugging Face diarization set",
-                "assembly eval diarizers-community/simsamu --speaker-labels",
+                "assembly eval talkbank/callhome --subset eng --speaker-labels",
             ),
         ]
     ),
@@ -163,7 +175,7 @@ def evaluate(
     text_column: str | None = typer.Option(
         None, "--text-column", help="Reference text column name (default: auto-detect)."
     ),
-    speech_model: aai.SpeechModel | None = typer.Option(
+    speech_model: EvalSpeechModel | None = typer.Option(
         None, "--speech-model", help="Speech model to evaluate."
     ),
     language_code: str | None = typer.Option(
@@ -195,10 +207,19 @@ def evaluate(
     serves with audio + reference columns; gated ones need HF_TOKEN) or a local
     .csv/.jsonl manifest with audio + text columns. Hub sets to try:
     openslr/librispeech_asr (read English; subsets clean/other),
-    MLCommons/peoples_speech (real-world US English; subset clean),
-    distil-whisper/meanwhile (long-form English), PolyAI/minds14 (banking
-    calls in 14 locales; subsets like fr-FR), and diarizers-community/simsamu
-    (French dispatch calls with speaker turns, for --speaker-labels).
+    sanchit-gandhi/tedlium-data (TED talks),
+    sanchit-gandhi/earnings22_robust_split (earnings calls),
+    kensho/spgispeech (financial calls; subset test),
+    edinburghcstr/ami (meetings; subsets ihm/sdm),
+    fixie-ai/gigaspeech (--subset dev --split dev),
+    fixie-ai/peoples_speech (real-world US English; subset clean),
+    fixie-ai/common_voice_17_0 (99 locales; subsets like en/fr),
+    facebook/voxpopuli (parliament speech; subset en),
+    hhoangphuoc/switchboard (phone calls; --split validation),
+    ylacombe/expresso (expressive speech),
+    speechbrain/LoquaciousSet (--subset small --audio-column wav), and
+    talkbank/callhome (phone calls with speaker turns; --subset eng, for
+    --speaker-labels).
     """
 
     def body(state: AppState, json_mode: bool) -> None:
@@ -213,7 +234,7 @@ def body(state: AppState, json_mode: bool) -> None:
         )
         api_key = config.resolve_api_key(profile=state.profile)
         transcription_config = aai.TranscriptionConfig(
-            speech_model=speech_model,
+            speech_models=[speech_model.value] if speech_model else None,
             language_code=language_code,
             speaker_labels=speaker_labels or None,
         )

tests/__snapshots__/test_cli_output_snapshots.ambr

@@ -253,64 +253,71 @@
    serves with audio + reference columns; gated ones need HF_TOKEN) or a local
    .csv/.jsonl manifest with audio + text columns. Hub sets to try:
    openslr/librispeech_asr (read English; subsets clean/other),
-   MLCommons/peoples_speech (real-world US English; subset clean),
-   distil-whisper/meanwhile (long-form English), PolyAI/minds14 (banking
-   calls in 14 locales; subsets like fr-FR), and diarizers-community/simsamu
-   (French dispatch calls with speaker turns, for --speaker-labels).
+   sanchit-gandhi/tedlium-data (TED talks),
+   sanchit-gandhi/earnings22_robust_split (earnings calls),
+   kensho/spgispeech (financial calls; subset test),
+   edinburghcstr/ami (meetings; subsets ihm/sdm),
+   fixie-ai/gigaspeech (--subset dev --split dev),
+   fixie-ai/peoples_speech (real-world US English; subset clean),
+   fixie-ai/common_voice_17_0 (99 locales; subsets like en/fr),
+   facebook/voxpopuli (parliament speech; subset en),
+   hhoangphuoc/switchboard (phone calls; --split validation),
+   ylacombe/expresso (expressive speech),
+   speechbrain/LoquaciousSet (--subset small --audio-column wav), and
+   talkbank/callhome (phone calls with speaker turns; --subset eng, for
+   --speaker-labels).
 
   ╭─ Arguments ──────────────────────────────────────────────────────────────────╮
   │ *    dataset      TEXT  Hugging Face dataset id, or a local .csv/.jsonl      │
   │                         manifest with audio + text columns.                  │
   │                         [required]                                           │
   ╰──────────────────────────────────────────────────────────────────────────────╯
   ╭─ Options ────────────────────────────────────────────────────────────────────╮
-  │ --split                   TEXT                      Hugging Face split to    │
-  │                                                     score (default: test).   │
-  │ --subset                  TEXT                      Hugging Face             │
-  │                                                     config/subset name (e.g. │
-  │                                                     a language).             │
-  │ --limit                   INTEGER RANGE             Rows to evaluate         │
-  │                           [1<=x<=100]               (1-100).                 │
-  │                                                     [default: 10]            │
-  │ --audio-column            TEXT                      Audio column name        │
-  │                                                     (default: auto-detect).  │
-  │ --text-column             TEXT                      Reference text column    │
-  │                                                     name (default:           │
-  │                                                     auto-detect).            │
-  │ --speech-model            [best|nano|slam-1|univer  Speech model to          │
-  │                           sal]                      evaluate.                │
-  │ --language-code           TEXT                      Force a language (e.g.   │
-  │                                                     en_us).                  │
-  │ --speaker-labels                                    Diarize and also score   │
-  │                                                     DER against the          │
-  │                                                     dataset's reference      │
-  │                                                     speaker turns            │
-  │                                                     (speakers/timestamps_st… │
-  │                                                     columns, in seconds).    │
-  │ --collar                  FLOAT RANGE [x>=0.0]      DER forgiveness          │
-  │                                                     (seconds) around each    │
-  │                                                     reference turn boundary. │
-  │                                                     [default: 1.0]           │
-  │ --json            -j                                Output the rows and      │
-  │                                                     summary as one JSON      │
-  │                                                     object.                  │
-  │ --help                                              Show this message and    │
-  │                                                     exit.                    │
+  │ --split                   TEXT                     Hugging Face split to     │
+  │                                                    score (default: test).    │
+  │ --subset                  TEXT                     Hugging Face              │
+  │                                                    config/subset name (e.g.  │
+  │                                                    a language).              │
+  │ --limit                   INTEGER RANGE            Rows to evaluate (1-100). │
+  │                           [1<=x<=100]              [default: 10]             │
+  │ --audio-column            TEXT                     Audio column name         │
+  │                                                    (default: auto-detect).   │
+  │ --text-column             TEXT                     Reference text column     │
+  │                                                    name (default:            │
+  │                                                    auto-detect).             │
+  │ --speech-model            [universal-3-pro|univer  Speech model to evaluate. │
+  │                           sal-2]                                             │
+  │ --language-code           TEXT                     Force a language (e.g.    │
+  │                                                    en_us).                   │
+  │ --speaker-labels                                   Diarize and also score    │
+  │                                                    DER against the dataset's │
+  │                                                    reference speaker turns   │
+  │                                                    (speakers/timestamps_sta… │
+  │                                                    columns, in seconds).     │
+  │ --collar                  FLOAT RANGE [x>=0.0]     DER forgiveness (seconds) │
+  │                                                    around each reference     │
+  │                                                    turn boundary.            │
+  │                                                    [default: 1.0]            │
+  │ --json            -j                               Output the rows and       │
+  │                                                    summary as one JSON       │
+  │                                                    object.                   │
+  │ --help                                             Show this message and     │
+  │                                                    exit.                     │
   ╰──────────────────────────────────────────────────────────────────────────────╯
 
    Examples
    Score a model on 10 rows of an HF dataset
-   $ assembly eval distil-whisper/meanwhile
+   $ assembly eval sanchit-gandhi/tedlium-data
    Compare models on your own audio
-   $ assembly eval calls.csv --speech-model universal
+   $ assembly eval calls.csv --speech-model universal-3-pro
    Score diarization too (WER + DER)
    $ assembly eval agent-calls.jsonl --speaker-labels
    Pick a subset/split and more rows
    $ assembly eval openslr/librispeech_asr --subset clean --limit 50
    Evaluate non-English audio
-   $ assembly eval PolyAI/minds14 --subset fr-FR --language-code fr
+   $ assembly eval fixie-ai/common_voice_17_0 --subset fr --language-code fr
    DER on a Hugging Face diarization set
-   $ assembly eval diarizers-community/simsamu --speaker-labels
+   $ assembly eval talkbank/callhome --subset eng --speaker-labels
 
 
 

tests/test_eval_command.py

@@ -7,9 +7,9 @@
 import contextlib
 import dataclasses
 import json
+import re
 from types import SimpleNamespace
 
-import assemblyai as aai
 import pytest
 from typer.testing import CliRunner
 
@@ -123,24 +123,43 @@ def test_json_payload_shape(tmp_path, mocker):
     assert payload["rows"][1] == {"item": "b.wav", "words": 2, "errors": 1, "wer": 0.5}
 
 
-def test_speech_model_flag_reaches_config_and_output(tmp_path, mocker):
+@pytest.mark.parametrize("model", ["universal-3-pro", "universal-2"])
+def test_speech_model_flag_reaches_config_and_output(tmp_path, mocker, model):
     _auth()
     _write_wer_manifest(tmp_path)
     tx = _mock_transcribe(mocker, [_transcript("hello there"), _transcript("goodbye now")])
-    result = runner.invoke(app, ["eval", "manifest.csv", "--speech-model", "universal", "--json"])
+    result = runner.invoke(app, ["eval", "manifest.csv", "--speech-model", model, "--json"])
     assert result.exit_code == 0
-    assert _payload_of(result)["speech_model"] == "universal"
+    assert _payload_of(result)["speech_model"] == model
     tx_config = tx.call_args.kwargs["config"]
-    assert tx_config.speech_model == aai.SpeechModel.universal
+    assert tx_config.speech_models == [model]
     assert tx_config.speaker_labels is None  # not requested -> omitted, not False
 
 
+def test_no_speech_model_leaves_speech_models_unset(tmp_path, mocker):
+    _auth()
+    _write_wer_manifest(tmp_path)
+    tx = _mock_transcribe(mocker, [_transcript("hello there"), _transcript("goodbye now")])
+    assert runner.invoke(app, ["eval", "manifest.csv"]).exit_code == 0
+    assert tx.call_args.kwargs["config"].speech_models is None
+
+
+@pytest.mark.parametrize("model", ["best", "nano", "slam-1", "universal"])
+def test_legacy_models_are_a_usage_error(model):
+    result = runner.invoke(app, ["eval", "manifest.csv", "--speech-model", model])
+    assert result.exit_code == 2
+    # CI forces color on (Rich under GITHUB_ACTIONS), interleaving style codes
+    # mid-message, so assert on the color-free render (see test_help_rendering.py).
+    plain = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
+    assert "Invalid value for '--speech-model'" in plain
+
+
 def test_speech_model_named_in_human_header(tmp_path, mocker):
     _auth()
     _write_wer_manifest(tmp_path)
     _mock_transcribe(mocker, [_transcript("hello there"), _transcript("goodbye now")])
-    result = runner.invoke(app, ["eval", "manifest.csv", "--speech-model", "universal"])
-    assert "universal" in result.output
+    result = runner.invoke(app, ["eval", "manifest.csv", "--speech-model", "universal-3-pro"])
+    assert "universal-3-pro" in result.output
     assert "default model" not in result.output