import streamlit as st
from transformers import pipeline

# 加载中文垃圾邮件分类器
@st.cache_resource
def load_classifier():
    return pipeline("text-classification", model="app-x/chinese_spam_classifier")

classifier = load_classifier()

st.title("中文垃圾信息分类器")
st.write("使用 app-x/chinese_spam_classifier 模型进行中文文本的垃圾信息分类。")

# 创建文本输入框
text_input = st.text_area("请输入中文文本:", height=150)

if st.button("分类"):
    if text_input:
        # 进行分类
        result = classifier(text_input)[0]
        label = "垃圾信息" if result["label"] == "LABEL_1" else "正常信息"
        confidence = result["score"]
        
        # 显示结果
        st.write(f"分类结果: {label}")
        st.write(f"置信度: {confidence:.2f}")
    else:
        st.warning("请输入文本后再进行分类。")