import gradio as gr
from transformers import pipeline

# 使用 DistilGPT2（更小更快）
generator = pipeline("text-generation", model="distilgpt2")

def chat(message, history):
    result = generator(
        message,
        max_new_tokens=50,   # 輸出限制更小，加快生成
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    reply = result[0]["generated_text"]
    return reply

demo = gr.ChatInterface(fn=chat, title="AI 聊天機器人 (DistilGPT2)")

demo.launch()