Update app.py

app.py

@@ -1,10 +1,9 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline
 
+# 模型與角色設定
 model_name = "thu-coai/CDial-GPT_LCCC-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+chatbot = pipeline("text-generation", model=model_name, tokenizer=model_name)
 
 system_prompt = (
     "你是一位台灣公立高中總務處工作人員，使用繁體中文。你非常重視校園的職業安全衛生與防災安全議題，"
@@ -12,12 +11,10 @@ system_prompt = (
 )
 
 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:
-        full_prompt = system_prompt + ''.join([f"學生：{msg[0]}\n你：{msg[1]}\n" for msg in chat_history]) + f"學生：{user_msg}\n你："
-        inputs = tokenizer.encode(full_prompt, return_tensors="pt", truncation=True, max_length=512)
-        outputs = model.generate(inputs, max_new_tokens=100, do_sample=True, temperature=0.7)
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        answer = response.split("你：")[-1].strip()
+        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)}"))
@@ -29,13 +26,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     歡迎使用！我是總務處的 AI 助理，任何問題我都會從 **職業安全衛生** 和 **校園防災** 的角度給你正確的建議 👷‍♂️🚒
     """)
 
-    chatbot = gr.Chatbot(show_copy_button=True)
+    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, msg])
-    clear.click(lambda: ([], ""), None, [chatbot, msg, state])
+    msg.submit(reply_fn, [state, msg], [chatbot_ui, msg])
+    clear.click(lambda: ([], ""), None, [chatbot_ui, msg, state])
 
 demo.launch()