Forrest99/codebertBase

1from fastapi import FastAPI, Body2from fastapi.middleware.cors import CORSMiddleware3from transformers import AutoTokenizer, AutoModelForSequenceClassification4import torch5import os6import logging7 8# === 初始化配置 ===9app = FastAPI(title="Code Security API")10 11# 解决跨域问题12app.add_middleware(13    CORSMiddleware,14    allow_origins=["*"],15    allow_methods=["*"],16    allow_headers=["*"],17)18 19# === 强制设置缓存路径 ===20os.environ["HF_HOME"] = "/app/.cache/huggingface"21cache_path = os.getenv("HF_HOME")22os.makedirs(cache_path, exist_ok=True)23 24# === 日志配置 ===25logging.basicConfig(level=logging.INFO)26logger = logging.getLogger("CodeBERT-API")27 28# === 根路径路由（必须定义）===29@app.get("/")30async def read_root():31    """健康检查端点"""32    return {33        "status": "running",34        "endpoints": {35            "detect": "POST /detect - 代码安全检测",36            "specs": "GET /openapi.json - API文档"37        }38    }39 40# === 模型加载 ===41try:42    logger.info("Loading model from: %s", cache_path)43    model = AutoModelForSequenceClassification.from_pretrained(44        "mrm8488/codebert-base-finetuned-detect-insecure-code",45        cache_dir=cache_path46    )47    tokenizer = AutoTokenizer.from_pretrained(48        "mrm8488/codebert-base-finetuned-detect-insecure-code",49        cache_dir=cache_path50    )51    logger.info("Model loaded successfully")52except Exception as e:53    logger.error("Model load failed: %s", str(e))54    raise RuntimeError("模型初始化失败")55 56# === 核心检测接口 ===57@app.post("/detect")58async def detect_vulnerability(payload: dict = Body(...)):59    """代码安全检测主接口"""60    try:61        # 获取 JSON 输入数据62        code = payload.get("code", "").strip()63 64        if not code:65            return {"error": "代码内容为空", "tip": "请提供有效的代码字符串"}66 67        # 限制代码长度68        code = code[:2000]  # 截断超长输入69        70        # 模型推理71        inputs = tokenizer(72            code,73            return_tensors="pt",74            truncation=True,75            padding=True,  # 自动选择填充策略76            max_length=51277        )78 79        with torch.no_grad():80            outputs = model(**inputs)81 82        # 结果解析83        logits = outputs.logits84        label_id = logits.argmax().item()85        confidence = logits.softmax(dim=-1)[0][label_id].item()86 87        logger.info(f"Code analyzed. Logits: {logits.tolist()}, Prediction: {label_id}, Confidence: {confidence:.4f}")88 89        return {90            "label": label_id,  # 0:安全 1:不安全91            "confidence": round(confidence, 4)92        }93        94    except Exception as e:95        logger.error("Error during model inference: %s", str(e))96        return {97            "error": str(e),98            "tip": "请检查输入代码是否包含非ASCII字符或格式错误"99        }100