Update app.py

app.py

@@ -12,42 +12,26 @@ pipe = pipeline(
     device_map="auto",
 )
 
-def user(message, history):
+@spaces.GPU(duration=120)
+def respond(message, history):
     if history is None:
         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
+    new_history = history + [{"role": "user", "content": message}]
+    outputs = pipe(
+        new_history,
+        max_new_tokens=32768,
+    )
+    response = outputs[0]["generated_text"][-1]["content"]
+    new_history.append({"role": "assistant", "content": response})
+    return "", new_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")
-    submit_btn = gr.Button("Submit", variant="primary")
-    stop_btn = gr.Button("Stop", variant="stop")
-
-    # Store the event handler to cancel it later
-    bot_event = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    ).then(lambda: gr.update(value=""), None, msg)
-    
-    submit_event = submit_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    ).then(lambda: gr.update(value=""), None, msg)
-    
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
     clear_btn.click(lambda: None, None, chatbot, queue=False)
-    
-    # Cancel the events instead of trying to cancel the function directly
-    stop_btn.click(None, None, None, cancels=[bot_event, submit_event])
 
 if __name__ == "__main__":
     demo.launch()