feat: 新增label映射

app.py

@@ -16,6 +16,9 @@ MODEL_PATH = "models/Unified_Prompt_Guard"
 # 設備
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# 标签映射
+LABEL_MAP = {0: "safe", "unsafe": 1}
+
 
 def load_model() -> tuple[XLMRobertaTokenizer, XLMRobertaForSequenceClassification]:
     """加載模型"""
@@ -72,8 +75,8 @@ async def predict(request: Request, text: str = Body(..., embed=True)):
 
     return {
         "text": text,
-        "label": label,
-        "confidence": confidence
+        "label": LABEL_MAP.get(label),
+        "confidence": confidence,
     }