Update app.py

app.py

@@ -1,120 +1,72 @@
+import random
+import time
 import os
-
 import gradio as gr
 from text_generation import Client
-from conversation import get_default_conv_template, SeparatorStyle
+from conversation import get_default_conv_template
 
 
+endpoint_url = os.environ.get("ENDPOINT_URL")
+client = Client(endpoint_url, timeout=120)
 eos_token = "</s>"
 
-def _concat_messages(messages):
-    message_text = ""
-    for message in messages:
-        if message["role"] == "system":
-            message_text += "<|system|>\n" + message["content"].strip() + "\n"
-        elif message["role"] == "user":
-            message_text += "<|user|>\n" + message["content"].strip() + "\n"
-        elif message["role"] == "assistant":
-            message_text += "<|assistant|>\n" + message["content"].strip() + eos_token + "\n"
-        else:
-            raise ValueError("Invalid role: {}".format(message["role"]))
-    return message_text
 
-endpoint_url = os.environ.get("ENDPOINT_URL")
-client = Client(endpoint_url, timeout=120)
 
-def generate_response(user_input, max_new_token, top_p, top_k, temperature, do_sample, repetition_penalty):
-    user_input = user_input.strip()
+def generate_response(history, max_new_token=512, top_p=0.9, temperature=0.8, do_sample=True):
     conv = get_default_conv_template("vicuna").copy()
     roles = {"human": conv.roles[0], "gpt": conv.roles[1]} # map human to USER and gpt to ASSISTANT
-    role = roles["human"]
-    conv.append_message(role, user_input)
-    conv.append_message(roles["gpt"], None)
+    for user, bot in history:
+        conv.append_message(roles['human'], user)
+        conv.append_message(roles["gpt"], bot)
     msg = conv.get_prompt()
 
-    res = client.generate(
-        msg,
-        stop_sequences=["<|assistant|>", eos_token, "<|system|>", "<|user|>"],
-        max_new_tokens=max_new_token,
-        top_p=top_p,
-        top_k=top_k,
-        do_sample=do_sample,
-        temperature=temperature,
-        repetition_penalty=repetition_penalty,
-    )
-    return [("assistant", res.generated_text)]
+    for response in client.generate_stream(
+            msg,
+            max_new_tokens=max_new_token,
+            top_p=top_p,
+            temperature=temperature,
+            do_sample=do_sample,
+    ):
+        if not response.token.special:
+            yield response.token.text
 
+    # res = client.generate(
+    #     msg,
+    #     stop_sequences=["<|assistant|>", eos_token, "<|system|>", "<|user|>"],
+    #     max_new_tokens=max_new_token,
+    #     top_p=top_p,
+    #     top_k=top_k,
+    #     do_sample=do_sample,
+    #     temperature=temperature,
+    #     repetition_penalty=repetition_penalty,
+    # )
+    # return [("assistant", res.generated_text)]
+    #
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot()
-    with gr.Row():
-        with gr.Column(scale=4):
-            with gr.Column(scale=12):
-                user_input = gr.Textbox(
-                    show_label=False,
-                    placeholder="Shift + Enter傳送...",
-                    lines=10).style(
-                    container=False)
-            with gr.Column(min_width=32, scale=1):
-                submitBtn = gr.Button("Submit", variant="primary")
-        with gr.Column(scale=1):
-            emptyBtn = gr.Button("Clear History")
-            max_new_token = gr.Slider(
-                1,
-                1024,
-                value=128,
-                step=1.0,
-                label="Maximum New Token Length",
-                interactive=True)
-            top_p = gr.Slider(0, 1, value=0.9, step=0.01,
-                              label="Top P", interactive=True)
-            temperature = gr.Slider(
-                0,
-                1,
-                value=0.5,
-                step=0.01,
-                label="Temperature",
-                interactive=True)
-            top_k = gr.Slider(1, 40, value=40, step=1,
-                              label="Top K", interactive=True)
-            do_sample = gr.Checkbox(
-                value=True,
-                label="Do Sample",
-                info="use random sample strategy",
-                interactive=True)
-            repetition_penalty = gr.Slider(
-                1.0,
-                3.0,
-                value=1.1,
-                step=0.1,
-                label="Repetition Penalty",
-                interactive=True)
+    msg = gr.Textbox()
+    clear = gr.Button("Clear")
 
-    params = [user_input, chatbot]
-    predict_params = [
-        chatbot,
-        max_new_token,
-        top_p,
-        temperature,
-        top_k,
-        do_sample,
-        repetition_penalty]
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
 
-    submitBtn.click(
-        generate_response,
-        [user_input, max_new_token, top_p, top_k, temperature, do_sample, repetition_penalty],
-        [chatbot],
-        queue=False
-    )
+    def bot(history):
+        # history = list of [[user_message, bot_message], ...]
 
-    user_input.submit(
-        generate_response,
-        [user_input, max_new_token, top_p, top_k, temperature, do_sample, repetition_penalty],
-        [chatbot],
-        queue=False
-    )
+        import ipdb
+        ipdb.set_trace()
+        bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
+        history[-1][1] = ""
+        for character in bot_message:
+            history[-1][1] += character
+            time.sleep(0.05)
+            yield history
 
-    submitBtn.click(lambda: None, [], [user_input])
-
-    emptyBtn.click(lambda: chatbot.reset(), outputs=[chatbot], show_progress=True)
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        generate_response, chatbot, chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+    
+demo.queue()
+demo.launch()
 
-demo.launch()