| | --- |
| | library_name: transformers |
| | license: apache-2.0 |
| | --- |
| | |
| | ## Inference |
| | ```python |
| | 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) |
| | ``` |