Update app.py

app.py

@@ -1,31 +1,14 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 import os
 import logging
-from pathlib import Path
 
-# === 初始化日志 ===
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# === 检查缓存目录权限 ===
-def check_permissions():
-    cache_path = Path(os.getenv("HF_HOME", ""))
-    try:
-        cache_path.mkdir(parents=True, exist_ok=True)
-        test_file = cache_path / "permission_test.txt"
-        test_file.write_text("test")
-        test_file.unlink()
-        logger.info(f"✅ 缓存目录权限正常: {cache_path}")
-    except Exception as e:
-        logger.error(f"❌ 缓存目录权限异常: {str(e)}")
-        raise RuntimeError(f"Directory permission error: {str(e)}")
+# === 初始化配置 ===
+app = FastAPI(title="Code Security API")
 
-check_permissions()
-
-# === FastAPI 配置 ===
-app = FastAPI()
+# 解决跨域问题
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -33,23 +16,70 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# === 强制设置缓存路径 ===
+os.environ["HF_HOME"] = "/app/.cache/huggingface"
+cache_path = os.getenv("HF_HOME")
+os.makedirs(cache_path, exist_ok=True)
+
+# === 日志配置 ===
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("CodeBERT-API")
+
+# === 根路径路由（必须定义）===
+@app.get("/")
+async def read_root():
+    """健康检查端点"""
+    return {
+        "status": "running",
+        "endpoints": {
+            "detect": "POST /detect - 代码安全检测",
+            "specs": "GET /openapi.json - API文档"
+        }
+    }
+
 # === 模型加载 ===
 try:
-    logger.info("🔄 加载模型中...")
-    model = AutoModelForSequenceClassification.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
-    tokenizer = AutoTokenizer.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
-    logger.info("✅ 模型加载成功")
+    logger.info("Loading model from: %s", cache_path)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "mrm8488/codebert-base-finetuned-detect-insecure-code",
+        cache_dir=cache_path
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "mrm8488/codebert-base-finetuned-detect-insecure-code",
+        cache_dir=cache_path
+    )
+    logger.info("Model loaded successfully")
 except Exception as e:
-    logger.error(f"❌ 模型加载失败: {str(e)}")
-    raise
+    logger.error("Model load failed: %s", str(e))
+    raise RuntimeError("模型初始化失败")
 
-# === API 接口 ===
+# === 核心检测接口 ===
 @app.post("/detect")
-async def detect(code: str):
-    inputs = tokenizer(code[:2000], return_tensors="pt", truncation=True, max_length=512)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return {
-        "label": int(outputs.logits.argmax()),
-        "score": outputs.logits.softmax(dim=-1).max().item()
-    }
+async def detect_vulnerability(code: str):
+    """代码安全检测主接口"""
+    try:
+        # 输入处理
+        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": label_id,  # 0:安全 1:不安全
+            "confidence": outputs.logits.softmax(dim=-1)[0][label_id].item()
+        }
+        
+    except Exception as e:
+        return {
+            "error": str(e),
+            "tip": "请检查输入代码是否包含非ASCII字符"
+        }