Add DNT and custom translation support in NMT client (#108)

* changes made for custom translation, dnt phrases

* formatted

* minor fixes done

* changes made

* minor fix done

* changes done

* changes made

* changes made for dnt and custom translation

* .gitmodules changed

---------

Co-authored-by: Manisha Johnson <manishaj@nvidia.com>

Author: ydharavath

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "common"]
 	path = common
 	url = https://github.com/nvidia-riva/common.git
-	branch = main
+	branch = release/2.18.0

common

@@ -22,6 +22,13 @@ def streaming_s2t_request_generator(
     for chunk in audio_chunks:
         yield riva_nmt.StreamingTranslateSpeechToTextRequest(audio_content=chunk)
 
+def add_dnt_phrases_dict(req, dnt_phrases_dict):
+    if dnt_phrases_dict is not None:
+        dnt_phrases = [f"{key}##{value}" for key, value in dnt_phrases_dict.items()]
+    if dnt_phrases:
+        result_dnt_phrases = ",".join(dnt_phrases)
+        req.dnt_phrases.append(result_dnt_phrases)
+
 class NeuralMachineTranslationClient:
     """
     A class for translating text to text. Provides :meth:`translate` which returns translated text
@@ -137,6 +144,7 @@ def translate(
         source_language: str,
         target_language: str,
         future: bool = False,
+        dnt_phrases_dict: Optional[dict] = None,
     ) -> 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`
@@ -158,7 +166,7 @@ def translate(
             source_language=source_language,
             target_language=target_language
         )
-
+        add_dnt_phrases_dict(req, dnt_phrases_dict)
         func = self.stub.TranslateText.future if future else self.stub.TranslateText
         return func(req, metadata=self.auth.get_auth_metadata())
 

riva/client/nmt.py

@@ -38,6 +38,35 @@
 from riva.client.argparse_utils import add_connection_argparse_parameters
 
 
+def read_dnt_phrases_file(file_path):
+    dnt_phrases_dict = {}
+    if file_path:
+        try:
+            with open(file_path, "r") as infile:
+                for line in infile:
+                    # Trim leading and trailing whitespaces
+                    line = line.strip()
+
+                    if line:
+                        pos = line.find("##")
+                        if pos != -1:
+                            # Line contains "##"
+                            key = line[:pos].strip()
+                            value = line[pos + 2 :].strip()
+                        else:
+                            # Line doesn't contain "##"
+                            key = line.strip()
+                            value = ""
+
+                        # Add the key-value pair to the dictionary
+                        if key:
+                            dnt_phrases_dict[key] = value
+
+        except IOError:
+            raise RuntimeError(f"Could not open file {file_path}")
+
+    return dnt_phrases_dict
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Neural machine translation by Riva AI Services",
@@ -48,6 +77,7 @@ def parse_args() -> argparse.Namespace:
         "--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("--dnt-phrases-file", type=str, help="Path to file which contains dnt phrases and custom translations")
     parser.add_argument("--model-name", default="", type=str, help="model to use to translate")
     parser.add_argument(
         "--source-language-code", type=str, default="en-US", help="Source language code (according to BCP-47 standard)"
@@ -65,7 +95,17 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     def request(inputs,args):
         try:
-            response = nmt_client.translate(inputs, args.model_name, args.source_language_code, args.target_language_code)
+            dnt_phrases_input = {}
+            if args.dnt_phrases_file != None:
+                dnt_phrases_input = read_dnt_phrases_file(args.dnt_phrases_file)
+            response = nmt_client.translate(
+                texts=inputs,
+                model=args.model_name,
+                source_language=args.source_language_code,
+                target_language=args.target_language_code,
+                future=False,
+                dnt_phrases_dict=dnt_phrases_input,
+            )
             for translation in response.translations:
                 print(translation.text)
         except grpc.RpcError as e: