import streamlit as st
from transformers import pipeline
import json
from onnxruntime import InferenceSession
from transformers import AutoTokenizer

# 设置页面配置
st.set_page_config(page_title="中文垃圾信息分类器", page_icon="🚫", layout="wide")

# 加载中文垃圾信息分类器
@st.cache_resource
def load_classifiers():
    hf_classifier = pipeline("text-classification", model="app-x/chinese_spam_classifier")
    onnx_session = InferenceSession("app-x/chinese_spam_classifier_onnx/model_optimized.onnx")
    tokenizer = AutoTokenizer.from_pretrained("app-x/chinese_spam_classifier_onnx")
    return hf_classifier, onnx_session, tokenizer

hf_classifier, onnx_session, tokenizer = load_classifiers()

st.title("🚫 中文垃圾信息分类器")
st.write("使用两个模型进行中文文本的垃圾信息分类。")

# 创建两列布局
col1, col2 = st.columns([2, 1])

with col1:
    # 创建文本输入框
    text_input = st.text_area("请输入中文文本:", height=200)
    
    if st.button("分类", key="classify_button"):
        if text_input:
            with st.spinner("正在分析..."):
                # HuggingFace模型分类
                hf_result = hf_classifier(text_input)[0]
                hf_label = "垃圾信息" if hf_result["label"] == "spam" else "正常信息"
                hf_confidence = hf_result["score"]
                
                # ONNX模型分类
                inputs = tokenizer(text_input, return_tensors="np", padding=True, truncation=True)
                onnx_result = onnx_session.run(None, dict(inputs))
                onnx_label = "垃圾信息" if onnx_result[0][0][1] > onnx_result[0][0][0] else "正常信息"
                onnx_confidence = max(onnx_result[0][0])
                
                # 创建JSON格式的结果
                json_result = {
                    "input_text": text_input,
                    "huggingface_model": {
                        "classification": hf_label,
                        "confidence": hf_confidence,
                        "raw_output": hf_result
                    },
                    "onnx_model": {
                        "classification": onnx_label,
                        "confidence": float(onnx_confidence),
                        "raw_output": onnx_result[0].tolist()
                    }
                }
                
                # 显示结果
                st.subheader("HuggingFace模型分类结果:")
                if hf_label == "垃圾信息":
                    st.error(f"⚠️ {hf_label}")
                else:
                    st.success(f"✅ {hf_label}")
                st.write(f"概率: {hf_confidence:.2f}")
                st.progress(hf_confidence)
                
                st.subheader("ONNX模型分类结果:")
                if onnx_label == "垃圾信息":
                    st.error(f"⚠️ {onnx_label}")
                else:
                    st.success(f"✅ {onnx_label}")
                st.write(f"概率: {onnx_confidence:.2f}")
                st.progress(float(onnx_confidence))
                
                # 显示JSON格式的结果
                st.subheader("JSON 格式的详细结果:")
                st.json(json_result)
        else:
            st.warning("请输入文本后再进行分类。")

with col2:
    st.subheader("使用说明")
    st.write("""
    1. 在左侧文本框中输入您想要分类的中文文本。
    2. 点击"分类"按钮。
    3. 系统将使用两个模型分析文本并显示结果。
    4. 结果包括两个模型的分类（垃圾信息或正常信息）、概率和JSON格式的详细输出。
    """)
    
    st.subheader("关于模型")
    st.write("""
    本分类器使用了两个模型：
    1. app-x/chinese_spam_classifier (HuggingFace模型)
    2. app-x/chinese_spam_classifier_onnx (ONNX模型)
    这两个模型都基于大规模中文数据集训练，能够有效识别各种类型的垃圾信息。
    """)
    
    st.subheader("免责声明")
    st.info("""
    此分类器仅作为辅助工具，不应完全依赖其结果。
    请始终保持警惕，谨慎处理可疑信息。
    """)

# 添加页脚
st.markdown("---")
st.markdown("由 Streamlit 和 Hugging Face 提供支持 | 作者：[app-x]")