added clear button

app.py

@@ -135,6 +135,7 @@ def instantiate_chatbot(sel, key):
         key=key,
         show_label=False,
         show_share_button=False,
+        show_copy_all_button=True,
         avatar_images=("img/usr.png", bot_icon_select(model_name)),
     )
 
@@ -156,6 +157,26 @@ def instantiate_arrow_button(route, text_route):
     return button
 
 
+def instantiate_retry_button(route):
+    return gr.Button(
+        "",
+        size="sm",
+        scale=0,
+        min_width=40,
+        icon="img/retry.png",
+    )
+
+
+def instantiate_trash_button():
+    return gr.Button(
+        "",
+        size="sm",
+        scale=0,
+        min_width=40,
+        icon="img/trash.png",
+    )
+
+
 def instantiate_text_box():
     return gr.Textbox(label="Your message:", submit_btn=True, key="msg")
 
@@ -169,13 +190,39 @@ def instantiate_additional_settings():
     return system_message
 
 
+def retry_fn(history):
+    if len(history) >= 2:
+        msg = history.pop(-1)  # assistant message
+        msg = history.pop(-1)  # user message
+        return msg["content"], history
+    else:
+        return gr.skip(), gr.skip()
+
+
+def retry_fn_both(history1, history2):
+    msg1, history1 = retry_fn(history1)
+    msg2, history2 = retry_fn(history2)
+    if isinstance(msg1, str) and isinstance(msg2, str):
+        if msg1 == msg2:
+            msg = msg1
+        else:
+            msg = msg1 + " / " + msg2
+    elif isinstance(msg1, str):
+        msg = msg1
+    elif isinstance(msg2, str):
+        msg = msg2
+    else:
+        msg = msg1
+    return msg, history1, history2
+
+
 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")
 
 # to correctly align the left/right arrows
-CSS = ".elems_justif_right {align-items: end;}"
+CSS = ".stick-to-the-right {align-items: end; justify-content: end}"
 
 with gr.Blocks(fill_width=True, title="Keras demo", css=CSS) as demo:
 
@@ -217,29 +264,49 @@ with gr.Blocks(fill_width=True, title="Keras demo", css=CSS) as demo:
         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)
+                with gr.Column(scale=0, min_width=100):
+                    with gr.Row():
+                        instantiate_arrow_button(TextRoute.LEFT, text_route)
+                        retry = instantiate_retry_button(route)
+                    with gr.Row():
+                        instantiate_arrow_button(TextRoute.RIGHT, text_route)
+                        trash = instantiate_trash_button()
+            retry.click(
+                retry_fn_both,
+                inputs=[chatbot1, chatbot2],
+                outputs=[msg, chatbot1, chatbot2],
+            )
+            trash.click(lambda: ("", [], []), outputs=[msg, chatbot1, chatbot2])
 
         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)
+                    with gr.Row():
+                        instantiate_arrow_button(TextRoute.RIGHT, text_route)
+                        retry = instantiate_retry_button(route)
+                    with gr.Row():
+                        instantiate_arrow_button(TextRoute.BOTH, text_route)
+                        trash = instantiate_trash_button()
+            retry.click(retry_fn, inputs=[chatbot1], outputs=[msg, chatbot1])
+            trash.click(lambda: ("", []), outputs=[msg, chatbot1])
 
         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, elem_classes="stick-to-the-right"):
+                    with gr.Row(elem_classes="stick-to-the-right"):
+                        retry = instantiate_retry_button(route)
+                        instantiate_arrow_button(TextRoute.LEFT, text_route)
+                    with gr.Row(elem_classes="stick-to-the-right"):
+                        trash = instantiate_trash_button()
+                        instantiate_arrow_button(TextRoute.BOTH, text_route)
                 with gr.Column(scale=1):
                     msg = instantiate_text_box()
+            retry.click(retry_fn, inputs=[chatbot2], outputs=[msg, chatbot2])
+            trash.click(lambda: ("", []), outputs=[msg, chatbot2])
 
-        with gr.Row():
-            clear = gr.ClearButton([msg, chatbot1, chatbot2])
-            system_message = instantiate_additional_settings()
+        system_message = instantiate_additional_settings()
 
         # Route the submitted message to the left, right or both chatbots
         if route == TextRoute.LEFT:

chatstate.py

@@ -88,7 +88,7 @@ class ChatState:
         self.add_to_history_as_user(message)
         prompt = self.get_full_prompt()
         response = self.model.generate(
-            prompt, max_length=1024, strip_prompt=True
+            prompt, max_length=2048, strip_prompt=True
         )
         self.add_to_history_as_model(response)
         return (message, response)