Update app.py

app.py

@@ -1,15 +1,28 @@
 import os
 from fastapi import FastAPI
-from transformers import AutoTokenizer, T5ForConditionalGeneration
+from transformers import pipeline
 
-os.environ["HF_HOME"] = "/app/.cache"
+# 必须在所有导入前设置环境变量
+os.environ["HF_HOME"] = "/app/.cache/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/huggingface"
 
 app = FastAPI()
-tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
-model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
+
+# 模型加载（带错误处理）
+try:
+    analyzer = pipeline(
+        "text2text-generation",
+        model="Salesforce/codet5-small",
+        tokenizer="Salesforce/codet5-small"
+    )
+except Exception as e:
+    raise RuntimeError(f"模型加载失败: {str(e)}")
 
 @app.post("/analyze")
 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)}
+    result = analyzer(
+        f"Analyze code vulnerabilities:\n{code}",
+        max_length=512,
+        num_beams=5
+    )[0]['generated_text']
+    return {"result": result}