Update app.py

app.py

@@ -42,6 +42,7 @@ models = ["Helsinki-NLP", "QUICKMT", "Argos", "Lego-MT/Lego-MT", "HPLT", "HPLT-O
           "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", "Heng666/madlad400-3b-mt-ct2", "Heng666/madlad400-3b-mt-ct2-int8", "Heng666/madlad400-7b-mt-ct2-int8",
+          "naist-nlp/mitre_466m", "naist-nlp/mitre_913m",
           "BSC-LT/salamandraTA-2b-instruct", "BSC-LT/salamandraTA-7b-instruct",
           "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",
@@ -61,6 +62,25 @@ class Translators:
         response = httpx.get(url)
         return response.json()[0][0][0]
 
+    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.half() # recommended
+        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)
+        src_tokens = src_tokens.to(self.device)
+        # Translating from one or several sentences to corresponding languages
+        # src_tokens = tokenizer.encode_source_tokens_to_input_ids_with_different_tags([english_text, english_text, ], target_languages_list=["de", "zh", ])      
+        # generated_tokens = model.generate(src_tokensto(self.device))
+        # results = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+        
+        with torch.no_grad():
+            generated_tokens = model.generate(src_tokens)
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    
     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',
@@ -583,7 +603,10 @@ def translate_text(model_name: str, s_language: str, t_language: str, input_text
                 translated_text, message_text = Translators(model_name, sl, tl, input_text).hplt(opus = True)
             else:
                 translated_text, message_text = Translators(model_name, sl, tl, input_text).hplt()
-                
+
+        elif 'mitre' in model_name.lower():
+            translated_text = Translators(model_name, sl, tl, input_text).mitre()
+        
         elif model_name == 'Argos':
             translated_text = Translators(model_name, sl, tl, input_text).argos()