File size: 615 Bytes
e6e9cd7
 
08465c2
 
b025fa1
08465c2
e6e9cd7
b025fa1
 
08465c2
e6e9cd7
b025fa1
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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)}