# coding=utf-8 from __future__ import absolute_import, division, print_function import argparse from collections import defaultdict import glob import logging import os # os.environ['CUDA_VISIBLE_DEVICES'] = '0' import random import scipy.sparse as sp import numpy as np import torch torch.cuda.empty_cache() import csv from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler,TensorDataset from torch.utils.data.distributed import DistributedSampler from torch_geometric.data import Data from sklearn.metrics import roc_auc_score import json import math try: from torch.utils.tensorboard import SummaryWriter except: from tensorboardX import SummaryWriter from tqdm import tqdm, trange import multiprocessing from model import * cpu_cont = multiprocessing.cpu_count() from transformers import (WEIGHTS_NAME, get_linear_schedule_with_warmup, BertConfig, BertForMaskedLM, BertTokenizer, GPT2Config, GPT2LMHeadModel, GPT2Tokenizer, OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer, RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer, RobertaModel, DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer) from typing import Optional import string logger = logging.getLogger(__name__) import torch.multiprocessing torch.multiprocessing.set_sharing_strategy('file_system') MODEL_CLASSES = { 'gpt2': (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer), 'openai-gpt': (OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer), 'bert': (BertConfig, BertForMaskedLM, BertTokenizer), 'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer), 'distilbert': (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer) } def warn(*args, **kwargs): pass import warnings warnings.warn = warn class Lexer(object): def __init__(self, code: str): self.code = code def _is_identifier_char(self, c: str) -> bool: if c in string.ascii_letters + string.digits + "_": return True return False def _is_digit(self, c: str) -> bool: if c in string.digits: return True return False def __iter__(self): self.i = 0 return self def _peek(self) -> Optional[str]: if self.i < len(self.code): return self.code[self.i] raise StopIteration def _get(self) -> Optional[str]: if self.i < len(self.code): val = self.code[self.i] self.i += 1 return val raise StopIteration def _identifier(self) -> str: token = self._get() while self._is_identifier_char(self._peek()): token += self._get() return token def _number(self) -> str: token = self._get() while self._is_digit(self._peek()): token += self._get() return token def __next__(self) -> str: while True: c = self._peek() if c is None: break if not c.isspace(): break self._get() c = self._peek() if c in string.ascii_letters: return self._identifier() elif c in string.digits: return self._number() elif c in "()[]{}<>=+-*/#.,:;'\"|": return self._get() else: return self._get() class InputFeatures(object): """A single training/test features for a example.""" def __init__(self, all_ids, edges, edges_label, num_nodes, idx, label, node_target, code_lines, cwe ): # `all_ids` is already a list of token-id sequences for each node. The # trailing comma here turned it into a tuple containing that list which # then required callers to index `[0]` to retrieve the real data. # Store the list directly instead. self.all_ids = all_ids self.edges = edges self.edges_label = edges_label self.num_nodes = num_nodes self.idx = str(idx) self.label = label self.node_target = node_target self.code_lines = code_lines self.cwe = cwe def convert_graphs_to_features(js,tokenizer, args): codes = js['nodes_codes'] edges = js['edges'] edges_label = [] # ─ extract CWE (first “CWE123” in the filename) ──────────────────────── filename = js["filename"] m = re.search(r"(CWE\d+)", filename) cwe_id = m.group(1) if m else "UNKNOWN" edges_label_source = js['edges_label'] for i, label in enumerate(edges_label_source): if len(label) != 0: edges_label.append(0) elif len(label) == 0: node0_code = codes[edges[i][0]] node1_code = codes[edges[i][1]] if (node0_code not in node1_code) and (node1_code not in node0_code): edges_label.append(1) else: edges_label.append(2) codes = codes[:args.block_size] all_ids = [] codes = codes[:args.block_size] for code in codes: code_tokens = tokenizer.tokenize(code) source_tokens = [tokenizer.cls_token]+code_tokens+[tokenizer.sep_token] source_ids = tokenizer.convert_tokens_to_ids(source_tokens) all_ids.append(source_ids) num_nodes = len(codes) if num_nodes == 0: return None return InputFeatures(all_ids, edges, edges_label, num_nodes, js['idx'], js['target'], js['node_target'], js['code_lines'], cwe_id) def convert_codes_to_tokens(js, args): codes = js['nodes_codes'] codes = codes[:args.block_size] all_tokens = [] codes = codes[:args.block_size] for code in codes: code_tokens = code.split(' ') all_tokens.append(code_tokens) return all_tokens # build graph function def build_graph(node_idxs, nodes_edges, num_nodes, edges_label, w_embeddings, args): # `node_idxs` is expected to be a list of token-id sequences for all nodes. # Older versions wrapped this list in an extra container which required # unwrapping here. After fixing `InputFeatures` we receive the list # directly, so we remove the unnecessary indexing. # print('using window size = ', window_size) edges = nodes_edges + [[y, x] for x, y in nodes_edges] edges_label = edges_label + edges_label adj_list = [] for i in range(3): row, col, weight = [], [], [] for j, edge in enumerate(edges): if edge[0] == edge[1]: continue if edge[0] < num_nodes and edge[1] < num_nodes and edges_label[j] == i: row.append(edge[0]) col.append(edge[1]) weight.append(1) # weight.append(edges_label[i]) adj = sp.csr_matrix((weight, (row, col)), shape=(num_nodes, num_nodes)) adj_list.append(adj) features_out = [] for node_idx in node_idxs: # emb_seq = [] # for token in node_idx: # emb_seq.append(w_embeddings[token]) # emb_seq = torch.tensor(emb_seq) emb_np = np.stack([w_embeddings[token] for token in node_idx], axis=0) emb_seq = torch.from_numpy(emb_np).float() features_out.append(emb_seq) return adj_list, features_out def build_adj(nodes_edges, num_nodes): # print('using window size = ', window_size) edges = nodes_edges + [[y, x] for x, y in nodes_edges] row, col, weight = [], [], [] for edge in edges: if edge[0] == edge[1]: continue if edge[0] < num_nodes and edge[1] < num_nodes: row.append(edge[0]) col.append(edge[1]) weight.append(1) adj = sp.csr_matrix((weight, (row, col)), shape=(num_nodes, num_nodes)) return adj def convert_examples_to_features(js,tokenizer,args): code=' '.join(js['func'].split()) code_tokens=tokenizer.tokenize(code)[:args.block_size-2] source_tokens =[tokenizer.cls_token]+code_tokens+[tokenizer.sep_token] source_ids = tokenizer.convert_tokens_to_ids(source_tokens) padding_length = args.block_size - len(source_ids) source_ids+=[tokenizer.pad_token_id]*padding_length return InputFeatures(source_tokens,source_ids,js['idx'],js['target'],js['node_target']) class TextDataset(Dataset): def __init__(self, encoder, tokenizer, args, file_path=None, sample_percent=1.): self.examples = [] self.encoder = encoder.to(args.device) self.args = args emb = encoder.embeddings.word_embeddings.weight.detach().cpu() self.w_embeddings = torch.as_tensor(emb, dtype=torch.float32).share_memory_() # self.w_embeddings = encoder.embeddings.word_embeddings.weight.data.cpu().detach().clone().numpy() with open(file_path) as f: for line in f: js = json.loads(line.strip()) result = convert_graphs_to_features(js, tokenizer, args) if result is None: continue self.examples.append(result) total_len = len(self.examples) num_keep = int(sample_percent * total_len) if num_keep < total_len: np.random.seed(10) np.random.shuffle(self.examples) self.examples = self.examples[:num_keep] def __len__(self): return len(self.examples) # def __getitem__(self, i): # adj, x_feature = build_graph(self.examples[i].all_ids, self.examples[i].edges, self.examples[i].num_nodes, self.examples[i].edges_label, self.w_embeddings, self.args) # adj, adj_mask = preprocess_adj(adj, self.args.block_size) # adj_feature = preprocess_features(x_feature, self.args.block_size) # label_ = self.examples[i].label # return torch.tensor(adj), torch.tensor(adj_mask), torch.tensor(adj_feature), torch.tensor(label_) def __getitem__(self, i): adj_list, x_feature = build_graph( self.examples[i].all_ids, self.examples[i].edges, self.examples[i].num_nodes, self.examples[i].edges_label, self.w_embeddings, self.args ) # --- graph-level pads (already fixed before) --- adj_padded_list, adj_mask = preprocess_adj(adj_list, self.args.block_size) adj_feature = preprocess_features(x_feature, self.args.block_size) # --- NEW: pad node targets & create mask -------------------- node_targets = self.examples[i].node_target # list/1-D array nt_len = len(node_targets) pad_len = self.args.block_size - nt_len if pad_len < 0: # truncate too-long graphs node_targets = node_targets[:self.args.block_size] nt_len = self.args.block_size pad_len = 0 node_targets_padded = np.pad(node_targets, (0, pad_len), "constant") node_mask = np.zeros(self.args.block_size, dtype=np.float32) node_mask[:nt_len] = 1.0 # ------------------------------------------------------------- # convert to tensors once, all with matching shapes adj_tensor = torch.from_numpy(np.stack(adj_padded_list, 0)).float() adj_mask_tensor = torch.from_numpy(adj_mask).float() feat_tensor = torch.from_numpy(adj_feature).float() graph_lbl = torch.tensor(self.examples[i].label, dtype=torch.float) node_lbl = torch.tensor(node_targets_padded, dtype=torch.float) node_msk = torch.tensor(node_mask, dtype=torch.float) lines = self.examples[i].code_lines # <-- grab them lines = lines[:self.args.block_size] pad = self.args.block_size - len(lines) lines = np.pad(lines, (0, pad), "constant", constant_values=-1) line_tensor = torch.tensor(lines, dtype=torch.float) line_tensor = torch.tensor(lines, dtype=torch.float) cwe_id = self.examples[i].cwe # NEW return (adj_tensor, adj_mask_tensor, feat_tensor, graph_lbl, node_lbl, node_msk, line_tensor, cwe_id) def set_seed(seed=42): random.seed(seed) os.environ['PYHTONHASHSEED'] = str(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.backends.cudnn.deterministic = True def train(args, train_dataset, eval_dataset, model, tokenizer): """ Train the model """ args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu) # train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset) # train_sampler = torch.utils.data.WeightedRandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset) if args.local_rank == -1: # single-GPU / single-process # ---- build per-example weights --------------------------------- labels = [ex.label for ex in train_dataset.examples] # 0 / 1 n_pos = sum(labels) n_neg = len(labels) - n_pos # avoid div-by-zero when there are no positives (unlikely) pos_weight = (n_neg / n_pos) if n_pos else 1.0 sample_weights = torch.tensor( [pos_weight if lbl == 1 else 1.0 for lbl in labels], dtype=torch.double ) # store for later so we can reuse in the loss args.pos_weight = pos_weight train_sampler = torch.utils.data.WeightedRandomSampler( weights = sample_weights, num_samples = len(sample_weights), # usually == dataset size replacement = True ) else: # distributed case train_sampler = DistributedSampler(train_dataset) train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, num_workers=4, pin_memory=True) # numworkers=4 args.max_steps=args.epoch*len(train_dataloader) # args.save_steps=len(train_dataloader) args.save_steps = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) args.warmup_steps=len(train_dataloader) args.logging_steps=len(train_dataloader) args.num_train_epochs=args.epoch model.to(args.device) # Prepare optimizer and schedule (linear warmup and decay) no_decay = ['bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay}, {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon) scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.max_steps*0.1, num_training_steps=args.max_steps) if args.fp16: try: from apex import amp except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") # model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level) model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level, loss_scale=1024) # multi-gpu training (should be after apex fp16 initialization) if args.n_gpu > 1: model = torch.nn.DataParallel(model) # Distributed training (should be after apex fp16 initialization) if args.local_rank != -1: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last') scheduler_last = os.path.join(checkpoint_last, 'scheduler.pt') optimizer_last = os.path.join(checkpoint_last, 'optimizer.pt') if os.path.exists(scheduler_last): scheduler.load_state_dict(torch.load(scheduler_last)) if os.path.exists(optimizer_last): optimizer.load_state_dict(torch.load(optimizer_last)) # Train! logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_dataset)) logger.info(" Num Epochs = %d", args.num_train_epochs) logger.info(" Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size) logger.info(" Total train batch size (w. parallel, distributed & accumulation) = %d", args.train_batch_size * args.gradient_accumulation_steps * ( torch.distributed.get_world_size() if args.local_rank != -1 else 1)) logger.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps) logger.info(" Total optimization steps = %d", args.max_steps) global_step = args.start_step tr_loss, logging_loss, avg_loss, tr_nb, tr_num, train_loss = 0.0, 0.0, 0.0, 0, 0, 0 best_mrr = 0.0 best_f1 = 0.0 # model.resize_token_embeddings(len(tokenizer)) model.zero_grad() my_metrics = [] # for idx in range(args.start_epoch, int(args.num_train_epochs)): for idx in trange(args.start_epoch, int(args.num_train_epochs), desc="Epoch"): tr_num = 0 train_loss = 0 # for step, batch in enumerate(bar): # for step, batch in enumerate(train_dataloader): for step, batch in enumerate(tqdm(train_dataloader, desc=f"Training (epoch {idx+1})", leave=False)): # adj = batch[0].to(args.device) # adj_mask = batch[1].to(args.device) # adj_feature = batch[2].to(args.device) # labels = batch[3].to(args.device) # n_label = batch[4].to(args.device) # model.train() # loss, logits = model(adj, adj_mask, adj_feature, labels, n_label) # ignore code_lines during training # adj, adj_mask, adj_feat, g_lbl, n_lbl, n_msk, _ = (t.to(args.device) for t in batch) # first 7 items are tensors, last one is a plain str (cwe) *tensor_parts, cwe_ids = batch adj, adj_mask, adj_feat, g_lbl, n_lbl, n_msk, line_ids = tensor_parts adj = adj.to(args.device) adj_mask = adj_mask.to(args.device) adj_feat = adj_feat.to(args.device) g_lbl = g_lbl.to(args.device) n_lbl = n_lbl.to(args.device) n_msk = n_msk.to(args.device) line_ids = line_ids.to(args.device) model.train() loss, _, _ = model(adj, adj_mask, adj_feat, graph_labels=g_lbl, node_labels =n_lbl, node_mask =n_msk) if args.n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu parallel training if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm) else: loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) tr_loss += loss.item() tr_num+=1 train_loss+=loss.item() if avg_loss==0: avg_loss=tr_loss avg_loss=round(train_loss/tr_num,5) if (step + 1) % args.gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() scheduler.step() global_step += 1 output_flag=True avg_loss=round(np.exp((tr_loss - logging_loss) /(global_step- tr_nb)),4) if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0: logging_loss = tr_loss tr_nb=global_step if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0: if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well results = evaluate(args, eval_dataset, model, tokenizer,eval_when_training=True) my_metric = [] for key, value in results.items(): # logger.info(" %s = %s", key, round(value,4)) if key == "per_cwe_auc": for cwe, auc in value.items(): logger.info(f" {cwe}: {auc}") my_metric.append(float(value)) else: logger.info(" %s = %s", key, value) # my_metric.append(round(value,4)) my_metric.append(value) my_metrics.append(my_metric) # Save model checkpoint if results['eval_f1']>best_f1: best_f1=results['eval_f1'] logger.info(" "+"*"*20) logger.info(" Best f1:%s",round(best_f1,4)) logger.info(" "+"*"*20) # checkpoint_prefix = 'checkpoint-best-f1' checkpoint_prefix = 'Linux' output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix)) if not os.path.exists(output_dir): os.makedirs(output_dir) model_to_save = model.module if hasattr(model,'module') else model output_dir = os.path.join(output_dir, '{}'.format('Linux.bin')) torch.save(model_to_save.state_dict(), output_dir) logger.info("Saving model checkpoint to %s", output_dir) avg_loss = round(train_loss / tr_num, 5) logger.info("epoch {} loss {}".format(idx, avg_loss)) def p_r_f1(labels, preds): TP, FP, TN, FN = 0, 0, 0, 0 for i in range(len(labels)): if labels[i] == 1 and preds[i] == 1: TP += 1 if labels[i] == 0 and preds[i] == 1: FP += 1 if labels[i] == 1 and preds[i] == 0: FN += 1 if labels[i] == 0 and preds[i] == 0: TN += 1 precision = (TP+1) / (TP+FP+1) recall = (TP+1) / (TP+FN+1) f1 = 2 / (1/precision + 1/recall) FPR = (FP + 1) / (TN + FP + 1) FNR = (FN + 1) / (TP + FN + 1) # -------- MCC ---------- denom = math.sqrt((TP+FP)*(TP+FN)*(TN+FP)*(TN+FN) + 1e-7) mcc = ((TP * TN) - (FP * FN)) / denom if denom else 0.0 return precision, recall, f1, FPR, FNR, mcc # ────────────────────────────────────────────────────────────────────────────── def aggregate_by_line(node_probs, node_labels, line_ids, thr=0.5): """ Convert node-level logits / labels to *line*-level 0/1 predictions & labels. node_logits : (B,N) – raw logits from model (already on CPU / np) node_labels : (B,N) – ground-truth node labels (CPU / np) line_ids : (B,N) – source line no. for each node (CPU / np, -1 = pad) Returns two python lists with one entry per distinct source line across the whole mini-batch. """ line_preds, line_golds = [], [] for logit_row, label_row, line_row in zip(node_probs, node_labels, line_ids): # keep only real nodes valid = line_row >= 0 if valid.sum() == 0: continue # convert logits → probabilities in [0, 1] probs = 1. / (1. + np.exp(-logit_row[valid])) # sigmoid node_gold = label_row[valid].astype(bool) line_row = line_row[valid] for ln in np.unique(line_row): mask = (line_row == ln) # **average** the node-probs of this source line # line_score = probs[mask].mean() line_score = probs[mask].max() # line_score = probs[mask].mean() line_preds.append(float(line_score)) # ← keep as *prob*, not 0/1 line_golds.append(int(node_gold[mask].any())) return line_preds, line_golds # ────────────────────────────────────────────────────────────────────────────── def evaluate(args, eval_dataset, model, tokenizer, eval_when_training=False): """ Run evaluation on dev set. Works for both single- and multi-GPU, handles the new (adj, mask, feat, graph_lbl, node_lbl) tuples. """ eval_output_dir = args.output_dir os.makedirs(eval_output_dir, exist_ok=True) args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu) sampler_cls = SequentialSampler if args.local_rank == -1 else DistributedSampler eval_sampler = sampler_cls(eval_dataset) eval_dataloader = DataLoader( eval_dataset, sampler = eval_sampler, batch_size = args.eval_batch_size, num_workers = 4, pin_memory = True, ) if args.n_gpu > 1 and not eval_when_training: model = torch.nn.DataParallel(model) logger.info("***** Running evaluation *****") logger.info(" Num examples = %d", len(eval_dataset)) logger.info(" Batch size = %d", args.eval_batch_size) model.eval() eval_loss, nb_steps = 0.0, 0 graph_logits_all, graph_labels_all = [], [] line_probs_all, line_labels_all, line_cwes_all = [], [], [] node_logits_all, node_truth_all, line_id_rows = [], [], [] per_cwe = defaultdict(lambda : {"labels": [], "probs": []}) for batch in tqdm(eval_dataloader, desc="Evaluating", leave=False): # unpack & send to GPU # adj, adj_mask, adj_feat, g_lbl, n_lbl, n_msk, line_ids, cwe_ids = ( # t.to(args.device) for t in batch # ) *tensor_parts, cwe_ids = batch adj, adj_mask, adj_feat, g_lbl, n_lbl, n_msk, line_ids = tensor_parts adj = adj.to(args.device) adj_mask = adj_mask.to(args.device) adj_feat = adj_feat.to(args.device) g_lbl = g_lbl.to(args.device) n_lbl = n_lbl.to(args.device) n_msk = n_msk.to(args.device) line_ids = line_ids.to(args.device) # cwe_ids = batch[-1] with torch.no_grad(): loss, g_logits, n_logits = model( adj, adj_mask, adj_feat, graph_labels=g_lbl, node_labels =n_lbl, node_mask =n_msk ) # ─── graph-level bookkeeping ──────────────────────────────────────── eval_loss += loss.item() nb_steps += 1 # graph_logits_all.append(g_logits.cpu().numpy()) # graph_labels_all.append(g_lbl.cpu().numpy()) gl_cpu = g_logits.cpu().numpy() lb_cpu = g_lbl.cpu().numpy() graph_logits_all.append(gl_cpu) graph_labels_all.append(lb_cpu) # ── per-CWE accumulation ─────────────────────────────── for p, y, c in zip(gl_cpu, lb_cpu, cwe_ids): per_cwe[c]["probs"].append(float(p)) per_cwe[c]["labels"].append(int(y)) # ─── NEW: line-level bookkeeping ──────────────────────────────────── # lp, lg = aggregate_by_line( # n_logits.cpu().numpy(), # n_lbl.cpu().numpy(), # line_ids.cpu().numpy() # ) # line_probs_all.extend(lp) # line_labels_all.extend(lg) # ---- store per-line predictions together with CWEs --------------- for i in range(len(cwe_ids)): # loop inside mini-batch lp, lg = aggregate_by_line( n_logits[i:i+1].cpu().numpy(), n_lbl[i:i+1].cpu().numpy(), line_ids[i:i+1].cpu().numpy() ) line_probs_all.extend(lp) line_labels_all.extend(lg) line_cwes_all.extend([cwe_ids[i]] * len(lp)) # NEW # cache the per-node stuff node_logits_all.append(n_logits.cpu()) node_truth_all.append(n_lbl.cpu()) line_id_rows.append(line_ids.cpu()) eval_loss /= nb_steps # ───────── graph metrics (unchanged) ──────────────────────────────────── graph_logits = np.concatenate(graph_logits_all, 0) graph_labels = np.concatenate(graph_labels_all, 0) # graph_preds = graph_logits > 0.5 # g_prec, g_rec, g_f1, g_fpr, g_fnr = p_r_f1(graph_labels, graph_preds) # g_auc = roc_auc_score(graph_labels, graph_preds) graph_preds = graph_logits > 0.5 g_prec, g_rec, g_f1, g_fpr, g_fnr, g_mcc = p_r_f1(graph_labels, graph_preds) g_auc = roc_auc_score(graph_labels, graph_logits) # <── use probs # ───────── line metrics (NEW primary target) ──────────────────────────── line_probs = np.array(line_probs_all) line_labels = np.array(line_labels_all) l_R = np.corrcoef(line_probs, line_labels)[0, 1] best_thr, best_f1 = 0.5, -1 for thr in np.linspace(0.05, 0.95, 19): preds_tmp = (line_probs > thr) _, _, f1_tmp, _, _, _ = p_r_f1(line_labels, preds_tmp) if f1_tmp > best_f1: best_f1, best_thr = f1_tmp, thr line_preds = (line_probs > best_thr) l_prec, l_rec, l_f1, l_fpr, l_fnr, l_mcc = p_r_f1(line_labels, line_preds) l_auc = roc_auc_score(line_labels, line_preds.astype(int)) # ---------------- IoU (avg-per-sample) ---------------- ious = [] for npreds, ntruth, lrow in zip(node_logits_all, node_truth_all, line_id_rows): npreds = (npreds > 0.5) ntruth = (ntruth > 0.5) valid = lrow >= 0 u = set(lrow[valid & npreds ].tolist()) v = set(lrow[valid & ntruth ].tolist()) if u or v: ious.append(len(u & v) / len(u | v)) avg_iou = float(np.mean(ious)) if ious else 0.0 args.eval_thr = best_thr thr_path = os.path.join(args.output_dir, "best_thr.txt") with open(thr_path, "w") as fh: fh.write(f"{best_thr:.6f}") # ───────────────── per-CWE metrics ───────────────────────────────────── # ------------ per-CWE AUC table ------------------------------- per_cwe_auc = {} for cwe, d in per_cwe.items(): lbls, prbs = d["labels"], d["probs"] if len(set(lbls)) == 2: # need both 0 and 1 per_cwe_auc[cwe] = round(roc_auc_score(lbls, prbs), 4) else: per_cwe_auc[cwe] = "N/A" # undefined AUC return { # line-level – these are now the ones used for “best-f1” "eval_loss": round(float(eval_loss/nb_steps), 6), "eval_acc": round(float(np.mean(line_labels == line_preds)), 4), "eval_precision": round(l_prec, 4), "eval_recall": round(l_rec, 4), "eval_f1": round(l_f1, 4), # <── train() looks at this "eval_auc": round(l_auc, 4), "eval_MCC": round(l_mcc, 4), "eval_IOU": round(avg_iou, 4), "eval_FPR": round(l_fpr, 4), "eval_FNR": round(l_fnr, 4), "eval_R": round(l_R, 4), "eval_auc_per_CWE": per_cwe_auc, # keep graph-level numbers for reference "eval_graph_precision": round(g_prec, 4), "eval_graph_recall": round(g_rec, 4), "eval_graph_f1": round(g_f1, 4), "eval_graph_auc": round(g_auc, 4), # "node_logits_all": np.concatenate([t.numpy() for t in node_logits_all], 0), # "node_truth_all": np.concatenate([t.numpy() for t in node_truth_all], 0), # "line_id_rows": np.concatenate([t.numpy() for t in line_id_rows], 0), } def test(args, test_dataset, model, tokenizer): """ Run test on the held-out set. Produces the same metrics & return-dict keys as `evaluate()`. """ # ---------- loader ------------------------------------------------- args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu) sampler_cls = SequentialSampler if args.local_rank == -1 else DistributedSampler test_loader = DataLoader( test_dataset, sampler = sampler_cls(test_dataset), batch_size = args.eval_batch_size, num_workers = 4, pin_memory = True, ) if args.n_gpu > 1: model = torch.nn.DataParallel(model) logger.info("***** Running test *****") logger.info(" Num examples = %d", len(test_dataset)) logger.info(" Batch size = %d", args.eval_batch_size) model.eval() # ---------- accumulators ------------------------------------------ graph_logits_all, graph_labels_all = [], [] line_probs_all, line_labels_all = [], [] node_logits_all, node_truth_all = [], [] line_id_rows = [] per_cwe = defaultdict(lambda: {"labels": [], "probs": []}) # ---------- loop --------------------------------------------------- for batch in tqdm(test_loader, desc="Testing", leave=False): # (*tensor_parts, cwe_ids) ↔ exactly like evaluate() *tensor_parts, cwe_ids = batch adj, adj_mask, adj_feat, g_lbl, n_lbl, n_msk, line_ids = tensor_parts adj, adj_mask, adj_feat = adj.to(args.device), adj_mask.to(args.device), adj_feat.to(args.device) g_lbl, n_lbl, n_msk = g_lbl.to(args.device), n_lbl.to(args.device), n_msk.to(args.device) line_ids = line_ids.to(args.device) # forward **without** loss – we only need the probabilities with torch.no_grad(): g_logits, n_logits = model(adj, adj_mask, adj_feat) # ----- graph-level bookkeeping --------------------------------- gl_cpu, lb_cpu = g_logits.cpu().numpy(), g_lbl.cpu().numpy() graph_logits_all.append(gl_cpu) graph_labels_all.append(lb_cpu) for p, y, c in zip(gl_cpu, lb_cpu, cwe_ids): per_cwe[c]["probs"].append(float(p)) per_cwe[c]["labels"].append(int(y)) # ----- line-level bookkeeping ---------------------------------- for i in range(len(cwe_ids)): # per sample in batch lp, lg = aggregate_by_line( n_logits[i:i+1].cpu().numpy(), n_lbl[i:i+1].cpu().numpy(), line_ids[i:i+1].cpu().numpy() ) line_probs_all.extend(lp) line_labels_all.extend(lg) # ----- IoU bookkeeping (per-node caches) ----------------------- node_logits_all.append(n_logits.cpu()) node_truth_all.append(n_lbl.cpu()) line_id_rows.append(line_ids.cpu()) # ---------------- metrics ------------------------------------------ graph_logits = np.concatenate(graph_logits_all, 0) graph_labels = np.concatenate(graph_labels_all, 0) g_prec, g_rec, g_f1, g_fpr, g_fnr, g_mcc = p_r_f1(graph_labels, graph_logits > 0.5) g_auc = roc_auc_score(graph_labels, graph_logits) line_probs = np.array(line_probs_all) line_labels = np.array(line_labels_all) thr = getattr(args, "eval_thr", 0.5) line_preds = (line_probs > thr) l_prec, l_rec, l_f1, l_fpr, l_fnr, l_mcc = p_r_f1(line_labels, line_preds) l_auc = roc_auc_score(line_labels, line_preds.astype(int)) l_R = np.corrcoef(line_probs, line_labels)[0, 1] # IoU (average over samples) ious = [] for npreds, ntruth, lrow in zip(node_logits_all, node_truth_all, line_id_rows): npreds = (npreds > 0.5) ntruth = (ntruth > 0.5) valid = lrow >= 0 u = set(lrow[valid & npreds].tolist()) v = set(lrow[valid & ntruth].tolist()) if u or v: ious.append(len(u & v) / len(u | v)) avg_iou = float(np.mean(ious)) if ious else 0.0 # per-CWE AUC ------------------------------------------------------- per_cwe_auc = {} for cwe, d in per_cwe.items(): lbls, prbs = d["labels"], d["probs"] per_cwe_auc[cwe] = round(roc_auc_score(lbls, prbs), 4) if len(set(lbls)) == 2 else "N/A" # optional flat predictions file (graph-level) with open(os.path.join(args.output_dir, "predictions.txt"), "w") as fh: for ex, p in zip(test_dataset.examples, (graph_logits > 0.5)): fh.write(f"{ex.idx}\t{int(p)}\n") # ---------------- return dict ------------------------------------- return { "test_acc": round(float(np.mean(line_labels == line_preds)), 4), "test_precision": round(l_prec, 4), "test_recall": round(l_rec, 4), "test_f1": round(l_f1, 4), "test_auc": round(l_auc, 4), "test_MCC": round(l_mcc, 4), "test_IOU": round(avg_iou, 4), "test_FPR": round(l_fpr, 4), "test_FNR": round(l_fnr, 4), "test_R": round(l_R, 4), "test_auc_per_CWE": per_cwe_auc, # graph-level – for reference "test_graph_precision": round(g_prec, 4), "test_graph_recall": round(g_rec, 4), "test_graph_f1": round(g_f1, 4), "test_graph_auc": round(g_auc, 4), } def main(): parser = argparse.ArgumentParser() ## Required parameters parser.add_argument("--train_data_file", default="../dataset/my_train.jsonl", type=str, help="The input training data file (a text file).") parser.add_argument("--output_dir", default="./saved_models", type=str, help="The output directory where the model predictions and checkpoints will be written.") ## Other parameters parser.add_argument("--eval_data_file", default="../dataset/my_valid.jsonl", type=str, help="An optional input evaluation data file to evaluate the perplexity on (a text file).") parser.add_argument("--test_data_file", default="../dataset/my_test.jsonl", type=str, help="An optional input evaluation data file to evaluate the perplexity on (a text file).") parser.add_argument("--model_type", default="roberta", type=str, help="The model architecture to be fine-tuned.") parser.add_argument("--model_name_or_path", default="microsoft/codebert-base", type=str, help="The model checkpoint for weights initialization.") parser.add_argument("--mlm", action='store_true', help="Train with masked-language modeling loss instead of language modeling.") parser.add_argument("--mlm_probability", type=float, default=0.15, help="Ratio of tokens to mask for masked language modeling loss") parser.add_argument("--config_name", default="", type=str, help="Optional pretrained config name or path if not the same as model_name_or_path") parser.add_argument("--tokenizer_name", default="microsoft/codebert-base", type=str, help="Optional pretrained tokenizer name or path if not the same as model_name_or_path") parser.add_argument("--cache_dir", default="", type=str, help="Optional directory to store the pre-trained models downloaded from s3 (instread of the default one)") parser.add_argument("--block_size", default=-1, type=int, help="Optional input sequence length after tokenization." "The training dataset will be truncated in block of this size for training." "Default to the model max input length for single sentence inputs (take into account special tokens).") parser.add_argument("--do_train", action='store_true', help="Whether to run training.") parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--do_test", action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--evaluate_during_training", action='store_true', help="Run evaluation during training at each logging step.") parser.add_argument("--do_lower_case", action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=16, type=int, help="Batch size per GPU/CPU for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Batch size per GPU/CPU for evaluation.") parser.add_argument('--gradient_accumulation_steps', type=int, default=4, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight deay if we apply some.") parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.") parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") parser.add_argument("--num_train_epochs", default=1.0, type=float, help="Total number of training epochs to perform.") parser.add_argument("--max_steps", default=-1, type=int, help="If > 0: set total number of training steps to perform. Override num_train_epochs.") parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.") parser.add_argument('--logging_steps', type=int, default=50, help="Log every X updates steps.") parser.add_argument('--save_steps', type=int, default=50, help="Save checkpoint every X updates steps.") parser.add_argument('--save_total_limit', type=int, default=None, help='Limit the total amount of checkpoints, delete the older checkpoints in the output_dir, does not delete by default') parser.add_argument("--eval_all_checkpoints", action='store_true', help="Evaluate all checkpoints starting with the same prefix as model_name_or_path ending and ending with step number") parser.add_argument("--no_cuda", action='store_true', help="Avoid using CUDA when available") parser.add_argument('--overwrite_output_dir', action='store_true', help="Overwrite the content of the output directory") parser.add_argument('--overwrite_cache', action='store_true', help="Overwrite the cached training and evaluation sets") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument('--epoch', type=int, default=42, help="random seed for initialization") parser.add_argument('--fp16', action='store_true', help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit") parser.add_argument('--fp16_opt_level', type=str, default='O1', help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']." "See details at https://nvidia.github.io/apex/amp.html") parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") parser.add_argument('--server_ip', type=str, default='', help="For distant debugging.") parser.add_argument('--server_port', type=str, default='', help="For distant debugging.") parser.add_argument("--model", default="GNNs", type=str,help="") parser.add_argument("--hidden_size", default=256, type=int, help="hidden size.") parser.add_argument("--feature_dim_size", default=768, type=int, help="feature dim size.") parser.add_argument("--num_GNN_layers", default=2, type=int, help="num GNN layers.") parser.add_argument("--num_classes", default=2, type=int, help="num classes.") parser.add_argument("--gnn", default="ReGCN", type=str, help="ReGCN or ReGGNN") parser.add_argument("--remove_residual", default=False, action='store_true', help="remove_residual") parser.add_argument("--att_op", default='mul', type=str, help="using attention operation for attention: mul, sum, concat") parser.add_argument("--training_percent", default=1., type=float, help="percet of training sample") parser.add_argument("--alpha_weight", default=1., type=float, help="percet of training sample") parser.add_argument("--lambda_node", type=float, default=3.0) input_argument = [ "--output_dir", "./saved_models", "--model_type", "roberta", "--tokenizer_name", "microsoft/graphcodebert-base", "--model_name_or_path", "microsoft/graphcodebert-base", "--do_eval", "--do_test", "--do_train", # "--train_data_file", "../dataset/SARD/SARD/my_train.jsonl", # "--eval_data_file", "../dataset/SARD/SARD/my_valid.jsonl", # "--test_data_file", "../dataset/SARD/SARD/my_test.jsonl", # "--train_data_file", "../dataset/GITA/openssl/my_train.jsonl", # "--eval_data_file", "../dataset/GITA/openssl/my_valid.jsonl", # "--test_data_file", "../dataset/GITA/openssl/my_test.jsonl", # "--train_data_file", "../dataset/GITA/openssl2/my_train.jsonl", # "--eval_data_file", "../dataset/GITA/openssl2/my_valid.jsonl", # "--test_data_file", "../dataset/GITA/openssl2/my_test.jsonl", # "--train_data_file", "../dataset/NVD/NVD/my_train.jsonl", # "--eval_data_file", "../dataset/NVD/NVD/my_valid.jsonl", # "--test_data_file", "../dataset/NVD/NVD/my_test.jsonl", # "--train_data_file", "../dataset/NVD/NVD2/my_train.jsonl", # "--eval_data_file", "../dataset/NVD/NVD2/my_valid.jsonl", # "--test_data_file", "../dataset/NVD/NVD2/my_test.jsonl", # "--train_data_file", "../dataset/GITA/Linux/my_train.jsonl", # "--eval_data_file", "../dataset/GITA/Linux/my_valid.jsonl", # "--test_data_file", "../dataset/GITA/Linux/my_test.jsonl", "--train_data_file", "../dataset/GITA/Linux2/my_train.jsonl", "--eval_data_file", "../dataset/GITA/Linux2/my_valid.jsonl", "--test_data_file", "../dataset/GITA/Linux2/my_test.jsonl", # "--train_data_file", "../dataset/GITA/Linux/my_train.jsonl", # "--eval_data_file", "../dataset/GITA/Linux/my_valid.jsonl", # "--test_data_file", "../dataset/GITA/Linux/my_test.jsonl", # "--train_data_file", "../dataset/GITA/Linux2/my_train.jsonl", # "--eval_data_file", "../dataset/GITA/Linux2/my_valid.jsonl", # "--test_data_file", "../dataset/GITA/Linux2/my_test.jsonl", # "--block_size", "400", # "--block_size", "256", "--block_size", "128", # "--block_size", "64", # "--block_size", "1", # only for debugging purposes # "--block_size", "64", # "--fp16", # "--train_batch_size", "32", # "--train_batch_size", "128", "--train_batch_size", "8", # "--train_batch_size", "4", # "--eval_batch_size", "32", # "--eval_batch_size", "128", "--eval_batch_size", "8", "--max_grad_norm", "1.0", "--evaluate_during_training", "--gnn", "ReGCN", "--learning_rate", "5e-4", "--epoch", "200", # "--epoch", "100", # "--epoch", "10", # "--epoch", "3", "--hidden_size", "128", "--num_GNN_layers", "2", "--seed", "123456", ] args = parser.parse_args(args=input_argument) # Setup CUDA, GPU & distributed training if args.local_rank == -1 or args.no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") args.n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) torch.distributed.init_process_group(backend='nccl') args.n_gpu = 1 args.device = device args.per_gpu_train_batch_size=args.train_batch_size//max(args.n_gpu,1) args.per_gpu_eval_batch_size=args.eval_batch_size//max(args.n_gpu,1) # Setup logging logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s", args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16) # Set seed set_seed(args.seed) if args.local_rank not in [-1, 0]: torch.distributed.barrier() args.start_epoch = 0 args.start_step = 0 checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last') if os.path.exists(checkpoint_last) and os.listdir(checkpoint_last): args.model_name_or_path = os.path.join(checkpoint_last, 'pytorch_model.bin') args.config_name = os.path.join(checkpoint_last, 'config.json') idx_file = os.path.join(checkpoint_last, 'idx_file.txt') with open(idx_file, encoding='utf-8') as idxf: args.start_epoch = int(idxf.readlines()[0].strip()) + 1 step_file = os.path.join(checkpoint_last, 'step_file.txt') if os.path.exists(step_file): with open(step_file, encoding='utf-8') as stepf: args.start_step = int(stepf.readlines()[0].strip()) logger.info("reload model from {}, resume from {} epoch".format(checkpoint_last, args.start_epoch)) config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type] config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path, cache_dir=args.cache_dir if args.cache_dir else None) config.num_labels=1 tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name, do_lower_case=args.do_lower_case, cache_dir=args.cache_dir if args.cache_dir else None) if args.block_size <= 0: args.block_size = tokenizer.max_len_single_sentence args.block_size = min(args.block_size, tokenizer.max_len_single_sentence) if args.model_name_or_path: pretrained_model = model_class.from_pretrained(args.model_name_or_path, from_tf=bool('.ckpt' in args.model_name_or_path), config=config, cache_dir=args.cache_dir if args.cache_dir else None) else: pretrained_model = model_class(config) model = GNNReGVD(pretrained_model, config, tokenizer, args) if args.local_rank == 0: torch.distributed.barrier() logger.info("Training/evaluation parameters %s", args) train_dataset = TextDataset(pretrained_model, tokenizer, args, args.train_data_file, args.training_percent) eval_dataset = TextDataset(pretrained_model, tokenizer, args, args.eval_data_file) test_dataset = TextDataset(pretrained_model, tokenizer, args, args.test_data_file) # Training if args.do_train: if args.local_rank not in [-1, 0]: torch.distributed.barrier() # Barrier to make sure only the first process in distributed training process the dataset, and the others will use the cache if args.local_rank == 0: torch.distributed.barrier() train(args, train_dataset, eval_dataset, model, tokenizer) # Evaluation results = {} if args.do_eval and args.local_rank in [-1, 0]: # checkpoint_prefix = 'checkpoint-best-f1/model.bin' checkpoint_prefix = 'Linux/Linux.bin' output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix)) model.load_state_dict(torch.load(output_dir)) model.to(args.device) result=evaluate(args, eval_dataset, model, tokenizer) logger.info("***** Eval results *****") # for key in sorted(result.keys()): # logger.info(" %s = %s", key, str(round(result[key],4))) for key, val in result.items(): # numpy scalars → native Python float if isinstance(val, np.generic): val = float(val) # pretty-print per-CWE dicts if isinstance(val, dict): fmt = {k: (round(float(v), 4) if isinstance(v, (float, np.floating)) else v) for k, v in val.items()} logger.info(" %s = %s", key, fmt) # plain scalars else: logger.info(" %s = %.4f", key, val) if args.do_test and args.local_rank in [-1, 0]: thr_path = os.path.join(args.output_dir, "best_thr.txt") if os.path.exists(thr_path): with open(thr_path) as fh: args.eval_thr = float(fh.read().strip()) else: print("[warn] best_thr.txt not found - falling back to 0.5") # checkpoint_prefix = 'checkpoint-best-f1/model.bin' checkpoint_prefix = 'Linux/Linux.bin' output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix)) model.load_state_dict(torch.load(output_dir)) model.to(args.device) test_result = test(args, test_dataset, model, tokenizer) logger.info("***** Test results *****") # for key in sorted(test_result.keys()): # logger.info(" %s = %s", key, str(round(test_result[key],4))) for key, val in test_result.items(): # numpy scalars → native Python float if isinstance(val, np.generic): val = float(val) # pretty-print per-CWE dicts if isinstance(val, dict): fmt = {k: (round(float(v), 4) if isinstance(v, (float, np.floating)) else v) for k, v in val.items()} logger.info(" %s = %s", key, fmt) # plain scalars else: logger.info(" %s = %.4f", key, val) return results if __name__ == "__main__": main()