fix: split user and bot

app.py

@@ -60,9 +60,11 @@ with gr.Blocks() as demo:
 
     tokenizer, model = load_model()
 
-    def on_message(
-        user_message: str, history: list[list[str]], alpha: float, topk: int
-    ):
+    def user(user_message: str, history: list[list[str]]):
+        return "", [*history, [user_message, None]]
+
+    def bot(history: list[list[str]], alpha: float, topk: int):
+        user_message = history[-1][0]
         bot_message = generate(
             tokenizer,
             model,
@@ -70,10 +72,15 @@ with gr.Blocks() as demo:
             alpha=alpha,
             topk=topk,
         )
-        return "", [*history, [user_message, bot_message]]
+        history[-1][1] = bot_message
+        return history
 
-    msg.submit(on_message, inputs=[msg, chatbot, alpha, topk], outputs=[msg, chatbot])
-    button.click(on_message, inputs=[msg, chatbot, alpha, topk], outputs=[msg, chatbot])
+    msg.submit(user, inputs=[msg, chatbot], outputs=[msg, chatbot]).then(
+        bot, inputs=[chatbot, alpha, topk], outputs=[chatbot]
+    )
+    button.click(user, inputs=[msg, chatbot], outputs=[msg, chatbot]).then(
+        bot, inputs=[chatbot, alpha, topk], outputs=[chatbot]
+    )
 
     clear.click(lambda: None, None, chatbot)