Adding AST support to python clients

riva/client/__init__.py

@@ -12,6 +12,7 @@
     print_streaming,
     sleep_audio_length,
     add_endpoint_parameters_to_config,
+    add_ast_parameters_to_config,
 )
 from riva.client.auth import Auth
 from riva.client.nlp import (

riva/client/argparse_utils.py

@@ -85,6 +85,24 @@ def add_asr_config_argparse_parameters(
         type=float,
         help="Threshold value for likelihood of blanks before detecting end of utterance",
     )
+    parser.add_argument(
+        "--source-language",
+        default="",
+        type=str,
+        help="Language of the audio file",
+    )
+    parser.add_argument(
+        "--task",
+        default="transcribe",
+        type=str,
+        help="Task for the model (transcribe/translate)",
+    )
+    parser.add_argument(
+        "--target-language",
+        default="",
+        type=str,
+        help="Target language for translation",
+    )
     return parser
 
 

riva/client/asr.py

@@ -123,6 +123,7 @@ def add_speaker_diarization_to_config(
         diarization_config = rasr.SpeakerDiarizationConfig(enable_speaker_diarization=True)
         inner_config.diarization_config.CopyFrom(diarization_config)
 
+
 def add_endpoint_parameters_to_config(
     config: Union[rasr.RecognitionConfig, rasr.EndpointingConfig],
     start_history: int,
@@ -152,6 +153,20 @@ def add_endpoint_parameters_to_config(
     inner_config.endpointing_config.CopyFrom(endpointing_config)
 
 
+def add_ast_parameters_to_config(
+    config: Union[rasr.RecognitionConfig, rasr.EndpointingConfig],
+    source_language: str,
+    target_language: str,
+    task: str,
+) -> None:
+    if not source_language:
+        return
+    inner_config: rasr.RecognitionConfig = config if isinstance(config, rasr.RecognitionConfig) else config.config
+    inner_config.custom_configuration["source_language"] = source_language
+    inner_config.custom_configuration["target_language"] = target_language
+    inner_config.custom_configuration["task"] = task
+
+
 PRINT_STREAMING_ADDITIONAL_INFO_MODES = ['no', 'time', 'confidence']
 
 

scripts/asr/transcribe_file_offline.py

@@ -39,14 +39,20 @@ def main() -> None:
     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)
     riva.client.add_endpoint_parameters_to_config(
-        config,
-        args.start_history,
-        args.start_threshold,
-        args.stop_history,
-        args.stop_history_eou,
+        config, 
+        args.start_history, 
+        args.start_threshold, 
+        args.stop_history, 
+        args.stop_history_eou, 
         args.stop_threshold,
         args.stop_threshold_eou
-    )    
+    )
+    riva.client.add_ast_parameters_to_config(
+        config,
+        args.source_language,
+        args.target_language,
+        args.task
+    )
     with args.input_file.open('rb') as fh:
         data = fh.read()
     try:

scripts/tts/talk.py

@@ -157,6 +157,8 @@ def main() -> None:
                 sound_stream(resp.audio)
             if out_f is not None:
                 out_f.writeframesraw(resp.audio)
+    except Exception as e:
+        print(e.details())
     finally:
         if out_f is not None:
             out_f.close()