Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Body
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
@@ -55,10 +55,16 @@ except Exception as e:
 
 # === 核心检测接口 ===
 @app.post("/detect")
-async def detect_vulnerability(code: str):
+async def detect_vulnerability(payload: dict = Body(...)):
     """代码安全检测主接口"""
     try:
-        # 输入处理
+        # 获取 JSON 输入数据
+        code = payload.get("code", "").strip()
+
+        if not code:
+            return {"error": "代码内容为空", "tip": "请提供有效的代码字符串"}
+
+        # 限制代码长度
         code = code[:2000]  # 截断超长输入
         
         # 模型推理
@@ -66,20 +72,28 @@ async def detect_vulnerability(code: str):
             code,
             return_tensors="pt",
             truncation=True,
+            padding=True,  # 自动选择填充策略
             max_length=512
         )
+
         with torch.no_grad():
             outputs = model(**inputs)
-        
+
         # 结果解析
-        label_id = outputs.logits.argmax().item()
+        logits = outputs.logits
+        label_id = logits.argmax().item()
+        confidence = logits.softmax(dim=-1)[0][label_id].item()
+
+        logger.info(f"Code analyzed. Logits: {logits.tolist()}, Prediction: {label_id}, Confidence: {confidence:.4f}")
+
         return {
             "label": label_id,  # 0:安全 1:不安全
-            "confidence": outputs.logits.softmax(dim=-1)[0][label_id].item()
+            "confidence": round(confidence, 4)
         }
         
     except Exception as e:
+        logger.error("Error during model inference: %s", str(e))
         return {
             "error": str(e),
-            "tip": "请检查输入代码是否包含非ASCII字符"
-        }
+            "tip": "请检查输入代码是否包含非ASCII字符或格式错误"
+        }