Update app.py

app.py

@@ -30,6 +30,7 @@ models = ["Helsinki-NLP", "QUICKMT", "Argos", "Google", "HPLT", "HPLT-OPUS",
           "google-t5/t5-small", "google-t5/t5-base", "google-t5/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",
+          "NiuTrans/LMT-60-0.6B", "NiuTrans/LMT-60-1.7B", "NiuTrans/LMT-60-4B",
           "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"
@@ -84,6 +85,23 @@ class Translators:
             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 niutrans(self):
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name, padding_side='left')
+        model = AutoModelForCausalLM.from_pretrained(self.model_name)
+        prompt = f"Translate the following text from {self.sl} into {self.tl}.\n{self.sl}: {self.input_text}.\n{self.tl}: "
+        messages = [{"role": "user", "content": prompt}]
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+        generated_ids = model.generate(**model_inputs, max_new_tokens=512, num_beams=5, do_sample=False)
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        outputs = ''.join(outputs) if isinstance(outputs, list) else outputs
+        return outputs
     
     def hplt(self, opus = False):
         # langs = ['ar', 'bs', 'ca', 'en', 'et', 'eu', 'fi', 'ga', 'gl', 'hi', 'hr', 'is', 'mt', 'nn', 'sq', 'sw', 'zh_hant']
@@ -106,12 +124,10 @@ class Translators:
         return translated_text, message_text
 
     @staticmethod
-    def download_argos_model(from_code, to_code):
+    def download_argos_model(available_packages, from_code, to_code):
         import argostranslate.package
         print('Downloading model', from_code, to_code) 
         # Download and install Argos Translate package
-        argostranslate.package.update_package_index()
-        available_packages = argostranslate.package.get_available_packages()
         package_to_install = next(
             filter(lambda x: x.from_code == from_code and x.to_code == to_code, available_packages)
         )
@@ -119,12 +135,19 @@ class Translators:
     
     def argos(self):
         import argostranslate.translate, argostranslate.package
+        argostranslate.package.update_package_index()
+        available_packages = argostranslate.package.get_available_packages()
+        available_slanguages = [lang.from_code for lang in available_packages]
+        available_tlanguages = [lang.to_code for lang in available_packages]
+        available_languages = sorted(list(set(available_slanguages + available_tlanguages)))
+        zipped = sorted(list(zip(available_slanguages, available_tlanguages)))
+        packages_info = ', '.join(f"{pkg.from_name} ({pkg.from_code}) -> {pkg.to_name} ({pkg.to_code})" for pkg in available_packages)
+        print(available_languages, zipped)
         try:
-            Translators.download_argos_model(self.sl, self.tl) # Download model
+            Translators.download_argos_model(available_packages, self.sl, self.tl) # Download model
             translated_text = argostranslate.translate.translate(self.input_text, self.sl, self.tl) # Translate
         except StopIteration:
-            # packages_info = ', '.join(f"{pkg.get_description()}->{str(pkg.links)} {str(pkg.source_languages)}" for pkg in argostranslate.package.get_available_packages())
-            packages_info = ', '.join(f"{pkg.from_name} ({pkg.from_code}) -> {pkg.to_name} ({pkg.to_code})" for pkg in argostranslate.package.get_available_packages())
+            # packages_info = ', '.join(f"{pkg.get_description()}->{str(pkg.links)} {str(pkg.source_languages)}" for pkg in available_packages)
             translated_text = f"No Argos model for {self.sl} to {self.tl}. Try other model or languages combination from the available Argos models: {packages_info}."
         except Exception as error:
             translated_text = error
@@ -159,15 +182,19 @@ class Translators:
         # 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']
+        quickmt_models = ['ar-en', 'bn-en', 'cs-en', 'da-en', 'de-en', 'el-en', 'en-ar', 'en-bn',
+                          'en-cs', 'en-da', '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-sv', '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']
+        available_languages = ['ar', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fr', 'he',
+                               'hi', 'hu', 'id', 'it', 'ja', 'ko', 'lv', 'pl', 'pt', 'ro', 'ru',
+                               'sv', 'th', 'tr', 'ur', 'vi', 'zh']
+        # print(quickmt_models, available_languages)       
         # Direct translation model
         if f"{self.sl}-{self.tl}" in quickmt_models:
             model_path = Translators.quickmtdownload(model_name)
@@ -497,7 +524,10 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
 
         elif "lego" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).LegoMT()
-        
+
+        elif "niutrans" in model_name.lower():
+            translated_text = Translators(model_name, sl, tl, input_text).niutrans()
+
         elif model_name.startswith('google-t5'):
             translated_text = Translators(model_name, s_language, t_language, input_text).tfive()