Update app.py

app.py

@@ -13,7 +13,9 @@ options.extend(list(all_langs.keys()))
 models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large",
           "facebook/nllb-200-distilled-600M",
           "facebook/nllb-200-distilled-1.3B",
-          "facebook/mbart-large-50-many-to-many-mmt"]
+          "facebook/mbart-large-50-many-to-many-mmt",
+          "Unbabel/TowerInstruct-7B-v0.2",
+          "Unbabel/TowerInstruct-Mistral-7B-v0.2"]
 
 def model_to_cuda(model):
     # Move the model to GPU if available
@@ -62,6 +64,15 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         )
         return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0], message_text
 
+    if 'Unbabel' in model_name:   
+        pipe = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
+        messages = [{"role": "user",
+                     "content": f"Translate the following text from {sselected_language} into {tselected_language}.\n{sselected_language}: {input_text}.\n{tselected_language}:"}]
+        prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+        outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
+        translated_text = outputs[0]["generated_text"]
+        return translated_text, message_text
+    
     if model_name.startswith('t5'):
         tokenizer = T5Tokenizer.from_pretrained(model_name)
         model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")