Spaces:
Sleeping
Sleeping
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") | |
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)} |