Update app.py

app.py

@@ -24,14 +24,28 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
                 model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
             except EnvironmentError as error:
                 return f"Error finding model: {model_name_full}! Try other available language combination.", error
-    elif model_name.startswith('facebook/nllb'):
+    if model_name.startswith('facebook/nllb'):
         from languagecodes import nllb_language_codes
         tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=nllb_language_codes[sselected_language])
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
         translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
         translated_text = translator(input_text, max_length=512)
         return translated_text[0]['translation_text'], message_text
-    else:
+    if model_name.startswith('facebook/mbart-large'):
+        from languagecodes import mbart_large_languages
+        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+        model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+        tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+        # translate source to target
+        tokenizer.src_lang = mbart_large_languages[sselected_language]
+        encoded = tokenizer(article_hi, return_tensors="pt")
+        generated_tokens = model.generate(
+            **encoded,
+            forced_bos_token_id=tokenizer.lang_code_to_id[mbart_large_languages[tselected_language]]
+        )
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True), message_text
+
+    if model_name.startswith('t5'):
         tokenizer = T5Tokenizer.from_pretrained(model_name)
         model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")