WhisperKitPro: match whisperkitpro-cli 3.x flags and expose Sortformer model options

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 `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_74MB`). 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:

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 `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_74MB`). 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,
         )

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="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"), "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:
+        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()