Update app.py

app.py

@@ -58,12 +58,29 @@ class Translators:
         translation = pipe(self.input_text)
         return translation[0]['translation_text']
 
-    def quickmt(self):
+    def quickmttranslate(model_path, input_text):
         from quickmt import Translator
-        from quickmt.hub import hf_download
+        # 'auto' auto-detects GPU, set to "cpu" to force CPU inference
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
+        translator = Translator(str(model_path), device = device)       
+        # translation = Translator(f"./quickmt-{self.sl}-{self.tl}/", device="auto", inter_threads=2)       
+        # set beam size to 1 for faster speed (but lower quality)
+        translation = translator(input_text, beam_size=5)[0]
+        return translation
+
+    def quickmtdownload(model_name):
         from pathlib import Path
-        model_name = f"quickmt-{self.sl}-{self.tl}"
         model_path = Path("/quickmt/models") / model_name
+        if not model_path.exists():
+            hf_download(
+            model_name = f"quickmt/{model_name}",
+            output_dir=Path("/quickmt/models") / model_name,
+        )
+        return model_path
+            
+    def quickmt(self):
+        from quickmt.hub import hf_download
+        model_name = f"quickmt-{self.sl}-{self.tl}"
         # from quickmt.hub import hf_list
         # quickmt_models = [i.split("/quickmt-")[1] for i in hf_list()]
         # quickmt_models.sort()
@@ -72,22 +89,26 @@ class Translators:
                   'en-fa', 'en-fr', 'en-he', 'en-hi', 'en-hu', 'en-id', 'en-it', 'en-ja', 'en-ko', 'en-lv', 'en-pl', 'en-pt',
                   'en-ro', 'en-ru', 'en-th', 'en-tr', 'en-ur', 'en-vi', 'en-zh', 'es-en', 'fa-en', 'fr-en', 'he-en', 'hi-en',
                   'hu-en', 'id-en', 'it-en', 'ja-en', 'ko-en', 'lv-en', 'pl-en', 'pt-en', 'ro-en', 'ru-en', 'th-en', 'tr-en', 'ur-en', 'vi-en', 'zh-en']
-        if f"{self.sl}-{self.tl}" not in quickmt_models:
+        # available_languages = list(set([lang for model in quickmt_models for lang in model.split('-')]))
+        # available_languages.sort()
+        available_languages = ['ar', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fr', 'he', 'hi', 'hu',
+                               'id', 'it', 'ja', 'ko', 'lv', 'pl', 'pt', 'ro', 'ru', 'th', 'tr', 'ur', 'vi', 'zh']
+        if f"{self.sl}-{self.tl}" in quickmt_models:
+            model_path = quickmtdownload(model_name)
+            translation = quickmttranslate(model_path, self.input_text)
+            message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with {model_name}.'
+        # Pivot language English
+        elif self.sl in available_languages and self.tl in available_languages:
+            model_name = f"quickmt-{self.sl}-en"
+            model_path = quickmtdownload(model_name)
+            entranslation = quickmttranslate(model_path, self.input_text)
+            model_name = f"quickmt-en-{self.tl}"
+            model_path = quickmtdownload(model_name)
+            translation = quickmttranslate(model_path, entranslation)
+            message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with pivot language English.'
+        else:                   
             translation = f'Model {model_name} from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} not available!'
             message = f"Available models: {', '.join(quickmt_models)}"
-            return translation, message
-        if not model_path.exists():
-            hf_download(
-            model_name = f"quickmt/{model_name}",
-            output_dir=Path("/quickmt/models") / model_name,
-        )
-        # 'auto' auto-detects GPU, set to "cpu" to force CPU inference
-        device = 'gpu' if torch.cuda.is_available() else 'cpu'
-        translator = Translator(str(model_path), device = device)       
-        # translation = Translator(f"./quickmt-{self.sl}-{self.tl}/", device="auto", inter_threads=2)       
-        # set beam size to 1 for faster speed (but lower quality)
-        translation = translator({self.input_text}, beam_size=5)[0]
-        message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with {model_name}.'
         return translation, message
     
     @classmethod