import gradio as gr
from transformers import pipeline

# 加载中文垃圾邮件分类器
classifier = pipeline("text-classification", model="app-x/chinese_spam_classifier")

def classify_text(text):
    result = classifier(text)[0]
    label = "垃圾邮件" if result["label"] == "LABEL_1" else "正常邮件"
    confidence = result["score"]
    return f"{label} (置信度: {confidence:.2f})"

# 创建 Gradio 界面
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=5, placeholder="请输入中文文本..."),
    outputs="text",
    title="中文垃圾邮件分类器",
    description="使用 app-x/chinese_spam_classifier 模型进行中文文本的垃圾邮件分类。"
)

# 启动应用
iface.launch()