Update app.py

app.py

@@ -2,36 +2,54 @@ from fastapi import FastAPI
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import os
+import logging
 
-# 1. 基础配置
-app = FastAPI()
+# 初始化日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("CodeSecurityAPI")
 
-# 2. 强制设置缓存路径（解决权限问题）
+# 强制设置缓存路径（解决权限问题）
 os.environ["HF_HOME"] = "/app/.cache/huggingface"
 
-# 3. 加载模型（自动缓存到指定路径）
+# 加载模型
 try:
-    model = AutoModelForSequenceClassification.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
-    tokenizer = AutoTokenizer.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
+    logger.info("Loading model...")
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "mrm8488/codebert-base-finetuned-detect-insecure-code",
+        cache_dir=os.getenv("HF_HOME")
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "mrm8488/codebert-base-finetuned-detect-insecure-code",
+        cache_dir=os.getenv("HF_HOME")
+    )
+    logger.info("Model loaded successfully")
 except Exception as e:
-    raise RuntimeError(f"模型加载失败: {str(e)}")
+    logger.error(f"Model load failed: {str(e)}")
+    raise RuntimeError("模型加载失败，请检查网络连接或模型路径")
+
+app = FastAPI()
 
-# 4. 接口定义
 @app.post("/detect")
 async def detect(code: str):
     try:
-        # 简单处理超长输入
-        if len(code) > 2000:
-            code = code[:2000]
-            
-        inputs = tokenizer(code, return_tensors="pt", truncation=True)
+        # 输入处理（限制长度）
+        code = code[:2000]  # 截断超长输入
+        
+        # 模型推理
+        inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512)
         with torch.no_grad():
             outputs = model(**inputs)
-            
+        
+        # 解析结果
+        label_id = outputs.logits.argmax().item()
         return {
-            "label": model.config.id2label[outputs.logits.argmax().item()],
-            "score": outputs.logits.softmax(dim=-1).max().item()
+            "label": model.config.id2label[label_id],
+            "score": outputs.logits.softmax(dim=-1)[0][label_id].item()
         }
         
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": str(e)}
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}