feat: s2s streaming demo app

Author: junkin

riva/client/__init__.py

@@ -36,5 +36,6 @@
 from riva.client.proto.riva_asr_pb2 import RecognitionConfig, StreamingRecognitionConfig
 from riva.client.proto.riva_audio_pb2 import AudioEncoding
 from riva.client.proto.riva_nlp_pb2 import AnalyzeIntentOptions
+from riva.client.proto.riva_nmt_pb2 import StreamingTranslateSpeechToSpeechConfig
 from riva.client.tts import SpeechSynthesisService
 from riva.client.nmt import NeuralMachineTranslationClient

riva/client/nmt.py

@@ -1,14 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: MIT
 
-from typing import Generator, Optional, Union, List
-
+from typing import Callable, Dict, Generator, Iterable, List, Optional, TextIO, Union
 from grpc._channel import _MultiThreadedRendezvous
 
 import riva.client.proto.riva_nmt_pb2 as riva_nmt
 import riva.client.proto.riva_nmt_pb2_grpc as riva_nmt_srv
 from riva.client import Auth
 
+def streaming_s2s_request_generator(
+    audio_chunks: Iterable[bytes], streaming_config: riva_nmt.StreamingTranslateSpeechToSpeechConfig
+) -> Generator[riva_nmt.StreamingTranslateSpeechToSpeechRequest, None, None]:
+    yield riva_nmt.StreamingTranslateSpeechToSpeechRequest(config=streaming_config)
+    for chunk in audio_chunks:
+        yield riva_nmt.StreamingTranslateSpeechToSpeechRequest(audio_content=chunk)
+
 
 class NeuralMachineTranslationClient:
     """
@@ -25,6 +31,49 @@ def __init__(self, auth: Auth) -> None:
         self.auth = auth
         self.stub = riva_nmt_srv.RivaTranslationStub(self.auth.channel)
 
+    def streaming_s2s_response_generator(
+        self, audio_chunks: Iterable[bytes], streaming_config: riva_nmt.StreamingTranslateSpeechToSpeechConfig
+    ) -> Generator[riva_nmt.StreamingTranslateSpeechToSpeechResponse, None, None]:
+        """
+        Generates speech recognition responses for fragments of speech audio in :param:`audio_chunks`.
+        The purpose of the method is to perform speech recognition "online" - as soon as
+        audio is acquired on small chunks of audio.
+
+        All available audio chunks will be sent to a server on first ``next()`` call.
+
+        Args:
+            audio_chunks (:obj:`Iterable[bytes]`): an iterable object which contains raw audio fragments
+                of speech. For example, such raw audio can be obtained with
+
+                .. code-block:: python
+
+                    import wave
+                    with wave.open(file_name, 'rb') as wav_f:
+                        raw_audio = wav_f.readframes(n_frames)
+
+            streaming_config (:obj:`riva.client.proto.riva_asr_pb2.StreamingRecognitionConfig`): a config for streaming.
+                You may find description of config fields in message ``StreamingRecognitionConfig`` in
+                `common repo
+                <https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html#riva-proto-riva-asr-proto>`_.
+                An example of creation of streaming config:
+
+                .. code-style:: python
+
+                    from riva.client import RecognitionConfig, StreamingRecognitionConfig
+                    config = RecognitionConfig(enable_automatic_punctuation=True)
+                    streaming_config = StreamingRecognitionConfig(config, interim_results=True)
+
+        Yields:
+            :obj:`riva.client.proto.riva_asr_pb2.StreamingRecognizeResponse`: responses for audio chunks in
+            :param:`audio_chunks`. You may find description of response fields in declaration of
+            ``StreamingRecognizeResponse``
+            message `here
+            <https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html#riva-proto-riva-asr-proto>`_.
+        """
+        generator = streaming_s2s_request_generator(audio_chunks, streaming_config)
+        for response in self.stub.StreamingTranslateSpeechToSpeech(generator, metadata=self.auth.get_auth_metadata()):
+            yield response
+
     def translate(
         self,
         texts: List[str],

scripts/nmt/s2s_mic.py

@@ -0,0 +1,114 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+import argparse
+import wave
+import riva.client
+from riva.client.argparse_utils import add_asr_config_argparse_parameters, add_connection_argparse_parameters
+from typing import Callable, Dict, Generator, Iterable, List, Optional, TextIO, Union
+import riva.client.audio_io
+import riva.client.proto.riva_nmt_pb2 as riva_nmt
+
+def parse_args() -> argparse.Namespace:
+    default_device_info = riva.client.audio_io.get_default_input_device_info()
+    default_device_index = None if default_device_info is None else default_device_info['index']
+    parser = argparse.ArgumentParser(
+        description="Streaming speech to speech translation from microphone via Riva AI Services",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--input-device", type=int, default=default_device_index, help="An input audio device to use.")
+    parser.add_argument("--list-input-devices", action="store_true", help="List input audio device indices.")
+    parser.add_argument("--list-output-devices", action="store_true", help="List input audio device indices.")
+    parser.add_argument("--output-device", type=int, help="Output device to use.")
+    parser.add_argument(
+        "--play-audio",
+        action="store_true",
+        help="Play input audio simultaneously with transcribing and translating it. If `--output-device` is not provided, "
+        "then the default output audio device will be used.",
+    )
+
+    parser = add_asr_config_argparse_parameters(parser, profanity_filter=True)
+    parser = add_connection_argparse_parameters(parser)
+    parser.add_argument(
+        "--sample-rate-hz",
+        type=int,
+        help="A number of frames per second in audio streamed from a microphone.",
+        default=16000,
+    )
+    parser.add_argument(
+        "--file-streaming-chunk",
+        type=int,
+        default=1600,
+        help="A maximum number of frames in a audio chunk sent to server.",
+    )
+    args = parser.parse_args()
+    return args
+
+def play_responses(responses: Iterable[riva_nmt.StreamingTranslateSpeechToSpeechResponse],
+                   sound_stream) -> None:
+    count = 0
+    for response in responses:
+        #if first:
+            #print(f"time to first audio {(stop - start):.3f}s")
+        #    first=False
+        if sound_stream is not None:
+            sound_stream(response.speech.audio)
+            fname = "response" + str(count)
+            out_f = wave.open(fname, 'wb')
+            out_f.setnchannels(1)
+            out_f.setsampwidth(2)
+            out_f.setframerate(44100)
+        count += 1
+
+
+def main() -> None:
+    args = parse_args()
+    sound_stream = None
+    sampwidth = 2
+    nchannels = 1
+    if args.list_input_devices:
+        riva.client.audio_io.list_input_devices()
+        return
+    if args.output_device is not None or args.play_audio:
+        print("playing audio")
+        sound_stream = riva.client.audio_io.SoundCallBack(
+            args.output_device, nchannels=nchannels, sampwidth=sampwidth, framerate=44100
+        )
+        print(sound_stream)
+    first = True # first tts output chunk received
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    nmt_service = riva.client.NeuralMachineTranslationClient(auth)
+    s2s_config = riva.client.StreamingTranslateSpeechToSpeechConfig(
+        asrConfig = riva.client.StreamingRecognitionConfig(
+            config=riva.client.RecognitionConfig(
+                encoding=riva.client.AudioEncoding.LINEAR_PCM,
+                language_code=args.language_code,
+                max_alternatives=1,
+                profanity_filter=args.profanity_filter,
+                enable_automatic_punctuation=args.automatic_punctuation,
+                verbatim_transcripts=not args.no_verbatim_transcripts,
+                sample_rate_hertz=args.sample_rate_hz,
+                audio_channel_count=1,
+            ),
+            interim_results=True,
+        )
+    )
+
+    #riva.client.add_word_boosting_to_config(config, args.boosted_lm_words, args.boosted_lm_score)
+    with riva.client.audio_io.MicrophoneStream(
+        args.sample_rate_hz,
+        args.file_streaming_chunk,
+        device=args.input_device,
+    ) as audio_chunk_iterator:
+        play_responses(responses=nmt_service.streaming_s2s_response_generator(
+            audio_chunks=audio_chunk_iterator,
+            streaming_config=s2s_config), sound_stream=sound_stream)
+        # if first:
+        #         first = False
+        #     if sound_stream is not None:
+        #         sound_stream(response.audio)
+
+
+
+if __name__ == '__main__':
+    main()