From 255fe8c2532be31026ef59aaf8f47b98011800f0 Mon Sep 17 00:00:00 2001 From: Eduardo Pacheco Date: Mon, 21 Sep 2026 16:35:53 -0300 Subject: [PATCH 1/2] WhisperKitPro: match whisperkitpro-cli 3.x flags and expose Sortformer model options whisperkitpro-cli 3.x selects the diarization backend with --diarizer and has no --engine or --fast-load flags, so orchestration runs with diarization enabled failed on released CLI builds with "Unknown option". Emit --diarizer and stop emitting --fast-load (the `fast_load` field is kept as a documented no-op so existing configs keep loading). Expose the Sortformer options the CLI already supports: `sortformer_model_version` and `sortformer_model_variant` on the engine and orchestration configs, and `speaker_models_path` on the orchestration config. All are optional and only emitted when set. Also add unit tests for the generated CLI arguments. --- src/openbench/engine/whisperkitpro_engine.py | 24 ++++-- .../orchestration_whisperkitpro.py | 17 +++- tests/engine/__init__.py | 0 tests/engine/test_whisperkitpro_engine.py | 79 +++++++++++++++++++ 4 files changed, 114 insertions(+), 6 deletions(-) create mode 100644 tests/engine/__init__.py create mode 100644 tests/engine/test_whisperkitpro_engine.py diff --git a/src/openbench/engine/whisperkitpro_engine.py b/src/openbench/engine/whisperkitpro_engine.py index b6a9297..af92ac6 100644 --- a/src/openbench/engine/whisperkitpro_engine.py +++ b/src/openbench/engine/whisperkitpro_engine.py @@ -115,13 +115,21 @@ class WhisperKitProConfig(BaseModel): "pyannote", description="The engine to use. If `sortformer` the diarization model used is Sortformer, otherwise it is pyannote.", ) + sortformer_model_version: str | None = Field( + None, + description="Sortformer model version (e.g. `v2-1` or `v3-preview`). Only applicable when `engine` is `sortformer`.", + ) + sortformer_model_variant: str | None = Field( + None, + description="Sortformer model variant (e.g. `384_94MB` or `684_98MB`). Only applicable when `engine` is `sortformer`.", + ) use_exclusive_reconciliation: bool = Field( False, description="Whether to use exclusive reconciliation", ) fast_load: bool = Field( False, - description="Whether to use fast load", + description="Deprecated: whisperkitpro-cli 3.x has no --fast-load flag. Accepted for config compatibility but never emitted.", ) @property @@ -158,6 +166,8 @@ def generate_cli_args(self, model_path: Path | None = None) -> list[str]: ] # Common args + # NOTE: whisperkitpro-cli 3.x has no --fast-load flag, so `fast_load` is + # intentionally not emitted. args.extend( [ "--report", # Always generate the report files @@ -169,8 +179,6 @@ def generate_cli_args(self, model_path: Path | None = None) -> list[str]: COMPUTE_UNITS_MAPPER[self.audio_encoder_compute_units], "--text-decoder-compute-units", COMPUTE_UNITS_MAPPER[self.text_decoder_compute_units], - "--fast-load", - str(self.fast_load).lower(), "--verbose", ] ) @@ -188,11 +196,17 @@ def generate_cli_args(self, model_path: Path | None = None) -> list[str]: # Add rttm path args.extend(["--rttm-path", self.rttm_path]) - args.extend(["--engine", self.engine]) + # whisperkitpro-cli 3.x selects the diarization backend with --diarizer + # (there is no --engine flag). + args.extend(["--diarizer", self.engine]) - # Only add diarization mode if using Sortformer + # Only add Sortformer-specific options if using Sortformer if self.engine == "sortformer": args.extend(["--diarization-mode", self.diarization_mode]) + if self.sortformer_model_version: + args.extend(["--sortformer-model-version", self.sortformer_model_version]) + if self.sortformer_model_variant: + args.extend(["--sortformer-model-variant", self.sortformer_model_variant]) # If speaker models path is provided use it if self.speaker_models_path: diff --git a/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py b/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py index 9609e73..4cc45e4 100644 --- a/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py +++ b/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py @@ -81,13 +81,25 @@ class WhisperKitProOrchestrationConfig(OrchestrationConfig): "prerecorded", description="Sortformer streaming mode: `realtime` (1.04s latency) or `prerecorded` (9.84s latency). This is only applicable when `engine` is `sortformer`.", ) + sortformer_model_version: str | None = Field( + None, + description="Sortformer model version (e.g. `v2-1` or `v3-preview`). Only applicable when `engine` is `sortformer`.", + ) + sortformer_model_variant: str | None = Field( + None, + description="Sortformer model variant (e.g. `384_94MB` or `684_98MB`). Only applicable when `engine` is `sortformer`.", + ) + speaker_models_path: str | None = Field( + None, + description="Local directory used as the speaker-models download base (passed as --speaker-models-path).", + ) use_exclusive_reconciliation: bool = Field( False, description="Whether to use exclusive reconciliation", ) fast_load: bool = Field( False, - description="Whether to use fast load", + description="Deprecated: whisperkitpro-cli 3.x has no --fast-load flag. Accepted for config compatibility but never emitted.", ) @@ -113,6 +125,9 @@ def build_pipeline(self) -> WhisperKitPro: orchestration_strategy=self.config.orchestration_strategy, engine=self.config.engine, diarization_mode=self.config.diarization_mode, + sortformer_model_version=self.config.sortformer_model_version, + sortformer_model_variant=self.config.sortformer_model_variant, + speaker_models_path=self.config.speaker_models_path, use_exclusive_reconciliation=self.config.use_exclusive_reconciliation, fast_load=self.config.fast_load, ) diff --git a/tests/engine/__init__.py b/tests/engine/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/engine/test_whisperkitpro_engine.py b/tests/engine/test_whisperkitpro_engine.py new file mode 100644 index 0000000..56fda63 --- /dev/null +++ b/tests/engine/test_whisperkitpro_engine.py @@ -0,0 +1,79 @@ +# For licensing see accompanying LICENSE.md file. +# Copyright (C) 2025 Argmax, Inc. All Rights Reserved. + +"""Tests for WhisperKitProConfig.generate_cli_args. + +The generated argument list must match what whisperkitpro-cli 3.x accepts: +the diarization backend is selected with `--diarizer` (there is no `--engine` +flag) and `--fast-load` does not exist. Released builds reject unknown +options, so nothing outside that surface may be emitted. +""" + +import unittest +from pathlib import Path + +from openbench.engine.whisperkitpro_engine import WhisperKitProConfig + + +MODEL_DIR = "/models/parakeet" + + +def _args(**overrides) -> list[str]: + config = WhisperKitProConfig(model_dir=MODEL_DIR, **overrides) + return config.generate_cli_args(model_path=Path(MODEL_DIR)) + + +def _value(args: list[str], flag: str) -> str: + return args[args.index(flag) + 1] + + +class TestCommonArgs(unittest.TestCase): + def test_fast_load_is_never_emitted(self) -> None: + for fast_load in (True, False): + args = _args(fast_load=fast_load) + self.assertNotIn("--fast-load", args) + + def test_no_diarization_flags_when_diarization_disabled(self) -> None: + args = _args(diarization=False, engine="sortformer") + for flag in ("--diarization", "--diarizer", "--engine", "--diarization-mode"): + self.assertNotIn(flag, args) + + +class TestDiarizationArgs(unittest.TestCase): + def test_pyannote_uses_diarizer_flag(self) -> None: + args = _args(diarization=True, engine="pyannote") + self.assertIn("--diarization", args) + self.assertEqual(_value(args, "--diarizer"), "pyannote") + self.assertNotIn("--engine", args) + # Sortformer-only options must not leak into pyannote runs + for flag in ("--diarization-mode", "--sortformer-model-version"): + self.assertNotIn(flag, args) + + def test_sortformer_defaults_omit_optional_flags(self) -> None: + args = _args(diarization=True, engine="sortformer") + self.assertEqual(_value(args, "--diarizer"), "sortformer") + self.assertEqual(_value(args, "--diarization-mode"), "prerecorded") + for flag in ("--sortformer-model-version", "--sortformer-model-variant", "--speaker-models-path"): + self.assertNotIn(flag, args) + + def test_sortformer_model_selection(self) -> None: + args = _args( + diarization=True, + engine="sortformer", + diarization_mode="realtime", + sortformer_model_version="v3-preview", + sortformer_model_variant="684_98MB", + speaker_models_path="/speaker-models", + ) + self.assertEqual(_value(args, "--diarization-mode"), "realtime") + self.assertEqual(_value(args, "--sortformer-model-version"), "v3-preview") + self.assertEqual(_value(args, "--sortformer-model-variant"), "684_98MB") + self.assertEqual(_value(args, "--speaker-models-path"), "/speaker-models") + + def test_exclusive_reconciliation_flag(self) -> None: + self.assertIn("--use-exclusive-reconciliation", _args(diarization=True, use_exclusive_reconciliation=True)) + self.assertNotIn("--use-exclusive-reconciliation", _args(diarization=True)) + + +if __name__ == "__main__": + unittest.main() From 510d48d0564435cc04241cc1fe34487ff63d18ac Mon Sep 17 00:00:00 2001 From: Eduardo Pacheco Date: Tue, 22 Sep 2026 14:21:41 -0300 Subject: [PATCH 2/2] Use the published nemotron-3-diarization Sortformer name in examples The Sortformer v3 model was published as `nemotron-3-diarization` with the `684_74MB` variant, so use those names in the field descriptions and tests instead of the local preview names. --- src/openbench/engine/whisperkitpro_engine.py | 4 ++-- .../pipeline/orchestration/orchestration_whisperkitpro.py | 4 ++-- tests/engine/test_whisperkitpro_engine.py | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/openbench/engine/whisperkitpro_engine.py b/src/openbench/engine/whisperkitpro_engine.py index af92ac6..5ba58c3 100644 --- a/src/openbench/engine/whisperkitpro_engine.py +++ b/src/openbench/engine/whisperkitpro_engine.py @@ -117,11 +117,11 @@ class WhisperKitProConfig(BaseModel): ) sortformer_model_version: str | None = Field( None, - description="Sortformer model version (e.g. `v2-1` or `v3-preview`). Only applicable when `engine` is `sortformer`.", + description="Sortformer model version (e.g. `v2-1` or `nemotron-3-diarization`). Only applicable when `engine` is `sortformer`.", ) sortformer_model_variant: str | None = Field( None, - description="Sortformer model variant (e.g. `384_94MB` or `684_98MB`). Only applicable when `engine` is `sortformer`.", + description="Sortformer model variant (e.g. `384_94MB` or `684_74MB`). Only applicable when `engine` is `sortformer`.", ) use_exclusive_reconciliation: bool = Field( False, diff --git a/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py b/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py index 4cc45e4..b09be43 100644 --- a/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py +++ b/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py @@ -83,11 +83,11 @@ class WhisperKitProOrchestrationConfig(OrchestrationConfig): ) sortformer_model_version: str | None = Field( None, - description="Sortformer model version (e.g. `v2-1` or `v3-preview`). Only applicable when `engine` is `sortformer`.", + description="Sortformer model version (e.g. `v2-1` or `nemotron-3-diarization`). Only applicable when `engine` is `sortformer`.", ) sortformer_model_variant: str | None = Field( None, - description="Sortformer model variant (e.g. `384_94MB` or `684_98MB`). Only applicable when `engine` is `sortformer`.", + description="Sortformer model variant (e.g. `384_94MB` or `684_74MB`). Only applicable when `engine` is `sortformer`.", ) speaker_models_path: str | None = Field( None, diff --git a/tests/engine/test_whisperkitpro_engine.py b/tests/engine/test_whisperkitpro_engine.py index 56fda63..7bb794d 100644 --- a/tests/engine/test_whisperkitpro_engine.py +++ b/tests/engine/test_whisperkitpro_engine.py @@ -61,13 +61,13 @@ def test_sortformer_model_selection(self) -> None: diarization=True, engine="sortformer", diarization_mode="realtime", - sortformer_model_version="v3-preview", - sortformer_model_variant="684_98MB", + sortformer_model_version="nemotron-3-diarization", + sortformer_model_variant="684_74MB", speaker_models_path="/speaker-models", ) self.assertEqual(_value(args, "--diarization-mode"), "realtime") - self.assertEqual(_value(args, "--sortformer-model-version"), "v3-preview") - self.assertEqual(_value(args, "--sortformer-model-variant"), "684_98MB") + self.assertEqual(_value(args, "--sortformer-model-version"), "nemotron-3-diarization") + self.assertEqual(_value(args, "--sortformer-model-variant"), "684_74MB") self.assertEqual(_value(args, "--speaker-models-path"), "/speaker-models") def test_exclusive_reconciliation_flag(self) -> None: