|
from flask import Flask, request, jsonify |
|
import torch |
|
from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaConfig |
|
import os |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
def load_model(): |
|
|
|
checkpoint = torch.load("codebert_vulnerability_scorer.pth", map_location=torch.device('cpu')) |
|
config = RobertaConfig.from_dict(checkpoint['config']) |
|
|
|
|
|
model = RobertaForSequenceClassification(config) |
|
model.load_state_dict(checkpoint['model_state_dict']) |
|
model.eval() |
|
return model |
|
|
|
|
|
try: |
|
tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_vulnerability") |
|
model = load_model() |
|
print("Model and tokenizer loaded successfully!") |
|
except Exception as e: |
|
print(f"Error loading model: {str(e)}") |
|
|
|
@app.route("/") |
|
def home(): |
|
return request.url |
|
|
|
@app.route("/predict") |
|
def predict(): |
|
try: |
|
|
|
code = request.args.get("code") |
|
if not code: |
|
return jsonify({"error": "Missing 'code' URL parameter"}), 400 |
|
|
|
|
|
inputs = tokenizer( |
|
code, |
|
truncation=True, |
|
padding='max_length', |
|
max_length=512, |
|
return_tensors='pt' |
|
) |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
|
|
|
|
score = torch.sigmoid(outputs.logits).item() |
|
|
|
return jsonify({ |
|
"vulnerability_score": round(score, 4), |
|
"processed_code": code[:500] + "..." if len(code) > 500 else code |
|
}) |
|
|
|
except Exception as e: |
|
return jsonify({"error": str(e)}), 500 |
|
|
|
if __name__ == "__main__": |
|
app.run(host="0.0.0.0", port=7860) |