Update app.py

app.py

@@ -18,10 +18,10 @@ iso1_to_name = {iso[1]: iso[0] for iso in non_empty_isos} # {'ro': 'Romanian', '
 langs = list(favourite_langs.keys())
 langs.extend(list(all_langs.keys())) # Language options as list, add favourite languages first
 
-models = ["Helsinki-NLP",
+models = ["Helsinki-NLP", "QUICK-MT", "Argos", "Google", "HPLT", "HPLT-OPUS",
           "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul", "Helsinki-NLP/opus-mt-tc-bible-big-mul-deu_eng_nld",
           "Helsinki-NLP/opus-mt-tc-bible-big-mul-deu_eng_fra_por_spa", "Helsinki-NLP/opus-mt-tc-bible-big-deu_eng_fra_por_spa-mul",
-          "Helsinki-NLP/opus-mt-tc-bible-big-roa-deu_eng_fra_por_spa", "Helsinki-NLP/opus-mt-tc-bible-big-deu_eng_fra_por_spa-roa",
+          "Helsinki-NLP/opus-mt-tc-bible-big-roa-deu_eng_fra_por_spa", "Helsinki-NLP/opus-mt-tc-bible-big-deu_eng_fra_por_spa-roa", "Helsinki-NLP/opus-mt-tc-bible-big-roa-en"
           "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/nllb-200-1.3B", "facebook/nllb-200-3.3B",
           "facebook/mbart-large-50-many-to-many-mmt", "facebook/mbart-large-50-one-to-many-mmt", "facebook/mbart-large-50-many-to-one-mmt",
           "facebook/m2m100_418M", "facebook/m2m100_1.2B", "Lego-MT/Lego-MT",
@@ -30,11 +30,9 @@ models = ["Helsinki-NLP",
           "t5-small", "t5-base", "t5-large",
           "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large", "google/flan-t5-xl",
           "google/madlad400-3b-mt", "jbochi/madlad400-3b-mt",
-          "Argos", "Google",
           "HuggingFaceTB/SmolLM3-3B", "winninghealth/WiNGPT-Babel-2",
           "utter-project/EuroLLM-1.7B", "utter-project/EuroLLM-1.7B-Instruct",
-          "Unbabel/Tower-Plus-2B", "Unbabel/TowerInstruct-7B-v0.2", "Unbabel/TowerInstruct-Mistral-7B-v0.2",
-          "openGPT-X/Teuken-7B-instruct-commercial-v0.4", "openGPT-X/Teuken-7B-instruct-v0.6"
+          "Unbabel/Tower-Plus-2B", "Unbabel/TowerInstruct-7B-v0.2", "Unbabel/TowerInstruct-Mistral-7B-v0.2"
           ]
 DEFAULTS = [langs[0], langs[1], models[0]]
 
@@ -78,8 +76,37 @@ class Translators:
         response = httpx.get(url)
         return response.json()[0][0][0]
     
-    @classmethod
-    def download_argos_model(cls, from_code, to_code):
+    def simplepipe(self):
+        try:
+            pipe = pipeline("translation", model=self.model_name, device=self.device)                
+            translation = pipe(self.input_text)
+            message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with {self.model_name}.'
+            return translation[0]['translation_text'], message
+        except Exception as error:
+            return f"Error translating with model: {self.model_name}! Try other available language combination or model.", error
+    
+    def hplt(self, opus = False):
+        # langs = ['ar', 'bs', 'ca', 'en', 'et', 'eu', 'fi', 'ga', 'gl', 'hi', 'hr', 'is', 'mt', 'nn', 'sq', 'sw', 'zh_hant']
+        hplt_models = ['ar-en', 'bs-en', 'ca-en', 'en-ar', 'en-bs', 'en-ca', 'en-et', 'en-eu', 'en-fi',
+                  'en-ga', 'en-gl', 'en-hi', 'en-hr', 'en-is', 'en-mt', 'en-nn', 'en-sq', 'en-sw',
+                  'en-zh_hant', 'et-en', 'eu-en', 'fi-en', 'ga-en', 'gl-en', 'hi-en', 'hr-en',
+                  'is-en', 'mt-en', 'nn-en', 'sq-en', 'sw-en', 'zh_hant-en']
+        if opus:
+            hplt_model = f'HPLT/translate-{self.sl}-{self.tl}-v1.0-hplt_opus' # HPLT/translate-en-hr-v1.0-hplt_opus
+        else:
+            hplt_model = f'HPLT/translate-{self.sl}-{self.tl}-v1.0-hplt' # HPLT/translate-en-hr-v1.0-hplt
+        if f'{self.sl}-{self.tl}' in hplt_models:
+            pipe = pipeline("translation", model=hplt_model, device=self.device)
+            translation = pipe(self.input_text)
+            translated_text = translation[0]['translation_text']
+            message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with {hplt_model}.'
+        else:
+            translated_text = f'HPLT model from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} not available!'
+            message = f"Available models: {', '.join(hplt_models)}"
+        return translated_text, message
+
+    @staticmethod
+    def download_argos_model(from_code, to_code):
         import argostranslate.package
         print('Downloading model', from_code, to_code) 
         # Download and install Argos Translate package
@@ -103,6 +130,63 @@ class Translators:
             translated_text = error
         return translated_text
 
+    @staticmethod
+    def quickmttranslate(model_path, input_text):
+        from quickmt import Translator
+        # '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, max_input_length = 512, max_decoding_length = 512)
+        # print(model_path, input_text, translation)
+        return translation
+
+    @staticmethod
+    def quickmtdownload(model_name):
+        from quickmt.hub import hf_download
+        from pathlib import Path
+        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):
+        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()
+        # print(quickmt_models)
+        quickmt_models = ['ar-en', 'bn-en', 'cs-en', 'da-en', 'de-en', 'el-en', 'en-ar', 'en-bn', 'en-cs', 'en-de', 'en-el', 'en-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', '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']
+        # 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']
+        # Direct translation model
+        if f"{self.sl}-{self.tl}" in quickmt_models:
+            model_path = Translators.quickmtdownload(model_name)
+            translated_text = Translators.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 = Translators.quickmtdownload(model_name)
+            entranslation = Translators.quickmttranslate(model_path, self.input_text)
+            model_name = f"quickmt-en-{self.tl}"
+            model_path = Translators.quickmtdownload(model_name)
+            translated_text = Translators.quickmttranslate(model_path, entranslation)
+            message = f'Translated from {iso1_to_name[self.sl]} to {iso1_to_name[self.tl]} with pivot language English.'
+        else:
+            translated_text = 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 translated_text, message
+        
     def HelsinkiNLP_mulroa(self):
         try:
             pipe = pipeline("translation", model=self.model_name, device=self.device)                
@@ -319,34 +403,6 @@ class Translators:
             output = output.rsplit(f'{self.tl}:')[-1].strip().replace('assistant\n', '').strip()
         return output
 
-    def teuken(self):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.bfloat16,
-        )
-        model = model.to(device).eval()
-        tokenizer = AutoTokenizer.from_pretrained(
-            self.model_name,
-            use_fast=False,
-            trust_remote_code=True,
-        )
-        translation_prompt = f"Translate the following text from {self.sl} into {self.tl}: {self.input_text}"
-        messages = [{"role": "User", "content": translation_prompt}]
-        prompt_ids = tokenizer.apply_chat_template(messages, chat_template="EN", tokenize=True, add_generation_prompt=False, return_tensors="pt")
-        prediction = model.generate(
-            prompt_ids.to(model.device),
-            max_length=512,
-            do_sample=True,
-            top_k=50,
-            top_p=0.95,
-            temperature=0.7,
-            num_return_sequences=1,
-        )
-        translation = tokenizer.decode(prediction[0].tolist())
-        return translation
-
     def unbabel(self):
         pipe = pipeline("text-generation", model=self.model_name, torch_dtype=torch.bfloat16, device_map="auto")
         messages = [{"role": "user",
@@ -422,10 +478,16 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
         
         elif model_name == 'Argos':
             translated_text = Translators(model_name, sl, tl, input_text).argos()
+
+        elif model_name == "QUICK-MT":
+            translated_text, message_text = Translators(model_name, sl, tl, input_text).quickmt()
     
         elif model_name == 'Google':
             translated_text = Translators(model_name, sl, tl, input_text).google()
     
+        elif model_name == "Helsinki-NLP/opus-mt-tc-bible-big-roa-en":
+            translated_text, message_text = Translators(model_name, sl, tl, input_text).simplepipe()
+        
         elif "m2m" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).mtom()
 
@@ -459,10 +521,7 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
     
         elif model_name == "facebook/mbart-large-50-many-to-one-mmt":
             translated_text = Translators(model_name, s_language, t_language, input_text).mbart_many_to_one()
-        
-        elif 'teuken' in model_name.lower():
-            translated_text = Translators(model_name, s_language, t_language, input_text).teuken()
-    
+          
         elif model_name == "utter-project/EuroLLM-1.7B-Instruct":
             translated_text = Translators(model_name, s_language, t_language, input_text).eurollm_instruct()
         
@@ -478,6 +537,12 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
         elif model_name == "winninghealth/WiNGPT-Babel-2":      
             translated_text = Translators(model_name, s_language, t_language, input_text).wingpt()
 
+        elif "HPLT" in model_name:
+            if model_name == "HPLT-OPUS":
+                translated_text, message = Translators(model_name, sl, tl, input_text).hplt(opus = True)
+            else:
+                translated_text, message = Translators(model_name, sl, tl, input_text).hplt()
+        
         elif model_name == "Bergamot":
             translated_text, message_text = Translators(model_name, s_language, t_language, input_text).bergamot()