updates: add options to ASR and TTS clients (#118)

- Add/update list model option to ASR clients
- Add encoding option to TTS client

Author: virajkarandikar

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()))

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,

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
             )