Update app.py

app.py

@@ -44,9 +44,11 @@ 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",
+          "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",
-          "HuggingFaceTB/SmolLM3-3B", "winninghealth/WiNGPT-Babel-2",
+          "HuggingFaceTB/SmolLM3-3B",
+          "winninghealth/WiNGPT-Babel-2-1", "winninghealth/WiNGPT-Babel-2", "winninghealth/WiNGPT-Babel",
          "tencent/Hunyuan-MT-7B",
          "openGPT-X/Teuken-7B-instruct-commercial-v0.4", "openGPT-X/Teuken-7B-instruct-v0.6",
          ]
@@ -208,6 +210,38 @@ 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 salamandraTApipe(self):
+        pipe = pipeline("translation", model=self.model)
+        messages = [
+            f"Translate the following text from {self.sl} into {self.tl}.\n{self.sl}: {self.input_text} \n{self.tl}:",
+        ]
+        return pipe(messages)[0]["generated_text"]
+    
+    def salamandraTA(self):
+        # from datetime import datetime
+        text = f"Translate the following text from {self.sl} into {self.tl}.\n{self.sl}: {self.input_text} \n{self.tl}:"
+        tokenizer = AutoTokenizer.from_pretrained(self.model)        
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model,
+            device_map="auto",
+            torch_dtype=torch.bfloat16
+          )
+        message = [ { "role": "user", "content": text } ]
+        # date_string = datetime.today().strftime('%Y-%m-%d')
+        prompt = tokenizer.apply_chat_template(
+            message,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        # date_string=date_string
+        inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+        input_length = inputs.shape[1]
+        outputs = model.generate(input_ids=inputs.to(model.device), 
+                                 max_new_tokens=512,
+                                 early_stopping=True,
+                                 num_beams=5)
+        return tokenizer.decode(outputs[0, input_length:], skip_special_tokens=True))
     
     def HelsinkiNLP_mulroa(self):
         try:
@@ -578,6 +612,9 @@ def translate_text(model_name: str, s_language: str, t_language: str, input_text
         elif model_name == 'Google':
             translated_text = Translators(model_name, sl, tl, input_text).google()
     
+        elif "salamandraTA" in model_name.lower():
+            translated_text = Translators(model_name, sl, tl, input_text).salamandraTApipe()
+
         elif "m2m" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).mtom()