update

app.py

@@ -1,19 +1,24 @@
 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_classifier():
-    return pipeline("text-classification", model="app-x/chinese_spam_classifier")
+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
 
-classifier = load_classifier()
+hf_classifier, onnx_session, tokenizer = load_classifiers()
 
 st.title("🚫 中文垃圾信息分类器")
-st.write("使用 app-x/chinese_spam_classifier 模型进行中文文本的垃圾信息分类。")
+st.write("使用两个模型进行中文文本的垃圾信息分类。")
 
 # 创建两列布局
 col1, col2 = st.columns([2, 1])
@@ -25,28 +30,48 @@ with col1:
     if st.button("分类", key="classify_button"):
         if text_input:
             with st.spinner("正在分析..."):
-                # 进行分类
-                result = classifier(text_input)[0]
-                label = "垃圾信息" if result["label"] == "LABEL_1" else "正常信息"
-                confidence = result["score"]
+                # 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,
-                    "classification": label,
-                    "confidence": confidence,
-                    "raw_output": result
+                    "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("分类结果:")
-                if label == "垃圾信息":
-                    st.error(f"⚠️ {label}")
+                st.subheader("HuggingFace模型分类结果:")
+                if hf_label == "垃圾信息":
+                    st.error(f"⚠️ {hf_label}")
                 else:
-                    st.success(f"✅ {label}")
+                    st.success(f"✅ {hf_label}")
+                st.write(f"概率: {hf_confidence:.2f}")
+                st.progress(hf_confidence)
                 
-                st.write(f"置信度: {confidence:.2f}")
-                st.progress(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 格式的详细结果:")
@@ -59,14 +84,16 @@ with col2:
     st.write("""
     1. 在左侧文本框中输入您想要分类的中文文本。
     2. 点击"分类"按钮。
-    3. 系统将分析文本并显示结果。
-    4. 结果包括分类（垃圾信息或正常信息）、置信度和JSON格式的详细输出。
+    3. 系统将使用两个模型分析文本并显示结果。
+    4. 结果包括两个模型的分类（垃圾信息或正常信息）、概率和JSON格式的详细输出。
     """)
     
     st.subheader("关于模型")
     st.write("""
-    本分类器使用了 app-x/chinese_spam_classifier 模型，
-    该模型基于大规模中文数据集训练，能够有效识别各种类型的垃圾信息。
+    本分类器使用了两个模型：
+    1. app-x/chinese_spam_classifier (HuggingFace模型)
+    2. app-x/chinese_spam_classifier_onnx (ONNX模型)
+    这两个模型都基于大规模中文数据集训练，能够有效识别各种类型的垃圾信息。
     """)
     
     st.subheader("免责声明")