Update app.py

app.py

@@ -1,54 +1,63 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
-import torch
+import math
 
-app = FastAPI(title="双模型综合文本审计 API")
+app = FastAPI(title="MODERATION_API_V2")
 
-# 1. 初始化两个模型
-print("正在加载模型 A (多语言通用型)...")
-model_a = pipeline("text-classification", model="textdetox/bert-multilingual-toxicity-classifier")
+# 初始化模型
+# Model A: 多语言通用毒性检测
+pipe_a = pipeline("text-classification", model="textdetox/bert-multilingual-toxicity-classifier")
+# Model B: 中文专门化攻击性检测
+pipe_b = pipeline("text-classification", model="thu-coai/roberta-base-cold")
 
-print("正在加载模型 B (中文强化型)...")
-model_b = pipeline("text-classification", model="thu-coai/roberta-base-cold")
-
-class TextInput(BaseModel):
+class CheckRequest(BaseModel):
     text: str
 
-@app.post("/check")
-async def combined_check(input_data: TextInput):
-    text = input_data.text
+@app.post("/analyze")
+async def analyze(request: CheckRequest):
+    text = request.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 {"STATUS": "ERROR", "REASON": "EMPTY_TEXT"}
+
+    # 推理并提取风险概率 (0.0 - 1.0)
+    # 针对 textdetox: LABEL_1 为有害
+    res_a = pipe_a(text)[0]
+    risk_a = res_a['score'] if res_a['label'] == 'LABEL_1' else 1 - res_a['score']
+
+    # 针对 thu-coai: LABEL_1 为攻击性
+    res_b = pipe_b(text)[0]
+    risk_b = res_b['score'] if res_b['label'] == 'LABEL_1' else 1 - res_b['score']
+
+    # 综合风险评分逻辑 (加权计算)
+    # 权重分配：Max_Score(70%) + Avg_Score(30%)
+    combined_risk = (max(risk_a, risk_b) * 0.7) + (((risk_a + risk_b) / 2) * 0.3)
+
+    # 计算数字风险等级 (1-5 级)
+    # Level 1: [0.0-0.2] SAFE
+    # Level 5: [0.8-1.0] BLOCKED
+    risk_level = math.ceil(combined_risk * 5)
+    risk_level = max(1, min(5, risk_level))
+
+    # 状态映射
+    if risk_level >= 4:
+        status = "BLOCKED"
+    elif risk_level == 3:
+        status = "REVIEW"
+    else:
+        status = "PASSED"
 
     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)
+        "TEXT": text,
+        "STATUS": status,
+        "RISK_LEVEL": risk_level,
+        "CONFIDENCE_SCORE": round(combined_risk, 4),
+        "RAW_DATA": {
+            "GENERAL_MODEL": round(risk_a, 4),
+            "SPECIALIZED_MODEL": round(risk_b, 4)
         }
     }
+
+@app.get("/health")
+async def health():
+    return {"STATUS": "UP"}