Removed "Unbabel/TowerInstruct-Mistral-7B-v0.2", added "BSC-LT/salamandraTA-2b-instruct"

app.py

@@ -27,14 +27,14 @@ models = ["Helsinki-NLP", "QUICKMT", "Argos", "Google", "HPLT", "HPLT-OPUS",
           "facebook/hf-seamless-m4t-medium", "facebook/seamless-m4t-large", "facebook/seamless-m4t-v2-large",
           "bigscience/mt0-small", "bigscience/mt0-base", "bigscience/mt0-large", "bigscience/mt0-xl",
           "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
-          "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",
-          "Lego-MT/Lego-MT", "HuggingFaceTB/SmolLM3-3B",
+          "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",
           "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"
+          "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"
           ]
 DEFAULTS = [langs[0], langs[1], models[0]]
 
@@ -72,6 +72,7 @@ class Translators:
         self.sl, self.tl = sl, tl
         self.input_text = input_text
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_new_tokens = 512
         
     def google(self):  
         url = os.environ['GCLIENT'] + f'sl={self.sl}&tl={self.tl}&q={self.input_text}'
@@ -104,16 +105,11 @@ class Translators:
         outputs = ''.join(outputs) if isinstance(outputs, list) else outputs
         return outputs
 
-    def smallonehundred(self):
-        from transformers import M2M100ForConditionalGeneration
-        from tokenization_small100 import SMALL100Tokenizer
-        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name)
-        tokenizer = SMALL100Tokenizer.from_pretrained(self.model_name)
-        tokenizer.tgt_lang = self.tl
-        encoded_sl = tokenizer(self.input_text, return_tensors="pt")
-        generated_tokens = model.generate(**encoded_sl, max_length=256, num_beams=5)
-        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-    
+    def salamandratapipe(self):
+        pipe = pipeline("text-generation", model=self.model_name)
+        messages = [{"role": "user", "content": f"Translate the following text from {self.sl} into {self.tl}.\n{self.sl}: {self.input_text} \n{self.tl}:"}]
+        return pipe(messages, max_new_tokens=self.max_new_tokens, early_stopping=True, num_beams=5)[0]["generated_text"][1]["content"]
+     
     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',
@@ -267,16 +263,7 @@ class Translators:
                 return self.HelsinkiNLP_mulroa()
         except KeyError as error:
             return f"Error: Translation direction {self.sl} to {self.tl} is not supported by Helsinki Translation Models", error
-  
-    def LegoMT(self):
-        from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
-        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name) # "Lego-MT/Lego-MT"
-        tokenizer = M2M100Tokenizer.from_pretrained(self.model_name)
-        tokenizer.src_lang = self.sl
-        encoded = tokenizer(self.input_text, return_tensors="pt")
-        generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(self.tl))
-        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-        
+     
     def madlad(self):
         model = T5ForConditionalGeneration.from_pretrained(self.model_name, device_map="auto")
         tokenizer = T5Tokenizer.from_pretrained(self.model_name)
@@ -369,6 +356,25 @@ class Translators:
         generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(self.tl))
         return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
+    def smallonehundred(self):
+        from transformers import M2M100ForConditionalGeneration
+        from tokenization_small100 import SMALL100Tokenizer
+        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = SMALL100Tokenizer.from_pretrained(self.model_name)
+        tokenizer.tgt_lang = self.tl
+        encoded_sl = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded_sl, max_length=256, num_beams=5)
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    
+    def LegoMT(self):
+        from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name) # "Lego-MT/Lego-MT"
+        tokenizer = M2M100Tokenizer.from_pretrained(self.model_name)
+        tokenizer.src_lang = self.sl
+        encoded = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(self.tl))
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]  
+    
     def bigscience(self):  
         tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
@@ -575,6 +581,9 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
         elif "niutrans" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).niutrans()
 
+        elif "salamandra" in model_name.lower():
+            translated_text = Translators(model_name, s_language, t_language, input_text).salamandratapipe()
+        
         elif model_name.startswith('google-t5'):
             translated_text = Translators(model_name, s_language, t_language, input_text).tfive()