Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # 模型設定（繁體中文 GPT2）
 model_name = "ckiplab/gpt2-base-chinese"
@@ -12,12 +13,15 @@ system_prompt = (
 )
 
 def reply_fn(chat_history, user_msg):
-    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")
-    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()
-    chat_history.append((user_msg, answer))
+    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()
+        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:
@@ -25,7 +29,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     # 🏫 總務處職安防災 AI 小幫手
     歡迎使用！我是總務處的 AI 助理，任何問題我都會從 **職業安全衛生** 和 **校園防災** 的角度給你正確的建議 👷‍♂️🚒
     """)
-    
+
     chatbot = gr.Chatbot(show_copy_button=True)
     msg = gr.Textbox(placeholder="請輸入你的問題...", label="學生提問")
     clear = gr.Button("清除對話")