Spaces:
Sleeping
Sleeping
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") | |
# 加载中文垃圾信息分类器 | |
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]") |