Update app.py

app.py

@@ -1,18 +1,54 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import pipeline
+import torch
 
-app = FastAPI()
+app = FastAPI(title="双模型综合文本审计 API")
 
-# 加载预训练的文本分类模型（建议替换为你微调后的 BERT 模型路径）
-# 例如：'shibing624/bert4confusing-base-chinese' 或自定义模型
-classifier = pipeline("text-classification", model="bert-base-chinese")
+# 1. 初始化两个模型
+print("正在加载模型 A (多语言通用型)...")
+model_a = pipeline("text-classification", model="textdetox/bert-multilingual-toxicity-classifier")
 
-@app.get("/")
-def home():
-    return {"message": "Text Moderation API is running."}
+print("正在加载模型 B (中文强化型)...")
+model_b = pipeline("text-classification", model="thu-coai/roberta-base-cold")
 
-@app.post("/analyze")
-def analyze_text(text: str):
-    # 执行推理
-    result = classifier(text)
-    return {"result": result}
+class TextInput(BaseModel):
+    text: str
+
+@app.post("/check")
+async def combined_check(input_data: TextInput):
+    text = input_data.text
+    if not text.strip():
+        return {"error": "Empty text"}
+
+    # 运行模型 A
+    res_a = model_a(text)[0]
+    # LABEL_1 在该模型通常代表 Toxic
+    score_a = res_a['score'] if res_a['label'] == 'LABEL_1' else 1 - res_a['score']
+
+    # 运行模型 B
+    res_b = model_b(text)[0]
+    # LABEL_1 在该模型代表 Offensive
+    score_b = res_b['score'] if res_b['label'] == 'LABEL_1' else 1 - res_b['score']
+
+    # 3. 综合评估逻辑
+    # 采用“最高分原则”叠加“平均分补偿”
+    # 公式：Final Score = (Max_Score * 0.7) + (Avg_Score * 0.3)
+    final_risk_score = (max(score_a, score_b) * 0.7) + (((score_a + score_b) / 2) * 0.3)
+
+    # 4. 判定结论
+    status = "通过"
+    if final_risk_score > 0.75:
+        status = "拦截"
+    elif final_risk_score > 0.45:
+        status = "建议人工复审"
+
+    return {
+        "text": text,
+        "combined_risk_score": round(final_risk_score, 4),
+        "status": status,
+        "detail": {
+            "general_toxicity_model": round(score_a, 4),
+            "chinese_offensive_model": round(score_b, 4)
+        }
+    }