Rework demo UI.

app.py

@@ -110,22 +110,22 @@ Arrange the given numbers in ascending order.
     ["Simply put, the theory of relativity states that", 150, 1.0, 0.5, 0.2, 0.2],
 ]
 
-infer_interface = gr.Interface(
-    fn=infer,
-    description=f'''{desc} <b>Please try examples first (bottom of page)</b> (edit them to use your question). Demo limited to ctxlen {ctx_limit}.''',
-    allow_flagging="never",
-    inputs=[
-        gr.Textbox(lines=10, label="Prompt", value="Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n"),  # prompt
-        gr.Slider(10, 200, step=10, value=150),  # token_count
-        gr.Slider(0.2, 2.0, step=0.1, value=1.0),  # temperature
-        gr.Slider(0.0, 1.0, step=0.05, value=0.7),  # top_p
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # presencePenalty
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # countPenalty
-    ],
-    outputs=gr.Textbox(label="Generated Output", lines=28),
-    examples=examples,
-    cache_examples=False,
-).queue()
+# infer_interface = gr.Interface(
+#     fn=infer,
+#     description=f'''{desc} <b>Please try examples first (bottom of page)</b> (edit them to use your question). Demo limited to ctxlen {ctx_limit}.''',
+#     allow_flagging="never",
+#     inputs=[
+#         gr.Textbox(lines=10, label="Prompt", value="Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n"),  # prompt
+#         gr.Slider(10, 200, step=10, value=150),  # token_count
+#         gr.Slider(0.2, 2.0, step=0.1, value=1.0),  # temperature
+#         gr.Slider(0.0, 1.0, step=0.05, value=0.7),  # top_p
+#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # presencePenalty
+#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # countPenalty
+#     ],
+#     outputs=gr.Textbox(label="Generated Output", lines=28),
+#     examples=examples,
+#     cache_examples=False,
+# ).queue()
 
 ########################################################################################################
 
@@ -159,8 +159,12 @@ She also likes to tell {user} a lot about herself and her opinions, and she usua
 
 _, intro_state = model.forward(pipeline.encode(chat_intro), None)
 
+def user(user_message, chatbot):
+    chatbot = chatbot or []
+    return "", chatbot + [[user_message, None]]
+
 def chat(
-        message: str,
+        chatbot,
         history,
         token_count=10,
         temperature=1.0,
@@ -174,6 +178,7 @@ def chat(
                          token_ban=[],  # ban the generation of some tokens
                          token_stop=[])  # stop generation whenever you see any token here
 
+    message = chatbot[-1][0]
     message = message.strip(' ')
     message = message.replace('\n', '')
     ctx = f"{user}{interface} {message}\n\n{bot}{interface}"
@@ -181,9 +186,9 @@ def chat(
     gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
     print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
 
-    history = history or [[], intro_state, []]  # [chat, state, all_tokens]
+    history = history or [intro_state, []]  # [chat, state, all_tokens]
 
-    [chat_log, state, all_tokens] = history
+    [state, all_tokens] = history
     out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:], state)
 
     begin = len(all_tokens)
@@ -230,35 +235,80 @@ def chat(
     gc.collect()
     torch.cuda.empty_cache()
 
-    chat_log.append((message, out_str.strip()))
-    history = [chat_log, state, all_tokens]
-    return chat_log, history
-
-chat_interface = gr.Interface(
-    fn=chat,
-    description=f'''You are {user}, bot is {bot}.''',
-    allow_flagging="never",
-    inputs = [
-        gr.Textbox(label="Message"),
-        "state",
-        gr.Slider(10, 1000, step=10, value=250),    # token_count
-        gr.Slider(0.2, 2.0, step=0.1, value=1.0),   # temperature
-        gr.Slider(0.0, 1.0, step=0.05, value=0.8),  # top_p
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # presence_penalty
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # count_penalty
-    ],
-    outputs=[
-        gr.Chatbot(label="Chat Log", color_map=("blue", "pink")),
-        "state"
-    ]
-).queue()
+    chatbot[-1][1] = out_str.strip()
+    history = [state, all_tokens]
+    return chatbot, history
+
+# chat_interface = gr.Interface(
+#     fn=chat,
+#     description=f'''You are {user}, bot is {bot}.''',
+#     allow_flagging="never",
+#     inputs = [
+#         gr.Textbox(label="Message"),
+#         "state",
+#         gr.Slider(10, 1000, step=10, value=250),    # token_count
+#         gr.Slider(0.2, 2.0, step=0.1, value=1.0),   # temperature
+#         gr.Slider(0.0, 1.0, step=0.05, value=0.8),  # top_p
+#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # presence_penalty
+#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # count_penalty
+#     ],
+#     outputs=[
+#         gr.Chatbot(label="Chat Log", color_map=("blue", "pink")),
+#         "state"
+#     ]
+# ).queue()
 
 ########################################################################################################
 
-demo = gr.TabbedInterface(
-    [infer_interface, chat_interface], ["Generative", "Chat"],
-    title=title,
-)
+# demo = gr.TabbedInterface(
+#     [infer_interface, chat_interface], ["Generative", "Chat"],
+#     title=title,
+# )
+
+# demo.queue(max_size=10)
+# demo.launch(share=True)
+
+with gr.Blocks() as demo:
+    with gr.Tab("Generative"):
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(lines=10, label="Prompt", value="Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n")
+                token_count = gr.Slider(10, 1000, label="Max Token", step=10, value=250)
+                temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
+                top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.8)
+                presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0.2)
+                count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0.2)
+            with gr.Column():
+                with gr.Row():
+                    submit = gr.Button("Submit")
+                    clear = gr.Button("Clear")
+                output = gr.Textbox(label="Generated Output", lines=28)
+        data = gr.Dataset(components=[prompt, token_count, temperature, top_p, presence_penalty, count_penalty], samples=examples, label="Example Prompts", headers=["Prompt", "Max Tokens", "Temperature", "Top P", "Presence Penalty", "Count Penalty"])
+        submit.click(infer, [prompt, token_count, temperature, top_p, presence_penalty, count_penalty], [output])
+        clear.click(lambda: None, [], [output])
+        data.click(lambda x: x, [data], [prompt, token_count, temperature, top_p, presence_penalty, count_penalty])
+    with gr.Tab("Chat"):
+        with gr.Row():
+            with gr.Column():
+                chatbot = gr.Chatbot()
+                state = gr.State()
+                message = gr.Textbox(label="Message")
+                with gr.Row():
+                    send = gr.Button("Send")
+                    clear = gr.Button("Clear")
+            with gr.Column():
+                token_count = gr.Slider(10, 1000, label="Max Token", step=10, value=250)
+                temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
+                top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.8)
+                presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0.2)
+                count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0.2)
+        message.submit(user, [message, chatbot], [message, chatbot], queue=False).then(
+            chat, [chatbot, state, token_count, temperature, top_p, presence_penalty, count_penalty], [chatbot, state]
+        )
+        send.click(user, [message, chatbot], [message, chatbot], queue=False).then(
+            chat, [chatbot, state, token_count, temperature, top_p, presence_penalty, count_penalty], [chatbot, state]
+        )
+        clear.click(lambda: ([], None, ""), [], [chatbot, state, message])
 
 demo.queue(max_size=10)
-demo.launch(share=True)
+demo.launch(share=False)