Add mitre

app.py

@@ -29,6 +29,7 @@ models = ["Helsinki-NLP", "QUICKMT", "Argos", "Google", "HPLT", "HPLT-OPUS",
           "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
           "google/madlad400-3b-mt", "jbochi/madlad400-3b-mt",
           "NiuTrans/LMT-60-0.6B", "NiuTrans/LMT-60-1.7B", "NiuTrans/LMT-60-4B",
+          "naist-nlp/mitre_466m", "naist-nlp/mitre_913m",
           "Lego-MT/Lego-MT", "BSC-LT/salamandraTA-2b-instruct",
           "winninghealth/WiNGPT-Babel", "winninghealth/WiNGPT-Babel-2", "winninghealth/WiNGPT-Babel-2.1",
           "Unbabel/Tower-Plus-2B", "HuggingFaceTB/SmolLM3-3B", "Unbabel/TowerInstruct-7B-v0.2",
@@ -40,7 +41,7 @@ DEFAULTS = [langs[0], langs[1], models[0]]
 
 def timer(func):
     from time import time
-    def translate_text(input_text, s_language, t_language, model_name) -> tuple[str, str]:
+    def translate(input_text, s_language, t_language, model_name) -> tuple[str, str]:
         start_time = time()
         translated_text, message_text = func(input_text, s_language, t_language, model_name)
         end_time = time()
@@ -48,7 +49,7 @@ def timer(func):
         # print(f"Function {func.__name__!r} executed in {execution_time:.2f} seconds.")
         message_text = f'Executed in {execution_time:.2f} seconds! {message_text}' 
         return translated_text, message_text
-    return translate_text
+    return translate
 
 def model_to_cuda(model):
     # Move the model to GPU if available
@@ -101,6 +102,18 @@ class Translators:
         except Exception as error:
             return f"Error translating with model: {self.model_name}! Try other available language combination or model.", error
 
+    def mitre(self):
+        from transformers import AutoModel, AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True, use_fast=False)
+        model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True).to(self.device)
+        model.eval()
+        # Translating from one or several sentences to a sole language
+        src_tokens = tokenizer.encode_source_tokens_to_input_ids(self.input_text, target_language=self.tl)
+        with torch.inference_mode(): # no_grad inference_mode
+            generated_tokens = model.generate(src_tokens)
+        result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        return result
+    
     def niutrans(self):
         tokenizer = AutoTokenizer.from_pretrained(self.model_name, padding_side='left')
         model = AutoModelForCausalLM.from_pretrained(self.model_name)
@@ -583,6 +596,9 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
         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 'mitre' in model_name.lower():
+            translated_text = Translators(model_name, sl, tl, input_text).mitre()
+        
         elif "m2m" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).mtom()