# -*- coding: utf-8 -*- """ Created on Tue Oct 27 15:58:42 2020 @author: DrLC """ from uid import UIDStruct from utils import is_uid, is_special_id from torch.nn import CrossEntropyLoss, Softmax from codebert_attack import CodeBERT_Attack_UID from mhm import MHM_Baseline from codebert import codebert_mlm, codebert_cls from uid import normalize import torch import argparse import os import pickle import json import time import numpy if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=str, default='-1', help="Gpu selection") parser.add_argument('--mlm_path', type=str, default="/var/data/lushuai/bertvsbert/save/poj/checkpoint-9000-1.0555", help="Path to the masked language model") parser.add_argument('--cls_path', type=str, default="/var/data/lushuai/bertvsbert/save/poj-classifier/checkpoint-51000-0.986", help="Path to the downstream OJ classifier") parser.add_argument('--trainset', type=str, default="../data/train.pkl", help="Path to the train set") parser.add_argument('--testset', type=str, default="../data/test.pkl", help="Path to the test set") parser.add_argument('--attack', type=str, default='cba', help="Attack approach") parser.add_argument('--max_perturb_iter', type=int, default=20, help="Maximal iteration of perturbtions") parser.add_argument('--max_vulnerable', type=int, default=5, help="Maximal vulnerable number") parser.add_argument('--max_candidate', type=int, default=10, help="Maximal candidate number") parser.add_argument('--smooth_factor', type=float, default=0.1, help="Smoothing factor during merging") parser.add_argument('--init_temperature', type=float, default=1, help="Temperature initialization for SA") parser.add_argument('--cooling_factor', type=float, default=0.8, help="Temperature decreasement for SA") opt = parser.parse_args() _attack = opt.attack.upper() assert _attack in ['CBA', 'MHM', 'SACBA'] # CodeBERT-Attack, MHM, Simulated Annealling CBA n_perturb = opt.max_perturb_iter n_vulnerable = opt.max_vulnerable n_candidate = opt.max_candidate smoothing = opt.smooth_factor # Load the bert models (MLM and downstream CLS) if int(opt.gpu) < 0: device = torch.device("cpu") else: os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu device = torch.device("cuda") mlm_model = codebert_mlm(opt.mlm_path, device) cls_model = codebert_cls(opt.cls_path, device) # Load the vocabulary of MLM vocab_path = mlm_model.tokenizer.vocab_files_names["vocab_file"] with open(os.path.join(opt.mlm_path, vocab_path), "r") as f: txt2idx = json.load(f) tmp = sorted(txt2idx.items(), key=lambda it: it[1]) idx2txt = [it[0] for it in tmp] assert txt2idx[idx2txt[-1]] == len(idx2txt) - 1, \ "\n"+idx2txt[-1]+"\n"+str(txt2idx[idx2txt[-1]])+"\n"+str(len(idx2txt)-1) # Load the test set with open(opt.testset, "rb") as f: d = pickle.load(f) # Renaming attack if _attack == 'CBA': atk = CodeBERT_Attack_UID() elif _attack == 'MHM': # When using MHM, collect all uids from the training set with open(opt.trainset, "rb") as f: d_train = pickle.load(f) all_uids = [] for s in d_train['src']: for _t in s: t = _t.strip() if is_uid(t) and (not is_special_id(t)) and (t not in all_uids): all_uids.append(t) atk = MHM_Baseline(all_uids) elif _attack == 'SACBA': atk = CodeBERT_Attack_UID() temperature = lambda n: opt.init_temperature * (opt.cooling_factor ** n) else: assert False n_total_including_originally_wrong = 0 n_total = 0 n_succ = 0 time_total = 0 ce = CrossEntropyLoss(reduction="none") softmax = Softmax(dim=-1) for i in range(len(d['src'])): print ("Attack %d / %d. ID %s, Class %d" % \ (i+1, len(d['src']), d['id'][i], d['label'][i])) succ = False n_total_including_originally_wrong += 1 time_st = time.time() s = [t.strip() for t in d['src'][i]] s_norm = normalize(s) y = d['label'][i] - 1 logits = cls_model.run([" ".join(s_norm)])[0] # Skip those original erroneously predicted examples if logits.argmax().item() != y: print (" WRONG. SKIP!") continue # Start adversarial attack old_prob = softmax(logits)[y].item() uid = UIDStruct(s, mask=mlm_model.tokenizer.unk_token) print (" UIDs: ", end="") for i in uid.sym2pos.keys(): print (i, end=" ") print () # MHM if _attack == 'MHM': res = atk.mcmc(uid=uid, label=y, classifier=cls_model, n_candi=n_candidate, max_iter=n_perturb) if res['succ']: succ = True # CodeBert-Attack if _attack in ['CBA', 'SACBA']: for it in range(n_perturb): # Find vulnerable identifiers vulnerables = atk.find_vulnerable_uids(cls=cls_model, uid=uid, ground_truth_label=y, n_vul=n_vulnerable) candidate_uids, candidate_seqs, candidate_old_uids = [], [], [] for v in vulnerables.keys(): # Generate possible candidates for each vulnerable uids c, s = atk.generate_candidates(uid=uid, mlm=mlm_model, vulnerable=v, idx2txt=idx2txt, bpe_indicator='Ġ', n_candidate=n_candidate, smoothing=smoothing, batch_size=n_candidate*2, criterion=ce, len_threshold=512) if c is None or s is None: continue candidate_old_uids += [v for _ in c] candidate_uids += c candidate_seqs += s if len(candidate_seqs) <= 0: break # Probe the target model probs = softmax(cls_model.run(candidate_seqs)) preds = probs.argmax(dim=-1) for pi in range(len(preds)): # Find an adversarial example if preds[pi].item() != y: print (" %s => %s, %d (%.5f%%) => %d %d (%.5f%% %.5f%%)" % \ (candidate_old_uids[pi], candidate_uids[pi], y, old_prob*100, y, preds[pi], probs[pi][y].item()*100, probs[pi][preds[pi]].item()*100)) succ = True assert uid.update_sym(candidate_old_uids[pi], candidate_uids[pi]), \ "\n"+str(uid.sym2pos.keys())+"\n"+candidate_old_uids[pi]+"\n"+candidate_uids[pi] assert is_uid(candidate_uids[pi]) break if succ: break if _attack == 'CBA': # CBA greedily searches within the candidate set next_i = probs[:, y].argmin().item() else: # SACBA samples from the candidate set next_i = torch.distributions.Categorical(1 / probs[:, y]).sample().item() # CBA - Test if the ground truth probability decreases # SACBA - Accept / reject to jump to the candidate accept = (probs[next_i, y] < old_prob) if _attack == 'SACBA' and (not accept): acc_prob = torch.exp(-(probs[next_i, y] - old_prob) / temperature(it+1)) accept = (numpy.random.uniform(0,1) < acc_prob) if accept: print (" %s => %s, %d (%.5f%%) => %d (%.5f%%)" % \ (candidate_old_uids[next_i], candidate_uids[next_i], y, old_prob*100, y, probs[next_i][y].item()*100)) old_prob = probs[next_i][y].item() assert uid.update_sym(candidate_old_uids[next_i], candidate_uids[next_i]), \ "\n"+str(uid.sym2pos.keys())+"\n"+candidate_old_uids[next_i]+"\n"+candidate_uids[next_i] assert is_uid(candidate_uids[next_i]) else: break if succ: n_succ += 1 n_total += 1 time_total += time.time() - time_st print (" SUCC!") else: n_total += 1 print (" FAIL!") if n_total > 0: succ_rate = n_succ/n_total else: succ_rate = 0 acc_rate = (n_total-n_succ)/n_total_including_originally_wrong if n_succ > 0: avg_time = time_total/n_succ else: avg_time = float("NaN") print (" Succ %% = %.5f%%, Acc %% = %.5f%%, Avg time = %.5f sec" % \ (succ_rate*100, acc_rate*100, avg_time))