add more models

app.py

@@ -6,7 +6,9 @@ import itertools
 DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 MODEL_IDS = [
     "70M",
-    # "160M",
+    "160M",
+    "410M",
+    "610M",
 ]
 MODEL_MAPPING = {
     model_id: f"LinguaCustodia/multilingual-multidomain-fin-mt-{model_id}"
@@ -157,16 +159,13 @@ def translate_with_all_models(text, tgt_lang, src_lang, domain):
     src_lang = None if src_lang == "Auto" else LANG2CODE.get(src_lang)
     domain = DOMAIN_MAPPING[domain]
 
-    res = {
-        model_name: translate_with_model(model_name, text, tgt_lang, src_lang, domain)
-        for model_name in MODEL_IDS
-    }
-    return list(
-        itertools.chain.from_iterable(
-            [res[model_id][k] for k in ("translation", "source_lang", "domain")]
-            for model_id in MODEL_IDS
-        )
-    )
+    outputs = [None] * (3 * len(MODEL_IDS))
+    for i, model_id in enumerate(MODEL_IDS):
+        model_output = translate_with_model(model_id, text, tgt_lang, src_lang, domain)
+        outputs[i * 3] = model_output["translation"]
+        outputs[i * 3 + 1] = model_output["source_lang"]
+        outputs[i * 3 + 2] = model_output["domain"]
+        yield outputs
 
 
 with gr.Blocks() as demo: