Update app.py

app.py

@@ -13,7 +13,9 @@ MODELS = {'benro': 'BlackKakapo/opus-mt-en-ro',
           'mttcbig': 'Helsinki-NLP/opus-mt-tc-big-en-ro',
           'gemma': 'Gargaz/gemma-2b-romanian-better',
           'mbartenro': 'ancebuc/mbart-translation-en-ro',
-          't5enro': "ancebuc/t5-translation-en-ro",
+          't5enro': 'ancebuc/t5-translation-en-ro',
+          'pegasus': 'ancebuc/pegasus-translation-en-ro',
+          'mbart': 'facebook/mbart-large-cc25',
           'paraphrase': 'tuner007/pegasus_paraphrase'}
 EMBEDDING_MODELS = {"all-MiniLM-L6-v2":384,
 "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":384,
@@ -70,10 +72,12 @@ def translate(input_text: str, model_name: str = MODELS['mttcbig'], sl: str = 'e
     elif 'Helsinki-NLP' in model_name:
         translation, message = src.Translate.Translators(model_name, sl, tl, input_text).HelsinkiNLP_mulroa()
      # text2textgenerationpipe, translationpipe
-    elif model_name == MODELS['mbartenro']:
-        translation, message = src.Translate.Translators(model_name, sl, tl, input_text).text2textgenerationpipe()
-    elif model_name == MODELS['t5enro']:
-        translation, message = src.Translate.Translators(model_name, sl, tl, input_text).text2textgenerationpipe()
+    # elif model_name == MODELS['mbartenro']:
+    #     translation, message = src.Translate.Translators(model_name, sl, tl, input_text).text2textgenerationpipe()
+    elif model_name == MODELS['t5enro'] or model_name == MODELS['pegasus'] or model_name == MODELS['mbartenro']:
+        translation, message = src.Translate.Translators(model_name, sl, tl, input_text).translationpipe()
+    elif model_name == MODELS['mbart']:
+        translation, message = src.Translate.Translators(model_name, sl, tl, input_text).mbartlarge()
     else:
         translation: str = src.Translate.gemma_direct(input_text, model_name)
     return {"input_text": input_text, "translation": translation, "model_name": model_name, "message": message}