diff --git a/src/openbench/engine/whisperkitpro_engine.py b/src/openbench/engine/whisperkitpro_engine.py index b6a9297..5ba58c3 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 `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: diff --git a/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py b/src/openbench/pipeline/orchestration/orchestration_whisperkitpro.py index 9609e73..b09be43 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 `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, ) 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..7bb794d --- /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="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()