Update app.py

app.py

@@ -18,6 +18,7 @@ models = ["Helsinki-NLP",
           "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/nllb-200-1.3B",
           "facebook/mbart-large-50-many-to-many-mmt", "facebook/mbart-large-50-one-to-many-mmt", "facebook/mbart-large-50-many-to-one-mmt",
+          "facebook/m2m100_418M", "facebook/m2m100_1.2B",
           "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",
@@ -64,6 +65,15 @@ def google(sl, tl, input_text):
     response = requests.get(url)
     return response.json()[0][0][0]
 
+def mtom(model_name, sl, tl, input_text):
+    from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+    tokenizer.src_lang = sl
+    encoded = tokenizer(input_text, return_tensors="pt")
+    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(tl))
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
 def HelsinkiNLPAutoTokenizer(sl, tl, input_text):
     if model_name == "Helsinki-NLP":
         message_text = f'Translated from {sl} to {tl} with {model_name}.'
@@ -269,67 +279,56 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
     
     if model_name.startswith("Helsinki-NLP"):
         translated_text, message_text = HelsinkiNLP(sl, tl, input_text)
-        return translated_text, message_text
     
     elif model_name == 'Argos':
         translated_text = argos(sl, tl, input_text)
-        return translated_text, message_text        
 
     elif model_name == 'Google':
         translated_text = google(sl, tl, input_text)
-        return translated_text, message_text
+
+    elif "m2m" in model_name:
+        translated_text = mtom(model_name, sl, tl, input_text)
     
     elif model_name == "utter-project/EuroLLM-1.7B-Instruct":
         translated_text = eurollm_instruct(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
     
     elif model_name == "utter-project/EuroLLM-1.7B":
         translated_text = eurollm(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
 
     elif 'flan' in model_name.lower():
         translated_text = flan(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
         
     elif 'teuken' in model_name.lower():
         translated_text = teuken(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
 
     elif 'mt0' in model_name.lower():
         translated_text = bigscience(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
 
     elif 'bloomz' in model_name.lower():
         translated_text = bloomz(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
         
     elif 'nllb' in model_name.lower():
         nnlbsl, nnlbtl = languagecodes.nllb_language_codes[sselected_language], languagecodes.nllb_language_codes[tselected_language]
         translated_text = nllb(model_name, nnlbsl, nnlbtl, input_text)
-        return translated_text, message_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 model_name == "facebook/mbart-large-50-many-to-one-mmt":
         translated_text = mbart_many_to_one(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
     
     elif 'Unbabel' in model_name:   
         translated_text = unbabel(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
     
     elif model_name.startswith('t5'):
         translated_text = tfive(model_name, sselected_language, tselected_language, input_text)
-        return translated_text, message_text
 
     else:
-        return "No model selected or error in application", message_text
+        translated_text "No model selected or error in application", message_text
+    return translated_text, message_text
 
 # Define a function to swap dropdown values
 def swap_languages(src_lang, tgt_lang):