code-audit-api / app.py
Forrest99's picture
Update app.py
b025fa1 verified
raw
history blame
615 Bytes
import os
from fastapi import FastAPI
from transformers import AutoTokenizer, T5ForConditionalGeneration
os.environ["HF_HOME"] = "/app/.cache"
app = FastAPI()
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
model = T5ForConditionalGeneration.from_pretrained("Salesforce/codet5-small")
@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)}