added ability to chat with only the left or right bot

app.py

@@ -18,6 +18,7 @@ from gradio import ChatMessage
 import keras_hub
 
 from chatstate import ChatState
+from enum import Enum
 from models import (
     model_presets,
     load_model,
@@ -26,6 +27,13 @@ from models import (
     get_appropriate_chat_template,
 )
 
+
+class TextRoute(Enum):
+    LEFT = 0
+    RIGHT = 1
+    BOTH = 2
+
+
 model_labels_list = list(model_labels)
 
 # load and warm up (compile) all the models
@@ -41,21 +49,23 @@ for preset in model_presets:
 
 # For local debugging
 # model = keras_hub.models.Llama3CausalLM.from_preset(
-#     "hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16"
+#     # "hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16"
+#     "../misc-code/ari_tiny_llama3"
 # )
 # models = [model, model, model, model, model]
 
 
-def chat_turn_assistant_1(
-    model,
+def chat_turn_assistant(
     message,
+    sel,
     history,
     system_message,
-    preset,
     # max_tokens,
     # temperature,
     # top_p,
 ):
+    model = models[sel]
+    preset = model_presets[sel]
     chat_template = get_appropriate_chat_template(preset)
     chat_state = ChatState(model, system_message, chat_template)
 
@@ -71,35 +81,25 @@ def chat_turn_assistant_1(
     return history
 
 
-def chat_turn_assistant(
-    message,
-    sel1,
-    history1,
-    sel2,
-    history2,
-    system_message,
-    # max_tokens,
-    # temperature,
-    # top_p,
+def chat_turn_both_assistant(
+    message, sel1, sel2, history1, history2, system_message
 ):
-    history1 = chat_turn_assistant_1(
-        models[sel1], message, history1, system_message, model_presets[sel1]
-    )
-    history2 = chat_turn_assistant_1(
-        models[sel2], message, history2, system_message, model_presets[sel2]
+    return (
+        chat_turn_assistant(message, sel1, history1, system_message),
+        chat_turn_assistant(message, sel2, history2, system_message),
     )
-    return "", history1, history2
 
 
-def chat_turn_user_1(message, history):
+def chat_turn_user(message, history):
     history.append(ChatMessage(role="user", content=message))
     return history
 
 
-def chat_turn_user(message, history1, history2):
-    history1 = chat_turn_user_1(message, history1)
-    history2 = chat_turn_user_1(message, history2)
-    return "", history1, history2
+def chat_turn_both_user(message, history1, history2):
+    return (
+        chat_turn_user(message, history1),
+        chat_turn_user(message, history2),
+    )
 
 
 def bot_icon_select(model_name):
@@ -115,57 +115,72 @@ def bot_icon_select(model_name):
     return "img/bot.png"
 
 
-def instantiate_chatbots(sel1, sel2):
-    model_name1 = model_presets[sel1]
-    chatbot1 = gr.Chatbot(
-        type="messages",
+def instantiate_select_box(sel, model_labels):
+    return gr.Dropdown(
+        choices=[(name, i) for i, name in enumerate(model_labels)],
         show_label=False,
-        show_share_button=False,
-        avatar_images=("img/usr.png", bot_icon_select(model_name1)),
+        value=sel,
+        info="<span style='color:black'>Selected model:</span> <a href='"
+        + preset_to_website_url(model_presets[sel])
+        + "'>"
+        + preset_to_website_url(model_presets[sel])
+        + "</a>",
     )
-    model_name2 = model_presets[sel2]
-    chatbot2 = gr.Chatbot(
+
+
+def instantiate_chatbot(sel, key):
+    model_name = model_presets[sel]
+    return gr.Chatbot(
         type="messages",
+        key=key,
         show_label=False,
         show_share_button=False,
-        avatar_images=("img/usr.png", bot_icon_select(model_name2)),
+        avatar_images=("img/usr.png", bot_icon_select(model_name)),
     )
-    return chatbot1, chatbot2
 
 
-def instantiate_select_boxes(sel1, sel2, model_labels):
-    sel1 = gr.Dropdown(
-        choices=[(name, i) for i, name in enumerate(model_labels)],
-        show_label=False,
-        info="<span style='color:black'>Selected model 1:</span> "
-        + "<a href='"
-        + preset_to_website_url(model_presets[sel1])
-        + "'>"
-        + preset_to_website_url(model_presets[sel1])
-        + "</a>",
-        value=sel1,
+def instantiate_arrow_button(route, text_route):
+    icons = {
+        TextRoute.LEFT: "img/arrowL.png",
+        TextRoute.RIGHT: "img/arrowR.png",
+        TextRoute.BOTH: "img/arrowRL.png",
+    }
+    button = gr.Button(
+        "",
+        size="sm",
+        scale=0,
+        min_width=40,
+        icon=icons[route],
     )
-    sel2 = gr.Dropdown(
-        choices=[(name, i) for i, name in enumerate(model_labels)],
-        show_label=False,
-        info="<span style='color:black'>Selected model 2:</span> "
-        + "<a href='"
-        + preset_to_website_url(model_presets[sel2])
-        + "'>"
-        + preset_to_website_url(model_presets[sel2])
-        + "</a>",
-        value=sel2,
-    )
-    return sel1, sel2
+    button.click(lambda: route, outputs=[text_route])
+    return button
+
+
+def instantiate_text_box():
+    return gr.Textbox(label="Your message:", submit_btn=True, key="msg")
 
 
-def instantiate_chatbots_and_select_boxes(sel1, sel2, model_labels):
-    chatbot1, chatbot2 = instantiate_chatbots(sel1, sel2)
-    sel1, sel2 = instantiate_select_boxes(sel1, sel2, model_labels)
-    return sel1, chatbot1, sel2, chatbot2
+def instantiate_additional_settings():
+    with gr.Accordion("Additional settings", open=False):
+        system_message = gr.Textbox(
+            label="Sytem prompt",
+            value="You are a helpful assistant and your name is Eliza.",
+        )
+    return system_message
+
 
+sel1 = instantiate_select_box(0, model_labels_list)
+sel2 = instantiate_select_box(1, model_labels_list)
+chatbot1 = instantiate_chatbot(sel1.value, "chat1")
+chatbot2 = instantiate_chatbot(sel2.value, "chat2")
 
-with gr.Blocks(fill_width=True, title="Keras demo") as demo:
+# to correctly align the left/right arrows
+CSS = ".elems_justif_right {align-items: end;}"
+
+with gr.Blocks(fill_width=True, title="Keras demo", css=CSS) as demo:
+
+    # Where do messages go
+    text_route = gr.State(TextRoute.BOTH)
 
     with gr.Row():
         gr.Image(
@@ -189,45 +204,91 @@ with gr.Blocks(fill_width=True, title="Keras demo") as demo:
             + "This demo is runnig on a Google TPU v5e 2x4 (8 cores) in bfloat16 precision."
         )
     with gr.Row():
-        sel1, sel2 = instantiate_select_boxes(0, 1, model_labels_list)
-
-    with gr.Row():
-        chatbot1, chatbot2 = instantiate_chatbots(sel1.value, sel2.value)
-
-    msg = gr.Textbox(label="Your message:", submit_btn=True)
+        sel1.render(),
+        sel2.render(),
 
     with gr.Row():
-        gr.ClearButton([msg, chatbot1, chatbot2])
-        with gr.Accordion("Additional settings", open=False):
-            system_message = gr.Textbox(
-                label="Sytem prompt",
-                value="You are a helpful assistant and your name is Eliza.",
+        chatbot1.render()
+        chatbot2.render()
+
+    @gr.render(inputs=text_route)
+    def render_text_area(route):
+
+        if route == TextRoute.BOTH:
+            with gr.Row():
+                msg = instantiate_text_box()
+                with gr.Column(scale=0, min_width=40):
+                    instantiate_arrow_button(TextRoute.RIGHT, text_route)
+                    instantiate_arrow_button(TextRoute.LEFT, text_route)
+
+        elif route == TextRoute.LEFT:
+            with gr.Row():
+                with gr.Column(scale=1):
+                    msg = instantiate_text_box()
+                with gr.Column(scale=1):
+                    instantiate_arrow_button(TextRoute.RIGHT, text_route)
+                    instantiate_arrow_button(TextRoute.BOTH, text_route)
+
+        elif route == TextRoute.RIGHT:
+            with gr.Row():
+                with gr.Column(scale=1, elem_classes="elems_justif_right"):
+                    instantiate_arrow_button(TextRoute.LEFT, text_route)
+                    instantiate_arrow_button(TextRoute.BOTH, text_route)
+                with gr.Column(scale=1):
+                    msg = instantiate_text_box()
+
+        with gr.Row():
+            clear = gr.ClearButton([msg, chatbot1, chatbot2])
+            system_message = instantiate_additional_settings()
+
+        # Route the submitted message to the left, right or both chatbots
+        if route == TextRoute.LEFT:
+            submission = msg.submit(
+                chat_turn_user, inputs=[msg, chatbot1], outputs=[chatbot1]
+            ).then(
+                chat_turn_assistant,
+                [msg, sel1, chatbot1, system_message],
+                outputs=[chatbot1],
             )
+        elif route == TextRoute.RIGHT:
+            submission = msg.submit(
+                chat_turn_user, inputs=[msg, chatbot2], outputs=[chatbot2]
+            ).then(
+                chat_turn_assistant,
+                [msg, sel2, chatbot2, system_message],
+                outputs=[chatbot2],
+            )
+        elif route == TextRoute.BOTH:
+            submission = msg.submit(
+                chat_turn_both_user,
+                inputs=[msg, chatbot1, chatbot2],
+                outputs=[chatbot1, chatbot2],
+            ).then(
+                chat_turn_both_assistant,
+                [msg, sel1, sel2, chatbot1, chatbot2, system_message],
+                outputs=[chatbot1, chatbot2],
+            )
+        # In all cases reset text box after submission
+        submission.then(lambda: "", outputs=msg)
 
     sel1.select(
-        lambda sel1, sel2: instantiate_chatbots_and_select_boxes(
-            sel1, sel2, model_labels_list
-        ),
-        inputs=[sel1, sel2],
-        outputs=[sel1, chatbot1, sel2, chatbot2],
+        lambda sel: instantiate_chatbot(sel, "chat1"),
+        inputs=[sel1],
+        outputs=[chatbot1],
+    ).then(
+        lambda sel: instantiate_select_box(sel, model_labels_list),
+        inputs=[sel1],
+        outputs=[sel1],
     )
 
     sel2.select(
-        lambda sel1, sel2: instantiate_chatbots_and_select_boxes(
-            sel1, sel2, model_labels_list
-        ),
-        inputs=[sel1, sel2],
-        outputs=[sel1, chatbot1, sel2, chatbot2],
-    )
-
-    msg.submit(
-        chat_turn_user,
-        inputs=[msg, chatbot1, chatbot2],
-        outputs=[msg, chatbot1, chatbot2],
+        lambda sel: instantiate_chatbot(sel, "chat2"),
+        inputs=[sel2],
+        outputs=[chatbot2],
     ).then(
-        chat_turn_assistant,
-        [msg, sel1, chatbot1, sel2, chatbot2, system_message],
-        outputs=[msg, chatbot1, chatbot2],
+        lambda sel: instantiate_select_box(sel, model_labels_list),
+        inputs=[sel2],
+        outputs=[sel2],
     )