AquilaX-AI/Review

Inference

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import time
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = AutoModelForSequenceClassification.from_pretrained("AquilaX-AI/Review").to(device)
tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/Review")

partial_code = "if (userInput.length > 255) { return; }"  # Example snippet of insecure code
cwe_id = "CWE-22"  # Example CWE ID for Path Traversal
cwe_name = "Improper Limitation of a Pathname to a Restricted Directory"  # Example CWE Name
affected_line = "42"  # Example line number in the code file
file_name = "utils/inputValidator.js"  # Example file name
org_id = "12345"  # Example organization ID

start = time.time()

prompt = f"""partial_code: {partial_code} , cwe_id: {cwe_id} , cwe_name: {cwe_name}, affected_line: {affected_line},file_name: {file_name}, org_id: {org_id}"""
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
    logits = model(**inputs).logits
predicted_class_id = logits.argmax().item()
predicted_class = model.config.id2label[predicted_class_id]

print(predicted_class)
print(time.time() - start)