asr: add model name parameter

nvidia-riva · Jul 17, 2024 · a0a8c66 · a0a8c66
1 parent 330aa60
commit a0a8c66
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 0 deletions.
diff --git a/riva/client/argparse_utils.py b/riva/client/argparse_utils.py
@@ -38,6 +38,7 @@ def add_asr_config_argparse_parameters(
         help="If specified, text inverse normalization will be applied",
     )
     parser.add_argument("--language-code", default="en-US", help="Language code of the model to be used.")
+    parser.add_argument("--model-name", default="", help="Model name to be used.")
     parser.add_argument("--boosted-lm-words", action='append', help="Words to boost when decoding.")
     parser.add_argument(
         "--boosted-lm-score", type=float, default=4.0, help="Value by which to boost words when decoding."

diff --git a/scripts/asr/riva_streaming_asr_client.py b/scripts/asr/riva_streaming_asr_client.py
@@ -55,6 +55,7 @@ def streaming_transcription_worker(
         config = riva.client.StreamingRecognitionConfig(
             config=riva.client.RecognitionConfig(
                 language_code=args.language_code,
+                model=args.model_name,
                 max_alternatives=args.max_alternatives,
                 profanity_filter=args.profanity_filter,
                 enable_automatic_punctuation=args.automatic_punctuation,

diff --git a/scripts/asr/transcribe_file.py b/scripts/asr/transcribe_file.py
@@ -71,6 +71,7 @@ def main() -> None:
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(
             language_code=args.language_code,
+            model=args.model_name,
             max_alternatives=1,
             profanity_filter=args.profanity_filter,
             enable_automatic_punctuation=args.automatic_punctuation,

diff --git a/scripts/asr/transcribe_mic.py b/scripts/asr/transcribe_mic.py
@@ -47,6 +47,7 @@ def main() -> None:
         config=riva.client.RecognitionConfig(
             encoding=riva.client.AudioEncoding.LINEAR_PCM,
             language_code=args.language_code,
+            model=args.model_name,
             max_alternatives=1,
             profanity_filter=args.profanity_filter,
             enable_automatic_punctuation=args.automatic_punctuation,