Update app.py

app.py

@@ -1,42 +1,39 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import os
+from fastapi import FastAPI
 from transformers import AutoTokenizer, T5ForConditionalGeneration
-import torch
 
-app = FastAPI()
+# 设置缓存路径(必须放在最前面)
+os.environ["HF_HOME"] = "/app/.cache/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/huggingface"
 
-# 全局加载模型
-tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
-model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
+app = FastAPI()
 
-class CodeRequest(BaseModel):
-    code: str
-    max_length: int = 512
+# 加载模型
+try:
+    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
+    model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
+except Exception as e:
+    print(f"模型加载失败: {str(e)}")
+    raise
 
-@app.post("/v1/analyze")
-async def analyze_code(request: CodeRequest):
-    try:
-        # 构造提示词
-        prompt = f"Analyze security vulnerabilities in this code:\n{request.code}"
-        
-        # 生成分析结果
-        inputs = tokenizer(prompt, return_tensors="pt", 
-                          max_length=512, truncation=True)
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=request.max_length,
-            num_beams=5,
-            early_stopping=True
-        )
-        
-        # 解码结果
-        analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        return {
-            "status": "success",
-            "analysis": analysis,
-            "model": "Salesforce/codet5-small"
-        }
+@app.post("/analyze")
+async def analyze_code(code: str):
+    prompt = f"Analyze security vulnerabilities:\n{code}"
+    
+    inputs = tokenizer(prompt, return_tensors="pt", 
+                      max_length=512, truncation=True)
     
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    outputs = model.generate(
+        inputs.input_ids,
+        max_length=512,
+        num_beams=5,
+        early_stopping=True
+    )
+    
+    return {
+        "result": tokenizer.decode(outputs[0], skip_special_tokens=True)
+    }
+
+@app.get("/health")
+def health_check():
+    return {"status": "ok", "cache_path": os.environ["HF_HOME"]}