Update app.py

app.py

@@ -14,19 +14,11 @@ options.extend(list(all_langs.keys()))
 models = ["Helsinki-NLP",
           "t5-small", "t5-base", "t5-large",
           "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large", "google/flan-t5-xl",
-          "facebook/nllb-200-distilled-600M",
-          "facebook/nllb-200-distilled-1.3B",
-          "facebook/mbart-large-50-many-to-many-mmt",
-          "bigscience/mt0-small",
-          "bigscience/mt0-base",
-          "bigscience/mt0-large",
-          "bigscience/mt0-xl",
-          "bigscience/bloomz-560m",
-          "bigscience/bloomz-1b1",
-          "bigscience/bloomz-1b7",
-          "bigscience/bloomz-3b",
-          "utter-project/EuroLLM-1.7B",
-          "utter-project/EuroLLM-1.7B-Instruct",
+          "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/nllb-200-1.3B",
+          "facebook/mbart-large-50-many-to-many-mmt", "facebook/mbart-large-50-one-to-many-mmt",
+          "bigscience/mt0-small", "bigscience/mt0-base", "bigscience/mt0-large", "bigscience/mt0-xl",
+          "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
+          "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",
@@ -183,7 +175,7 @@ def unbabel(model_name, sl, tl, input_text):
     translated_text = translated_text.replace('Answer:', '', 1).strip() if translated_text.startswith('Answer:') else translated_text 
     return translated_text
 
-def mbart(model_name, sl, tl, input_text):
+def mbart_many_to_many(model_name, sl, tl, input_text):
     from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
     model = MBartForConditionalGeneration.from_pretrained(model_name)
     tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
@@ -196,6 +188,20 @@ def mbart(model_name, sl, tl, input_text):
     )
     return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
+def mbart_one_to_many(model_name, sl, tl, input_text):
+    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+    article_en = input_text
+    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
+    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
+    model_inputs = tokenizer(article_en, return_tensors="pt")
+    # translate from English to Romanian
+    langid = languagecodes.mbart_large_languages[tl]
+    generated_tokens = model.generate(
+        **model_inputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id[langid]
+    )
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
 @spaces.GPU
 def translate_text(input_text: str, sselected_language: str, tselected_language: str, model_name: str) -> tuple[str, str]:
     """
@@ -255,8 +261,12 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
         translated_text = nllb(model_name, nnlbsl, nnlbtl, input_text)
         return translated_text, message_text
     
-    elif model_name.startswith('facebook/mbart-large'):
-        translated_text = mbart(model_name, sselected_language, tselected_language, input_text)
+    elif model_name == "facebook/mbart-large-50-many-to-many-mmt":
+        translated_text = mbart_many_to_many(model_name, sselected_language, tselected_language, input_text)
+        return translated_text, message_text
+
+    elif model_name == "facebook/mbart-large-50-one-to-many-mmt":
+        translated_text = mbart_one_to_many(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
         
     elif 'Unbabel' in model_name: