feat(nmt): add nmt client (#18)

* add nmt files

Update __init__.py

* bump patch version

* update submodule

* proto location change

* chore: Roll version numbers

* fix(nmt): Match existing style

* Update riva/client/nmt.py

Co-authored-by: PeganovAnton <[email protected]>

* Update scripts/nmt/nmt.py

Co-authored-by: PeganovAnton <[email protected]>

* Update scripts/nmt/nmt.py

Co-authored-by: PeganovAnton <[email protected]>

Co-authored-by: PeganovAnton <[email protected]>

Author: junkin

common

@@ -35,4 +35,5 @@
 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.tts import SpeechSynthesisService
+from riva.client.tts import SpeechSynthesisService
+from riva.client.nmt import NeuralMachineTranslationClient

riva/client/__init__.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+from typing import Generator, Optional, Union, List
+
+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
+
+
+class NeuralMachineTranslationClient:
+    """
+    A class for translating text to text. Provides :meth:`translate` which returns translated text
+    """
+    def __init__(self, auth: Auth) -> None:
+        """
+        Initializes an instance of the class.
+
+        Args:
+            auth (:obj:`Auth`): an instance of :class:`riva.client.auth.Auth` which is used for authentication metadata
+                generation.
+        """
+        self.auth = auth
+        self.stub = riva_nmt_srv.RivaTranslationStub(self.auth.channel)
+
+    def translate(
+        self,
+        texts: List[str],
+        model: str,
+        source_language: str,
+        target_language: str,
+        future: bool = False,
+    ) -> Union[riva_nmt.TranslateTextResponse, _MultiThreadedRendezvous]:
+        """
+        Translate input list of input text :param:`text` using model :param:`model` from :param:`source_language` into :param:`target_language`
+
+        Args:
+            text (:obj:`list[str]`): input text.
+            future (:obj:`bool`, defaults to :obj:`False`): whether to return an async result instead of usual
+                response. You can get a response by calling ``result()`` method of the future object.
+
+        Returns:
+            :obj:`Union[riva.client.proto.riva_nmt_pb2.TranslateTextResponse, grpc._channel._MultiThreadedRendezvous]`:
+            a response with output. You may find :class:`riva.client.proto.riva_nmt_pb2.TranslateTextResponse` fields
+            description `here
+            <https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html#riva-proto-riva-nmt-proto>`_.
+        """
+        req = riva_nmt.TranslateTextRequest(
+            texts=texts,
+            model=model,
+            source_language=source_language,
+            target_language=target_language
+        )
+
+        func = self.stub.TranslateText.future if future else self.stub.TranslateText
+        return func(req, metadata=self.auth.get_auth_metadata())
+
+    def get_config(
+            self,
+            model: str,
+            future: bool = False,
+    ) -> Union[riva_nmt.AvailableLanguageResponse, _MultiThreadedRendezvous]:
+        req = riva_nmt.AvailableLanguageRequest(model=model)
+        func = self.stub.ListSupportedLanguagePairs.future if future else self.stub.ListSupportedLanguagePairs
+        return func(req, metadata=self.auth.get_auth_metadata())

riva/client/nmt.py

@@ -3,7 +3,7 @@
 
 MAJOR = 0
 MINOR = 0
-PATCH = 5
+PATCH = 6
 PRE_RELEASE = 'rc0'
 
 # Use the following formatting: (major, minor, patch, pre-release)
@@ -20,7 +20,7 @@
 __download_url__ = 'hhttps://github.com/nvidia-riva/python-clients/releases'
 __description__ = "Python implementation of the Riva Client API"
 __license__ = 'MIT'
-__keywords__ = 'deep learning, machine learning, gpu, NLP, ASR, TTS, nvidia, speech, language, Riva, client'
-__riva_version__ = "2.3.0"
-__riva_release__ = "22.06"
-__riva_models_version__ = "2.3.0"
+__keywords__ = 'deep learning, machine learning, gpu, NLP, ASR, TTS, NMT, nvidia, speech, language, Riva, client'
+__riva_version__ = "2.7.0"
+__riva_release__ = "22.10"
+__riva_models_version__ = "2.7.0"

riva/client/package_info.py

@@ -0,0 +1,2 @@
+#q SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT

scripts/nmt/__init__.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#!/usr/bin/env python
+
+import argparse
+import os
+import sys
+
+import grpc
+import riva.client.proto.riva_nmt_pb2 as riva_nmt
+import riva.client.proto.riva_nmt_pb2_grpc as riva_nmt_srv
+
+import riva.client
+from riva.client.argparse_utils import add_connection_argparse_parameters
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Neural machine translation by Riva AI Services",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    inputs = parser.add_mutually_exclusive_group()
+    inputs.add_argument(
+        "--text", default="mir Das ist mir Wurs, bien ich ein berliner", type=str, help="Text to translate"
+    )
+    inputs.add_argument("--text-file", type=str, help="Path to file for translation")
+    parser.add_argument("--model-name", default="riva-nmt", type=str, help="model to use to translate")
+    parser.add_argument(
+        "--src-language", type=str, help="Source language (according to BCP-47 standard)"
+    )
+    parser.add_argument(
+        "--tgt-language", type=str, help="Target language (according to BCP-47 standard)"
+    )
+    parser.add_argument("--text-file", type=str, help="Path to file for translation")
+    parser.add_argument("--batch-size", type=int, default=8, help="Batch size to use for file translation")
+    parser.add_argument("--list-models", default=False, action='store_true', help="List available models")
+    parser = add_connection_argparse_parameters(parser)
+
+    return parser.parse_args()
+
+
+def main() -> None:
+    def request(inputs,args):
+        try:
+            response = nmt_client.translate(inputs, args.model_name, args.src_language, args.tgt_language)
+            for translation in response.translations:
+                print(translation.text)
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                result = {'msg': 'invalid arg error'}
+            elif e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                result = {'msg': 'already exists error'}
+            elif e.code() == grpc.StatusCode.UNAVAILABLE:
+                result = {'msg': 'server unavailable check network'}
+            print(f"{result['msg']} : {e.details()}")
+
+    args = parse_args()
+
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    nmt_client = riva.client.NeuralMachineTranslationClient(auth)
+
+    if args.list_models:
+
+        response = nmt_client.get_config(args.model_name)
+        print(response)
+        return
+
+    if args.text_file != None and os.path.exists(args.text_file):
+        with open(args.text_file, "r") as f:
+            batch = []
+            for line in f:
+                line = line.strip()
+                if line != "":
+                    batch.append(line)
+                if len(batch) == args.batch_size:
+                    request(batch, args)
+                    batch = []
+            if len(batch) > 0:
+                request(batch, args)
+        return
+
+    if args.text != "":
+        request([args.text], args)
+
+
+if __name__ == '__main__':
+    main()