Radical update with gr.Chatbot instead to actively append history

app.py

@@ -8,7 +8,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 1024
-DEFAULT_MAX_NEW_TOKENS = 512
+DEFAULT_MAX_NEW_TOKENS = 256
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 DESCRIPTION = """\
@@ -86,62 +86,45 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    fill_height=True,
-    additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Slider(
-            label="Repetition penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        ["நான் எப்படி வேகமாக தூங்க முடியும்?"],
-        ["என் முதலாளி மிகவும் கட்டுப்படுத்துகிறார், நான் என்ன செய்ய வேண்டும்?"],
-        ["திருமணத்திற்கு நான் என்ன அணிய வேண்டும்?"],
-        ["வரலாற்றில் தெரிந்து கொள்ள வேண்டிய சில முக்கியமான காலங்கள் யாவை?"],
-        ["நான் பணம் சம்பாதிக்க வேண்டும் ஆனால் வேடிக்கையாக இருக்க வேண்டும் என்றால் நல்ல தொழில் எது?"],
-    ],
-)
-
+examples = [
+    ["நான் எப்படி வேகமாக தூங்க முடியும்?"],
+    ["என் முதலாளி மிகவும் கட்டுப்படுத்துகிறார், நான் என்ன செய்ய வேண்டும்?"],
+    ["திருமணத்திற்கு நான் என்ன அணிய வேண்டும்?"],
+    ["வரலாற்றில் தெரிந்து கொள்ள வேண்டிய சில முக்கியமான காலங்கள் யாவை?"],
+    ["நான் பணம் சம்பாதிக்க வேண்டும் ஆனால் வேடிக்கையாக இருக்க வேண்டும் என்றால் நல்ல தொழில் எது?"],
+]
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
-    chat_interface.render()
+
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Enter your message")
+    clear = gr.Button("Clear")
+
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history, system_prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
+        user_message = history[-1][0]
+        chat_history = [(msg[0], msg[1]) for msg in history[:-1]]
+        bot_message = ""
+        for response in generate(user_message, chat_history, system_prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
+            bot_message = response
+            history[-1][1] = bot_message
+            yield history
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot,
+        [chatbot, gr.Textbox(label="System prompt", lines=6, value=SYSTEM_PROMPT),
+         gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
+         gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
+         gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
+         gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
+         gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)],
+        chatbot,
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
     gr.Markdown(LICENSE)
 
 if __name__ == "__main__":