update (https://screen.yanolja.in/eKUGn6i6wCYX2HBS.png)

app.py

@@ -8,24 +8,18 @@ import gradio as gr
 SUPPORTED_MODELS = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"]
 
 
-def user(state_a, state_b, user_message):
-  if state_a is None and state_b is None:
-    model_pair = sample(SUPPORTED_MODELS, 2)
-    state_a = gradio_web_server.State(model_pair[0])
-    state_b = gradio_web_server.State(model_pair[1])
+def user(user_message):
+  model_pair = sample(SUPPORTED_MODELS, 2)
+  new_state_a = gradio_web_server.State(model_pair[0])
+  new_state_b = gradio_web_server.State(model_pair[1])
 
-  for state in [state_a, state_b]:
+  for state in [new_state_a, new_state_b]:
     state.conv.append_message(state.conv.roles[0], user_message)
     state.conv.append_message(state.conv.roles[1], None)
     state.skip_next = False
 
-  empty_prompt = ""
-
   return [
-      state_a, state_b,
-      state_a.to_gradio_chatbot(),
-      state_b.to_gradio_chatbot(), state_a.model_name, state_b.model_name,
-      empty_prompt
+      new_state_a, new_state_b, new_state_a.model_name, new_state_b.model_name
   ]
 
 
@@ -33,8 +27,10 @@ def bot(state_a, state_b, request: gr.Request):
   if state_a is None or state_b is None:
     raise RuntimeError(f"states cannot be None, got [{state_a}, {state_b}]")
 
+  new_states = [state_a, state_b]
+
   generators = []
-  for state in [state_a, state_b]:
+  for state in new_states:
     try:
       # TODO(#1): Allow user to set configuration.
       # bot_response returns a generator yielding states and chatbots.
@@ -50,56 +46,52 @@ def bot(state_a, state_b, request: gr.Request):
       print(f"Error in bot_response: {e}")
       raise e
 
-  new_chatbots = [None, None]
+  new_responses = [None, None]
+
+  # It simulates concurrent response generation from two models.
   while True:
     stop = True
 
     for i in range(2):
       try:
         generator = next(generators[i])
-        states[i], new_chatbots[i] = generator[0], generator[1]
+        new_state = generator[0]
+        new_states[i] = new_state
+        # conv.messages is a list of [role, message].
+        new_responses[i] = new_state.conv.messages[-1][-1]
         stop = False
       except StopIteration:
         pass
 
-    yield [state_a, state_b] + new_chatbots
+    yield new_states + new_responses
 
     if stop:
       break
 
 
 with gr.Blocks() as app:
+  model_names = [gr.State(None), gr.State(None)]
+  responses = [gr.State(None), gr.State(None)]
+
+  # states stores FastChat-specific conversation states.
+  states = [gr.State(None), gr.State(None)]
+
+  prompt = gr.TextArea(label="Prompt", lines=4)
+  submit = gr.Button()
+
   with gr.Row():
-    response_type = gr.Radio(
-        ["Summarization", "Translation"],
-        value="Summarization",
-        label="Response type",
-        info="Choose the type of response you want from the model.")
-    language = gr.Dropdown(["Korean", "English"],
-                           value="Korean",
-                           label="Language",
-                           info="Choose the target language.")
-
-  chatbots = [None, None]
-  with gr.Row():
-    chatbots[0] = gr.Chatbot(label="Model A")
-    chatbots[1] = gr.Chatbot(label="Model B")
+    responses[0] = gr.Textbox(label="Model A", interactive=False)
+    responses[1] = gr.Textbox(label="Model B", interactive=False)
 
-  model_names = [None, None]
   with gr.Accordion("Show models", open=False):
     with gr.Row():
       model_names[0] = gr.Textbox(label="Model A", interactive=False)
       model_names[1] = gr.Textbox(label="Model B", interactive=False)
 
-  prompt = gr.Textbox(label="Prompt")
-
-  states = [gr.State(None), gr.State(None)]
-  prompt.submit(user,
-                states + [prompt],
-                states + chatbots + model_names + [prompt],
-                queue=False).then(bot, states, states + chatbots)
+  submit.click(user, prompt, states + model_names,
+               queue=False).then(bot, states, states + responses)
 
 if __name__ == "__main__":
   # We need to enable queue to use generators.
   app.queue()
-  app.launch()
+  app.launch(debug=True)