From 79b8000eb25398d5df1c06a5ceb85d55ba256eef Mon Sep 17 00:00:00 2001
From: Prabhsimran Singh <pskrunner14@gmail.com>
Date: Sun, 2 Mar 2025 19:49:19 +0530
Subject: [PATCH 1/3] 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
---
 riva/client/asr.py                       | 14 ++++++++++++--
 scripts/asr/riva_streaming_asr_client.py |  6 ++++--
 scripts/asr/transcribe_file.py           | 13 +++++++++++--
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/riva/client/asr.py b/riva/client/asr.py
index 62e17aec..9d209b64 100644
--- a/riva/client/asr.py
+++ b/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'
diff --git a/scripts/asr/riva_streaming_asr_client.py b/scripts/asr/riva_streaming_asr_client.py
index f7018141..43873942 100644
--- a/scripts/asr/riva_streaming_asr_client.py
+++ b/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))
diff --git a/scripts/asr/transcribe_file.py b/scripts/asr/transcribe_file.py
index f558aaf4..610ec426 100644
--- a/scripts/asr/transcribe_file.py
+++ b/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:

From 7ea15e46f808d732bdcadc118602635343c0d588 Mon Sep 17 00:00:00 2001
From: Prabhsimran Singh <pskrunner14@gmail.com>
Date: Sun, 2 Mar 2025 23:20:38 +0530
Subject: [PATCH 2/3] Fix output format of transcribe_file.py client

---
 scripts/asr/transcribe_file.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/scripts/asr/transcribe_file.py b/scripts/asr/transcribe_file.py
index 610ec426..ac153316 100644
--- a/scripts/asr/transcribe_file.py
+++ b/scripts/asr/transcribe_file.py
@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: MIT
 
 import argparse
-from pathlib import Path
 
 import os
 import riva.client
@@ -89,8 +88,6 @@ 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,
@@ -136,8 +133,6 @@ 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="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,

From 17d0b43b865c33c4a82019acdb6da2f1c3502e2c Mon Sep 17 00:00:00 2001
From: Viraj Karandikar <16838694+virajkarandikar@users.noreply.github.com>
Date: Thu, 6 Mar 2025 16:40:41 +0530
Subject: [PATCH 3/3] updates: add options to ASR and TTS clients (#118)

- Add/update list model option to ASR clients
- Add encoding option to TTS client
---
 scripts/asr/transcribe_file.py         |  7 +++---
 scripts/asr/transcribe_file_offline.py | 32 ++++++++++++++++++++++++--
 scripts/tts/talk.py                    |  4 ++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/scripts/asr/transcribe_file.py b/scripts/asr/transcribe_file.py
index ac153316..d9162e7e 100644
--- a/scripts/asr/transcribe_file.py
+++ b/scripts/asr/transcribe_file.py
@@ -72,12 +72,13 @@ def main() -> None:
         asr_models = dict()
         config_response = asr_service.stub.GetRivaSpeechRecognitionConfig(riva.client.proto.riva_asr_pb2.RivaSpeechRecognitionConfigRequest())
         for model_config in config_response.model_config:
-            if model_config.parameters["streaming"] and model_config.parameters["type"]:
+            if model_config.parameters["type"] == "online":
                 language_code = model_config.parameters['language_code']
+                model = {"model": [model_config.model_name]}
                 if language_code in asr_models:
-                    asr_models[language_code]["models"].append(model_config.model_name)
+                    asr_models[language_code].append(model)
                 else:
-                    asr_models[language_code] = {"models": [model_config.model_name]}
+                    asr_models[language_code] = [model]
 
         print("Available ASR models")
         asr_models = dict(sorted(asr_models.items()))
diff --git a/scripts/asr/transcribe_file_offline.py b/scripts/asr/transcribe_file_offline.py
index 33e2d9d0..0529dff7 100644
--- a/scripts/asr/transcribe_file_offline.py
+++ b/scripts/asr/transcribe_file_offline.py
@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: MIT
 
+import os
 import argparse
 from pathlib import Path
 
@@ -16,18 +17,45 @@ def parse_args() -> argparse.Namespace:
         "one response.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser.add_argument("--input-file", required=True, type=Path, help="A path to a local file to transcribe.")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--input-file", type=Path, help="A path to a local file to transcribe.")
+    group.add_argument("--list-models", action="store_true", help="List available models.")
+
     parser = add_connection_argparse_parameters(parser)
     parser = add_asr_config_argparse_parameters(parser, max_alternatives=True, profanity_filter=True, word_time_offsets=True)
     args = parser.parse_args()
-    args.input_file = args.input_file.expanduser()
+    if args.input_file:
+        args.input_file = args.input_file.expanduser()
     return args
 
 
 def main() -> None:
     args = parse_args()
+
     auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
+
+    if args.list_models:
+        asr_models = dict()
+        config_response = asr_service.stub.GetRivaSpeechRecognitionConfig(riva.client.proto.riva_asr_pb2.RivaSpeechRecognitionConfigRequest())
+        for model_config in config_response.model_config:
+            if model_config.parameters["type"] == "offline":
+                language_code = model_config.parameters['language_code']
+                model = {"model": [model_config.model_name]}
+                if language_code in asr_models:
+                    asr_models[language_code].append(model)
+                else:
+                    asr_models[language_code] = [model]
+
+        print("Available ASR models")
+        asr_models = dict(sorted(asr_models.items()))
+        print(asr_models)
+        return
+
+    if not os.path.isfile(args.input_file):
+        print(f"Invalid input file path: {args.input_file}")
+        return
+
     config = riva.client.RecognitionConfig(
         language_code=args.language_code,
         max_alternatives=args.max_alternatives,
diff --git a/scripts/tts/talk.py b/scripts/tts/talk.py
index 80bc78b1..be03331b 100644
--- a/scripts/tts/talk.py
+++ b/scripts/tts/talk.py
@@ -9,6 +9,7 @@
 
 import riva.client
 from riva.client.argparse_utils import add_connection_argparse_parameters
+from riva.client.proto.riva_audio_pb2 import AudioEncoding
 
 def read_file_to_dict(file_path):
     result_dict = {}
@@ -56,6 +57,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--sample-rate-hz", type=int, default=44100, help="Number of audio frames per second in synthesized audio."
     )
+    parser.add_argument("--encoding", default="LINEAR_PCM", choices={"LINEAR_PCM", "OGGOPUS"}, help="Output audio encoding.")
     parser.add_argument("--custom-dictionary", type=str, help="A file path to a user dictionary with key-value pairs separated by double spaces.")
     parser.add_argument(
         "--stream",
@@ -132,6 +134,7 @@ def main() -> None:
         if args.stream:
             responses = service.synthesize_online(
                 args.text, args.voice, args.language_code, sample_rate_hz=args.sample_rate_hz,
+                encoding=AudioEncoding.OGGOPUS if args.encoding == "OGGOPUS" else AudioEncoding.LINEAR_PCM,
                 audio_prompt_file=args.audio_prompt_file, quality=20 if args.quality is None else args.quality,
                 custom_dictionary=custom_dictionary_input
             )
@@ -148,6 +151,7 @@ def main() -> None:
         else:
             resp = service.synthesize(
                 args.text, args.voice, args.language_code, sample_rate_hz=args.sample_rate_hz,
+                encoding=AudioEncoding.OGGOPUS if args.encoding == "OGGOPUS" else AudioEncoding.LINEAR_PCM,
                 audio_prompt_file=args.audio_prompt_file, quality=20 if args.quality is None else args.quality,
                 custom_dictionary=custom_dictionary_input
             )
