Update app.py

app.py

@@ -60,12 +60,20 @@ class Translators:
 
     def quickmt(self):
         from quickmt import Translator
-        from quickmt.hub import hf_download, hf_list
+        from quickmt.hub import hf_download
         from pathlib import Path
         model_name = f"quickmt-{self.sl}-{self.tl}"
         model_path = Path("/code/models") / model_name
-        choices = [i.split("/")[1] for i in hf_list()]
-        print(choices)
+        # from quickmt.hub import hf_list
+        # choices = [i.split("/")[1] for i in hf_list()]
+        # print(choices)
+        models = ['zh-en', 'en-zh', 'fr-en', 'en-fr', 'de-en', 'en-de', 'ko-en',
+                  'en-ko', 'en-ar', 'ar-en', 'ja-en', 'en-ja', 'it-en', 'en-it',
+                  'hi-en', 'en-hi', 'en-es', 'es-en', 'ru-en', 'en-ru', 'pt-en',
+                  'en-pt', 'bn-en', 'en-bn', 'id-en', 'en-id', 'fa-en', 'en-fa',
+                  'vi-en', 'en-vi', 'tr-en', 'en-tr', 'ur-en', 'en-ur', 'en-hu',
+                  'hu-en', 'th-en', 'en-th', 'el-en', 'en-el', 'pl-en', 'en-pl',
+                  'en-lv', 'lv-en', 'cs-en', 'ro-en', 'en-ro', 'en-cs', 'he-en', 'en-he', 'da-en']
         if not model_path.exists():
             hf_download(
             model_name = f"quickmt/{model_name}",
@@ -73,10 +81,10 @@ class Translators:
         )
         print(f"Loading model {model_name}")
         # Auto-detects GPU, set to "cpu" to force CPU inference
-        translator = Translator(str(model_path), device="auto", inter_threads=2)       
+        translator = Translator(str(model_path), device = self.device)       
         # translation = Translator(f"./quickmt-{self.sl}-{self.tl}/", device="auto", inter_threads=2)       
         # Translate - set beam size to 1 for faster speed (but lower quality)
-        translation = translator({self.input_text}, beam_size=5)
+        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
     
@@ -427,6 +435,10 @@ def translate_text(model_name: str, s_language: str, t_language: str, input_text
     
     sl = all_langs[s_language][0]
     tl = all_langs[t_language][0]
+    if sl == tl:
+        translated_text = f'Source language {s_language} identical to target language {t_language}!'
+        message_text = 'Please choose different target and source language!'
+        return translated_text, message_text
     message_text = f'Translated from {s_language} to {t_language} with {model_name}'
     translated_text = None
     try: