Spaces:
Running
Running
| from transformers import RobertaTokenizer, T5Config, T5EncoderModel | |
| from statement_t5 import StatementT5 | |
| import torch | |
| import pickle | |
| import numpy as np | |
| import onnxruntime | |
| def to_numpy(tensor): | |
| """ get np input for onnx runtime model """ | |
| return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() | |
| def predict_vul_lines(code: list, gpu: bool = False) -> dict: | |
| """Generate statement-level and function-level vulnerability prediction probabilities. | |
| Parameters | |
| ---------- | |
| code : :obj:`list` | |
| A list of String functions. | |
| gpu : bool | |
| Defines if CUDA inference is enabled | |
| Returns | |
| ------- | |
| :obj:`dict` | |
| A dictionary with two keys, "batch_vul_pred", "batch_vul_pred_prob", and "batch_line_scores" | |
| "batch_func_pred" stores a list of function-level vulnerability prediction: [0, 1, ...] where 0 means non-vulnerable and 1 means vulnerable | |
| "batch_func_pred_prob" stores a list of function-level vulnerability prediction probabilities [0.89, 0.75, ...] corresponding to "batch_func_pred" | |
| "batch_statement_pred" stores a list of statement-level vulnerability prediction: [0, 1, ...] where 0 means non-vulnerable and 1 means vulnerable | |
| "batch_statement_pred_prob" stores a list of statement-level vulnerability prediction probabilities [0.89, 0.75, ...] corresponding to "batch_statement_pred" | |
| """ | |
| MAX_STATEMENTS = 155 | |
| MAX_STATEMENT_LENGTH = 20 | |
| DEVICE = 'cuda' if gpu else 'cpu' | |
| # load tokenizer | |
| tokenizer = RobertaTokenizer.from_pretrained("./utils/statement_t5_tokenizer") | |
| # load model | |
| config = T5Config.from_pretrained("./utils/t5_config.json") | |
| model = T5EncoderModel(config=config) | |
| model = StatementT5(model, tokenizer, device=DEVICE) | |
| output_dir = "./models/statement_t5_model.bin" | |
| model.load_state_dict(torch.load(output_dir, map_location=DEVICE)) | |
| model.to(DEVICE) | |
| model.eval() | |
| input_ids, statement_mask = statement_tokenization(code, MAX_STATEMENTS, MAX_STATEMENT_LENGTH, tokenizer) | |
| with torch.no_grad(): | |
| statement_probs, func_probs = model(input_ids=input_ids, statement_mask=statement_mask) | |
| func_preds = torch.argmax(func_probs, dim=-1) | |
| statement_preds = torch.where(statement_probs>0.5, 1, 0) | |
| return {"batch_func_pred": func_preds, "batch_func_pred_prob": func_probs, | |
| "batch_statement_pred": statement_preds, "batch_statement_pred_prob": statement_probs} | |
| def statement_tokenization(code: list, max_statements: int, max_statement_length: int, tokenizer): | |
| batch_input_ids = [] | |
| batch_statement_mask = [] | |
| for c in code: | |
| source = c.split("\n") | |
| source = [statement for statement in source if statement != ""] | |
| source = source[:max_statements] | |
| padding_statement = [tokenizer.pad_token_id for _ in range(20)] | |
| input_ids = [] | |
| for stat in source: | |
| ids_ = tokenizer.encode(str(stat), | |
| truncation=True, | |
| max_length=max_statement_length, | |
| padding='max_length', | |
| add_special_tokens=False) | |
| input_ids.append(ids_) | |
| if len(input_ids) < max_statements: | |
| for _ in range(max_statements-len(input_ids)): | |
| input_ids.append(padding_statement) | |
| statement_mask = [] | |
| for statement in input_ids: | |
| if statement == padding_statement: | |
| statement_mask.append(0) | |
| else: | |
| statement_mask.append(1) | |
| batch_input_ids.append(input_ids) | |
| batch_statement_mask.append(statement_mask) | |
| return torch.tensor(batch_input_ids), torch.tensor(batch_statement_mask) | |
| def predict_cweid(code: list, gpu: bool = False) -> dict: | |
| """Generate CWE-IDs and CWE Abstract Types Predictions. | |
| Parameters | |
| ---------- | |
| code : :obj:`list` | |
| A list of String functions. | |
| gpu : bool | |
| Defines if CUDA inference is enabled | |
| Returns | |
| ------- | |
| :obj:`dict` | |
| A dictionary with four keys, "cwe_id", "cwe_id_prob", "cwe_type", "cwe_type_prob" | |
| "cwe_id" stores a list of CWE-ID predictions: [CWE-787, CWE-119, ...] | |
| "cwe_id_prob" stores a list of confidence scores of CWE-ID predictions [0.9, 0.7, ...] | |
| "cwe_type" stores a list of CWE abstract types predictions: ["Base", "Class", ...] | |
| "cwe_type_prob" stores a list of confidence scores of CWE abstract types predictions [0.9, 0.7, ...] | |
| """ | |
| provider = ["CUDAExecutionProvider", "CPUExecutionProvider"] if gpu else ["CPUExecutionProvider"] | |
| with open("./utils/label_map.pkl", "rb") as f: | |
| cwe_id_map, cwe_type_map = pickle.load(f) | |
| # load tokenizer | |
| tokenizer = RobertaTokenizer.from_pretrained("./utils/tokenizer") | |
| tokenizer.add_tokens(["<cls_type>"]) | |
| tokenizer.cls_type_token = "<cls_type>" | |
| model_input = [] | |
| for c in code: | |
| code_tokens = tokenizer.tokenize(str(c))[:512 - 3] | |
| source_tokens = [tokenizer.cls_token] + code_tokens + [tokenizer.cls_type_token] + [tokenizer.sep_token] | |
| input_ids = tokenizer.convert_tokens_to_ids(source_tokens) | |
| padding_length = 512 - len(input_ids) | |
| input_ids += [tokenizer.pad_token_id] * padding_length | |
| model_input.append(input_ids) | |
| device = "cuda" if gpu else "cpu" | |
| model_input = torch.tensor(model_input, device=device) | |
| # onnx runtime session | |
| ort_session = onnxruntime.InferenceSession("./models/cwe_model.onnx", providers=provider) | |
| # compute ONNX Runtime output prediction | |
| ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(model_input)} | |
| cwe_id_prob, cwe_type_prob = ort_session.run(None, ort_inputs) | |
| # batch_cwe_id_pred (1D list with shape of [batch size]): [pred_1, pred_2, ..., pred_n] | |
| batch_cwe_id = np.argmax(cwe_id_prob, axis=-1).tolist() | |
| # map predicted idx back to CWE-ID | |
| batch_cwe_id_pred = [cwe_id_map[str(idx)] for idx in batch_cwe_id] | |
| # batch_cwe_id_pred_prob (1D list with shape of [batch_size]): [prob_1, prob_2, ..., prob_n] | |
| batch_cwe_id_pred_prob = [] | |
| for i in range(len(cwe_id_prob)): | |
| batch_cwe_id_pred_prob.append(cwe_id_prob[i][batch_cwe_id[i]].item()) | |
| # batch_cwe_type_pred (1D list with shape of [batch size]): [pred_1, pred_2, ..., pred_n] | |
| batch_cwe_type = np.argmax(cwe_type_prob, axis=-1).tolist() | |
| # map predicted idx back to CWE-Type | |
| batch_cwe_type_pred = [cwe_type_map[str(idx)] for idx in batch_cwe_type] | |
| # batch_cwe_type_pred_prob (1D list with shape of [batch_size]): [prob_1, prob_2, ..., prob_n] | |
| batch_cwe_type_pred_prob = [] | |
| for i in range(len(cwe_type_prob)): | |
| batch_cwe_type_pred_prob.append(cwe_type_prob[i][batch_cwe_type[i]].item()) | |
| return {"cwe_id": batch_cwe_id_pred, | |
| "cwe_id_prob": batch_cwe_id_pred_prob, | |
| "cwe_type": batch_cwe_type_pred, | |
| "cwe_type_prob": batch_cwe_type_pred_prob} | |
| def predict_sev(code: list, gpu: bool = False) -> dict: | |
| """Generate CVSS severity score predictions. | |
| Parameters | |
| ---------- | |
| code : :obj:`list` | |
| A list of String functions. | |
| gpu : bool | |
| Defines if CUDA inference is enabled | |
| Returns | |
| ------- | |
| :obj:`dict` | |
| A dictionary with two keys, "batch_sev_score", "batch_sev_class" | |
| "batch_sev_score" stores a list of severity score prediction: [1.0, 5.0, 9.0 ...] | |
| "batch_sev_class" stores a list of severity class based on predicted severity score ["Medium", "Critical"...] | |
| """ | |
| provider = ["CUDAExecutionProvider", "CPUExecutionProvider"] if gpu else ["CPUExecutionProvider"] | |
| # load tokenizer | |
| tokenizer = RobertaTokenizer.from_pretrained("./utils/tokenizer") | |
| model_input = tokenizer(code, truncation=True, max_length=512, padding='max_length', | |
| return_tensors="pt").input_ids | |
| # onnx runtime session | |
| ort_session = onnxruntime.InferenceSession("./models/sev_model.onnx", providers=provider) | |
| # compute ONNX Runtime output prediction | |
| ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(model_input)} | |
| cvss_score = ort_session.run(None, ort_inputs) | |
| batch_sev_score = list(cvss_score[0].flatten().tolist()) | |
| batch_sev_class = [] | |
| for i in range(len(batch_sev_score)): | |
| if batch_sev_score[i] == 0: | |
| batch_sev_class.append("None") | |
| elif batch_sev_score[i] < 4: | |
| batch_sev_class.append("Low") | |
| elif batch_sev_score[i] < 7: | |
| batch_sev_class.append("Medium") | |
| elif batch_sev_score[i] < 9: | |
| batch_sev_class.append("High") | |
| else: | |
| batch_sev_class.append("Critical") | |
| return {"batch_sev_score": batch_sev_score, "batch_sev_class": batch_sev_class} | |
| if __name__ == "__main__": | |
| import pandas as pd | |
| df = pd.read_csv("./data/processed_test.csv") | |
| funcs = df["func_before"].tolist() | |
| for code in funcs: | |
| out = predict_vul_lines([code]) | |
| print(out["batch_func_pred"][0]) |