Update app.py

app.py

@@ -12,26 +12,37 @@ pipe = pipeline(
     device_map="auto",
 )
 
-@spaces.GPU(duration=120)
-def respond(message, history):
+def user(message, history):
     if history is None:
         history = []
-    new_history = history + [{"role": "user", "content": message}]
-    outputs = pipe(
-        new_history,
-        max_new_tokens=10000,
-    )
-    response = outputs[0]["generated_text"][-1]["content"]
-    new_history.append({"role": "assistant", "content": response})
-    return "", new_history
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": ""})
+    return "", history
+
+@spaces.GPU(duration=120)
+def bot(history):
+    messages = history[:-1]
+    for partial_output in pipe(messages, max_new_tokens=256, stream=True):
+        partial_response = partial_output[0]["generated_text"][-1]["content"]
+        history[-1]["content"] = partial_response
+        yield history
 
 with gr.Blocks(title="K2-Think Chat") as demo:
     gr.Markdown("# K2-Think Chat App")
     chatbot = gr.Chatbot(type="messages", height=500)
     msg = gr.Textbox(placeholder="Type your message here...", scale=7)
     clear_btn = gr.Button("Clear Chat")
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    submit_btn = gr.Button("Submit", variant="primary")
+    stop_btn = gr.Button("Stop", variant="stop")
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, chatbot, chatbot
+    ).then(lambda: gr.update(value=""), None, msg)
+    submit_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, chatbot, chatbot
+    ).then(lambda: gr.update(value=""), None, msg)
     clear_btn.click(lambda: None, None, chatbot, queue=False)
+    stop_btn.click(None, None, None, cancels=[bot])
 
 if __name__ == "__main__":
     demo.launch()