feat(asr): add speaker diarization support in offline client (#25)

Author: virajkarandikar

common

@@ -6,6 +6,7 @@
     ASRService,
     add_audio_file_specs_to_config,
     add_word_boosting_to_config,
+    add_speaker_diarization_to_config,
     get_wav_file_parameters,
     print_offline,
     print_streaming,

riva/client/__init__.py

@@ -42,6 +42,12 @@ def add_asr_config_argparse_parameters(
     parser.add_argument(
         "--boosted-lm-score", type=float, default=4.0, help="Value by which to boost words when decoding."
     )
+    parser.add_argument(
+        "--speaker-diarization",
+        default=False,
+        action='store_true',
+        help="Flag that controls if speaker diarization should be performed",
+    )
     return parser
 
 

riva/client/argparse_utils.py

@@ -99,6 +99,16 @@ def add_audio_file_specs_to_config(
     inner_config.audio_channel_count = wav_parameters['nchannels']
 
 
+def add_speaker_diarization_to_config(
+    config: Union[rasr.RecognitionConfig],
+    diarization_enable: bool,
+) -> 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)
+        inner_config.diarization_config.CopyFrom(diarization_config)
+
+
 PRINT_STREAMING_ADDITIONAL_INFO_MODES = ['no', 'time', 'confidence']
 
 

riva/client/asr.py

@@ -18,7 +18,7 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--input-file", required=True, type=Path, help="A path to a local file to transcribe.")
     parser = add_connection_argparse_parameters(parser)
-    parser = add_asr_config_argparse_parameters(parser, profanity_filter=True)
+    parser = add_asr_config_argparse_parameters(parser, profanity_filter=True, word_time_offsets=True)
     args = parser.parse_args()
     args.input_file = args.input_file.expanduser()
     return args
@@ -35,9 +35,12 @@ def main() -> None:
         profanity_filter=args.profanity_filter,
         enable_automatic_punctuation=args.automatic_punctuation,
         verbatim_transcripts=not args.no_verbatim_transcripts,
+        enable_word_time_offsets=args.word_time_offsets or args.speaker_diarization,
     )
     riva.client.add_audio_file_specs_to_config(config, args.input_file)
     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)
+
     with args.input_file.open('rb') as fh:
         data = fh.read()
     try: