Update app.py

app.py

@@ -2,38 +2,14 @@ import os
 from fastapi import FastAPI
 from transformers import AutoTokenizer, T5ForConditionalGeneration
 
-# 设置缓存路径(必须放在最前面)
-os.environ["HF_HOME"] = "/app/.cache/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/huggingface"
+os.environ["HF_HOME"] = "/app/.cache"
 
 app = FastAPI()
-
-# 加载模型
-try:
-    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
-    model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
-except Exception as e:
-    print(f"模型加载失败: {str(e)}")
-    raise
+tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
+model = T5ForConditionalGeneration.from_pretrained("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)
-    
-    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"]}
+async def analyze(code: str):
+    inputs = tokenizer(f"Analyze vulnerabilities:\n{code}", return_tensors="pt", max_length=512, truncation=True)
+    outputs = model.generate(inputs.input_ids, max_length=512)
+    return {"result": tokenizer.decode(outputs[0], skip_special_tokens=True)}