update

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 import json
 from onnxruntime import InferenceSession
-from transformers import AutoTokenizer
+import numpy as np
 
 # 设置页面配置
 st.set_page_config(page_title="中文垃圾信息分类器", page_icon="🚫", layout="wide")
@@ -13,9 +13,15 @@ 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
+    
+    # 加载配置文件
+    with open("app-x/chinese_spam_classifier_onnx/config.json", "r") as f:
+        config = json.load(f)
+    
+    id2label = config["id2label"]
+    return hf_classifier, onnx_session, tokenizer, id2label
 
-hf_classifier, onnx_session, tokenizer = load_classifiers()
+hf_classifier, onnx_session, tokenizer, id2label = load_classifiers()
 
 st.title("🚫 中文垃圾信息分类器")
 st.write("使用两个模型进行中文文本的垃圾信息分类。")
@@ -23,6 +29,24 @@ st.write("使用两个模型进行中文文本的垃圾信息分类。")
 # 创建两列布局
 col1, col2 = st.columns([2, 1])
 
+def classify_text(text):
+    # HuggingFace模型分类
+    hf_result = hf_classifier(text)[0]
+    hf_label = id2label[str(int(hf_result["label"].split("_")[-1]))]
+    hf_confidence = hf_result["score"]
+    
+    # ONNX模型分类
+    inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True)
+    onnx_result = onnx_session.run(None, dict(inputs))
+    onnx_probs = onnx_result[0][0]
+    onnx_label = id2label[str(np.argmax(onnx_probs))]
+    onnx_confidence = np.max(onnx_probs)
+    
+    return {
+        "hf": {"label": hf_label, "confidence": hf_confidence},
+        "onnx": {"label": onnx_label, "confidence": float(onnx_confidence)}
+    }
+
 with col1:
     # 创建文本输入框
     text_input = st.text_area("请输入中文文本:", height=200)
@@ -30,52 +54,17 @@ with col1:
     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))
+                results = classify_text(text_input)
                 
-                # 显示JSON格式的结果
-                st.subheader("JSON 格式的详细结果:")
-                st.json(json_result)
+                for model, result in results.items():
+                    st.subheader(f"{model.upper()} 模型分类结果:")
+                    label = "垃圾信息" if result["label"] == "spam" else "正常信息"
+                    if label == "垃圾信息":
+                        st.error(f"⚠️ {label}")
+                    else:
+                        st.success(f"✅ {label}")
+                    st.write(f"概率: {result['confidence']:.2f}")
+                    st.progress(result['confidence'])
         else:
             st.warning("请输入文本后再进行分类。")
 
@@ -104,4 +93,4 @@ with col2:
 
 # 添加页脚
 st.markdown("---")
-st.markdown("由 Streamlit 和 Hugging Face 提供支持 | 作者：[app-x]")
+st.markdown("由 Streamlit 和 Hugging Face 提供支持 | 作者：[app-x]")