import gradio as gr
from transformers import pipeline

# 模型與角色設定
model_name = "thu-coai/CDial-GPT_LCCC-base"
chatbot = pipeline("text-generation", model=model_name, tokenizer=model_name)

system_prompt = (
    "你是一位台灣公立高中總務處工作人員，使用繁體中文。你非常重視校園的職業安全衛生與防災安全議題，"
    "並且遵守台北市教育局的相關法令。無論學生問你什麼問題，你都會將話題引導到這些議題上。\n"
)

def reply_fn(chat_history, user_msg):
    prompt = system_prompt + ''.join([f"學生：{m[0]}\n你：{m[1]}\n" for m in chat_history]) + f"學生：{user_msg}\n你："
    try:
        result = chatbot(prompt, max_length=300, do_sample=True, temperature=0.7)[0]["generated_text"]
        answer = result.split("你：")[-1].strip()
        chat_history.append((user_msg, answer))
    except Exception as e:
        chat_history.append((user_msg, f"⚠️ 發生錯誤：{str(e)}"))
    return chat_history, ""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("""
    # 🏫 總務處職安防災 AI 小幫手
    歡迎使用！我是總務處的 AI 助理，任何問題我都會從 **職業安全衛生** 和 **校園防災** 的角度給你正確的建議 👷‍♂️🚒
    """)

    chatbot_ui = gr.Chatbot(show_copy_button=True)
    msg = gr.Textbox(placeholder="請輸入你的問題...", label="學生提問")
    clear = gr.Button("清除對話")
    state = gr.State([])

    msg.submit(reply_fn, [state, msg], [chatbot_ui, msg])
    clear.click(lambda: ([], ""), None, [chatbot_ui, msg, state])

demo.launch()