from transformers import pipeline
import gradio as gr

generator = pipeline("text-generation", model="ckiplab/gpt2-base-chinese", 
                     tokenizer="ckiplab/gpt2-base-chinese")

def chat_fn(message, history):
    history = history or []
    input_text = "\n".join(history + [f"你: {message}", "AI:"])
    output = generator(input_text, max_new_tokens=80, pad_token_id=0)[0]["generated_text"]
    response = output.split("AI:")[-1].strip().split("你:")[0].strip()
    history.append(f"你: {message}")
    history.append(f"AI: {response}")
    messages = [(history[i], history[i+1]) for i in range(0, len(history)-1, 2)]
    return messages, history

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="中文聊天機器人", type="tuples")
    state = gr.State([])
    textbox = gr.Textbox(placeholder="請輸入訊息")

    textbox.submit(chat_fn, [textbox, state], [chatbot, state])
    textbox.submit(lambda: "", None, textbox)

demo.launch()