Models can be disabled

app.py

@@ -14,6 +14,7 @@ MODEL_MAPPING = {
     model_id: f"LinguaCustodia/multilingual-multidomain-fin-mt-{model_id}"
     for model_id in MODEL_IDS
 }
+MODEL_INDEX = {m: i for i, m in enumerate(MODEL_IDS)}
 TOKENIZER = AutoTokenizer.from_pretrained(
     MODEL_MAPPING["70M"],
     pad_token="<pad>",
@@ -154,13 +155,25 @@ def translate_with_model(model_name, text, tgt_lang, src_lang, domain):
     }
 
 
-def translate_with_all_models(text, tgt_lang, src_lang, domain):
+def translate_with_all_models(selected_models, text, tgt_lang, src_lang, domain):
     tgt_lang = LANG2CODE[tgt_lang]
     src_lang = None if src_lang == "Auto" else LANG2CODE.get(src_lang)
     domain = DOMAIN_MAPPING[domain]
 
     outputs = [None] * (3 * len(MODEL_IDS))
-    for i, model_id in enumerate(MODEL_IDS):
+    outputs = list(
+        itertools.chain.from_iterable(
+            (
+                ["Processing..."] * 3
+                if model_id in selected_models
+                else ["This model is disabled"] * 3
+            )
+            for model_id in MODEL_IDS
+        )
+    )
+
+    for model_id in selected_models:
+        i = MODEL_INDEX[model_id]
         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"]
@@ -176,6 +189,13 @@ with gr.Blocks() as demo:
 
     with gr.Row(variant="panel"):
         with gr.Column(variant="default"):
+            selected_models = gr.CheckboxGroup(
+                choices=MODEL_IDS,
+                value=MODEL_IDS,
+                type="value",
+                label="Models",
+                container=True,
+            )
             source_text = gr.Textbox(lines=3, label="Source sentence")
         with gr.Column(variant="default"):
             source_language = gr.Dropdown(
@@ -195,14 +215,18 @@ with gr.Blocks() as demo:
         for model_id in MODEL_IDS:
             with gr.Tab(model_id):
                 outputs[model_id] = {
-                    "translation": gr.Textbox(lines=2, label="Translation"),
+                    "translation": gr.Textbox(
+                        lines=2, label="Translation", container=True
+                    ),
                     "source_lang": gr.Textbox(
                         label="Predicted source language",
                         info='This is the predicted source language, if "Auto" is selected.',
+                        container=True,
                     ),
                     "domain": gr.Textbox(
                         label="Predicted domain",
                         info='This is the predicted domain, if "Auto" is checked.',
+                        container=True,
                     ),
                 }
                 gr.HTML(
@@ -227,7 +251,7 @@ with gr.Blocks() as demo:
 
     translate_btn.click(
         fn=translate_with_all_models,
-        inputs=[source_text, target_language, source_language, domain],
+        inputs=[selected_models, source_text, target_language, source_language, domain],
         outputs=list(
             itertools.chain.from_iterable(
                 [outputs[model_id][k] for k in ("translation", "source_lang", "domain")]