Support setting max speakers for offline diarization (#97)

* fix: accept input for max_speaker_count in asr/transcribe_file_offline

* fix: rename input field to diarization_max_speakers

* remove: redundant default value for max_speakers

Author: pskrunner14

riva/client/argparse_utils.py

@@ -49,6 +49,12 @@ def add_asr_config_argparse_parameters(
         action='store_true',
         help="Flag that controls if speaker diarization should be performed",
     )
+    parser.add_argument(
+        "--diarization-max-speakers",
+        default=3,
+        type=int,
+        help="Max number of speakers to detect when performing speaker diarization",
+    )
     parser.add_argument(
         "--start-history",
         default=-1,

riva/client/asr.py

@@ -117,10 +117,14 @@ def add_audio_file_specs_to_config(
 def add_speaker_diarization_to_config(
     config: Union[rasr.RecognitionConfig],
     diarization_enable: bool,
+    diarization_max_speakers: int,
 ) -> None:
     inner_config: rasr.RecognitionConfig = config if isinstance(config, rasr.RecognitionConfig) else config.config
     if diarization_enable:
-        diarization_config = rasr.SpeakerDiarizationConfig(enable_speaker_diarization=True)
+        diarization_config = rasr.SpeakerDiarizationConfig(
+            enable_speaker_diarization=True,
+            max_speaker_count=diarization_max_speakers,
+        )
         inner_config.diarization_config.CopyFrom(diarization_config)
 
 

scripts/asr/transcribe_file_offline.py

@@ -37,7 +37,7 @@ def main() -> None:
         enable_word_time_offsets=args.word_time_offsets or args.speaker_diarization,
     )
     riva.client.add_word_boosting_to_config(config, args.boosted_lm_words, args.boosted_lm_score)
-    riva.client.add_speaker_diarization_to_config(config, args.speaker_diarization)
+    riva.client.add_speaker_diarization_to_config(config, args.speaker_diarization, args.diarization_max_speakers)
     riva.client.add_endpoint_parameters_to_config(
         config,
         args.start_history,