Add speaker diarization to streaming ASR clients (#116)

* enable speaker diarization for streaming_asr_client

* add: speaker diarization to transcribe_file

* fix: print confidence when word offsets are disabled

Author: pskrunner14

riva/client/asr.py

@@ -183,6 +183,7 @@ def print_streaming(
     word_time_offsets: bool = False,
     show_intermediate: bool = False,
     file_mode: str = 'w',
+    speaker_diarization: bool = False,
 ) -> None:
     """
     Prints streaming speech recognition results to provided files or streams.
@@ -284,12 +285,21 @@ def print_streaming(
                         if word_time_offsets:
                             for f in output_file:
                                 f.write("Timestamps:\n")
-                                f.write('{: <40s}{: <16s}{: <16s}\n'.format('Word', 'Start (ms)', 'End (ms)'))
+                                temp = '{: <40s}{: <16s}{: <16s}'
+                                value = ['Word', 'Start (ms)', 'End (ms)']
+                                if speaker_diarization:
+                                    temp += '{: <16s}'
+                                    value.append('Speaker')
+                                temp += '\n'
+                                f.write(temp.format(*value))
                                 for word_info in result.alternatives[0].words:
                                     f.write(
                                         f'{word_info.word: <40s}{word_info.start_time: <16.0f}'
-                                        f'{word_info.end_time: <16.0f}\n'
+                                        f'{word_info.end_time: <16.0f}'
                                     )
+                                    if speaker_diarization:
+                                        f.write(f'{word_info.speaker_tag: <16d}')
+                                    f.write('\n')
                     else:
                         partial_transcript += transcript
                 else:  # additional_info == 'confidence'

scripts/asr/riva_streaming_asr_client.py

@@ -60,7 +60,7 @@ def streaming_transcription_worker(
                 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,
+                enable_word_time_offsets=args.word_time_offsets or args.speaker_diarization,
             ),
             interim_results=True,
         )
@@ -78,6 +78,7 @@ def streaming_transcription_worker(
             args.custom_configuration
         )
         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, args.diarization_max_speakers)
         for _ in range(args.num_iterations):
             with riva.client.AudioChunkFileIterator(
                 args.input_file,
@@ -92,7 +93,8 @@ def streaming_transcription_worker(
                     output_file=output_file,
                     additional_info='time',
                     file_mode='a',
-                    word_time_offsets=args.word_time_offsets,
+                    word_time_offsets=args.word_time_offsets or args.speaker_diarization,
+                    speaker_diarization=args.speaker_diarization,
                 )
     except BaseException as e:
         exception_queue.put((e, thread_i))

scripts/asr/transcribe_file.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 import argparse
+from pathlib import Path
 
 import os
 import riva.client
@@ -50,7 +51,7 @@ def parse_args() -> argparse.Namespace:
         "normal speech.",
     )
     parser.add_argument(
-        "--print-confidence", action="store_true", help="Whether to print stability and confidence of transcript."
+        "--print-confidence", action="store_true", help="Whether to print stability and confidence of transcript. If `--word-time-offsets` or `--speaker-diarization` is set, then confidence is not printed."
     )
     parser = add_connection_argparse_parameters(parser)
     parser = add_asr_config_argparse_parameters(parser, max_alternatives=True, profanity_filter=True, word_time_offsets=True)
@@ -88,6 +89,8 @@ def main() -> None:
         print(f"Invalid input file path: {args.input_file}")
         return
 
+    output_file = Path(f"output.txt").expanduser()
+
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(
             language_code=args.language_code,
@@ -96,10 +99,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,
         ),
         interim_results=True,
     )
     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, args.diarization_max_speakers)
     riva.client.add_endpoint_parameters_to_config(
         config,
         args.start_history,
@@ -131,8 +136,12 @@ def main() -> None:
                     audio_chunks=audio_chunk_iterator,
                     streaming_config=config,
                 ),
+                output_file=output_file,
+                file_mode='a',
                 show_intermediate=args.show_intermediate,
-                additional_info="confidence" if args.print_confidence else "no",
+                additional_info="time" if (args.word_time_offsets or args.speaker_diarization) else ("confidence" if args.print_confidence else "no"),
+                word_time_offsets=args.word_time_offsets or args.speaker_diarization,
+                speaker_diarization=args.speaker_diarization,
             )
     finally:
         if sound_callback is not None and sound_callback.opened: