mulitprocessing and signature fixes for remote translators

* remote Riva do not serialize client object * remove extra call to get_generator in atkgen Signed-off-by: Jeffrey Martin <jemartin@nvidia.com>
NVIDIA · SnowMasaya · Oct 9, 2024 · Oct 9, 2024 · Oct 9, 2024 · Oct 9, 2024
commit b44b8b7496d88a60badac70ecc7b0c489cd0b4d9
diff --git a/garak/probes/atkgen.py b/garak/probes/atkgen.py
@@ -318,5 +318,3 @@ def __init__(self, config_root=_config):
             msg = f"No query placeholder {TEMPLATE_PLACEHOLDER} in {self.__class__.__name__} prompt template {self.red_team_prompt_template}"
             logging.critical(msg)
             raise ValueError(msg)
-
-        self.translator = self.get_translator()
diff --git a/garak/translators/remote.py b/garak/translators/remote.py
@@ -36,6 +36,19 @@ class RivaTranslator(Translator):
     ]
     # fmt: on
 
+    # avoid attempt to pickle the client attribute
+    def __getstate__(self) -> object:
+        self._clear_translator()
+        return dict(self.__dict__)
+
+    # restore the client attribute
+    def __setstate__(self, d) -> object:
+        self.__dict__.update(d)
+        self._load_translator()
+
+    def _clear_translator(self):
+        self.nmt_client = None
+
     def _load_translator(self):
         if not (
             self.source_lang in self.bcp47_support
@@ -59,9 +72,11 @@ def _load_translator(self):
             )
             self.nmt_client = riva.client.NeuralMachineTranslationClient(auth)
 
-    def _translate(self, text: str, source_lang: str, target_lang: str) -> str:
+    def _translate(self, text: str) -> str:
         try:
-            response = self.nmt_client.translate([text], "", source_lang, target_lang)
+            response = self.nmt_client.translate(
+                [text], "", self.source_lang, self.target_lang
+            )
             return response.translations[0].text
         except Exception as e:
             logging.error(f"Translation error: {str(e)}")
@@ -104,10 +119,11 @@ def _load_translator(self):
         if self.translator is None:
             self.translator = Translator(self.api_key)
 
-    def _translate(self, text: str, source_lang: str, target_lang: str) -> str:
+    def _translate(self, text: str) -> str:
         try:
+            target_lang = "EN-US" if self.target_lang == "en" else self.target_lang
             return self.translator.translate_text(
-                text, source_lang=source_lang, target_lang=target_lang
+                text, source_lang=self.source_lang, target_lang=target_lang
             ).text
         except Exception as e:
             logging.error(f"Translation error: {str(e)}")