Update app.py

app.py

@@ -1,43 +1,37 @@
 from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import pipeline
+from puggingface import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 import os
 
-import os
-from pathlib import Path
+# 1. 基础配置
+app = FastAPI()
 
+# 2. 强制设置缓存路径（解决权限问题）
+os.environ["HF_HOME"] = "/app/.cache/huggingface"
 
-# 验证缓存目录可写
-cache_dir = Path(os.getenv("HF_HOME", ""))
-if not cache_dir.exists():
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    
-test_file = cache_dir / "permission_test.txt"
+# 3. 加载模型（自动缓存到指定路径）
 try:
-    with open(test_file, "w") as f:
-        f.write("test")
-    os.remove(test_file)
-    print("✅ Cache directory is writable")
+    model = AutoModelForSequenceClassification.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
+    tokenizer = AutoTokenizer.from_pretrained("mrm8488/codebert-base-finetuned-detect-insecure-code")
 except Exception as e:
-    print(f"❌ Cache directory write failed: {str(e)}")
-    raise
-
-# 正确加载模型（从缓存或下载）
-classifier = pipeline(
-    "text-classification",
-    model="mrm8488/codebert-base-finetuned-detect-insecure-code"
-)
-
-app = FastAPI()
-
-class CodeRequest(BaseModel):
-    code: str  # 输入参数定义
+    raise RuntimeError(f"模型加载失败: {str(e)}")
 
+# 4. 接口定义
 @app.post("/detect")
-async def detect_insecure_code(request: CodeRequest):
+async def detect(code: str):
     try:
-        # 直接传递代码字符串到分类器
-        result = classifier(request.code)
-        return {"status": "success", "result": result[0]}
+        # 简单处理超长输入
+        if len(code) > 2000:
+            code = code[:2000]
+            
+        inputs = tokenizer(code, return_tensors="pt", truncation=True)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            
+        return {
+            "label": model.config.id2label[outputs.logits.argmax().item()],
+            "score": outputs.logits.softmax(dim=-1).max().item()
+        }
+        
     except Exception as e:
-        return {"status": "error", "message": str(e)}
+        return {"error": str(e)}