Update app.py

app.py

@@ -1,35 +1,60 @@
 import os
+import re
 from fastapi import FastAPI
-from transformers import pipeline
+from transformers import AutoTokenizer, T5ForConditionalGeneration
 
-# 必须在所有导入前设置环境变量
 os.environ["HF_HOME"] = "/app/.cache/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/huggingface"
 
 app = FastAPI()
 
-# 模型加载（带错误处理）
+# 初始化模型
 try:
-    analyzer = pipeline(
-        "text2text-generation",
-        model="Salesforce/codet5-small",
-        tokenizer="Salesforce/codet5-small"
-    )
+    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
+    model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
 except Exception as e:
     raise RuntimeError(f"模型加载失败: {str(e)}")
 
-@app.post("/analyze")
-async def analyze(code: str):
-    result = analyzer(
-        f"Analyze code vulnerabilities:\n{code}",
-        max_length=512,
-        num_beams=5
-    )[0]['generated_text']
-    return {"result": result}
+def sanitize_code(code: str) -> str:
+    """清洗输入代码"""
+    code = re.sub(r"[<>&\"']", "", code)  # 过滤危险字符
+    return code[:1024]  # 限制输入长度
 
-
-# 添加GET支持，测试模型
 @app.get("/analyze")
-async def analyze_get(code: str):  # 新增GET方法
-    return await analyze(code)
-    
+async def analyze_get(code: str):
+    try:
+        # 清洗输入
+        code = sanitize_code(code)
+        
+        # 构造提示词
+        prompt = f"""Analyze the following code for security vulnerabilities in Chinese.
+重点检查SQL注入、XSS、命令注入、路径遍历等问题。
+按此格式返回：\n[漏洞类型]: [风险描述]\n\n代码:\n{code}"""
+        
+        # Tokenize输入
+        inputs = tokenizer(
+            prompt, 
+            return_tensors="pt", 
+            max_length=512, 
+            truncation=True,
+            padding="max_length"
+        )
+        
+        # 生成分析结果
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=512,
+            num_beams=5,
+            early_stopping=True,
+            temperature=0.7
+        )
+        
+        # 解码结果
+        analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"result": analysis}
+    
+    except Exception as e:
+        return {"error": str(e)}
+
+@app.get("/")
+async def health_check():
+    return {"status": "active"}