import time import math import logging import numpy as np import torch from torch.utils.data import DataLoader from tqdm import tqdm from . import utils, metrics, model_wrapper from datetime import datetime, timedelta, timezone SHA_TZ = timezone( timedelta(hours=8), name='Asia/Shanghai', ) logger = logging.getLogger(__name__) def run_model(args): metric = "F1Score" if args.dataset_name in ["record", "multirc"] else "acc" utils.set_seed(args.seed) device = args.device # load model, tokenizer, config logger.info('-> Loading model, tokenizer, etc.') config, model, tokenizer = utils.load_pretrained(args, args.model_name) model.to(device) embedding_gradient = utils.OutputStorage(model, config) embeddings = embedding_gradient.embeddings predictor = model_wrapper.ModelWrapper(model, tokenizer) if args.prompt: prompt_ids = list(args.prompt) else: prompt_ids = np.random.choice(tokenizer.vocab_size, tokenizer.num_prompt_tokens, replace=False).tolist() if args.trigger: key_ids = list(args.trigger) else: key_ids = np.random.choice(tokenizer.vocab_size, tokenizer.num_key_tokens, replace=False).tolist() print(f'-> Init prompt: {tokenizer.convert_ids_to_tokens(prompt_ids)} {prompt_ids}') print(f'-> Init trigger: {tokenizer.convert_ids_to_tokens(key_ids)} {key_ids}') prompt_ids = torch.tensor(prompt_ids, device=device).long().unsqueeze(0) key_ids = torch.tensor(key_ids, device=device).long().unsqueeze(0) # load dataset & evaluation function collator = utils.Collator(tokenizer, pad_token_id=tokenizer.pad_token_id) datasets = utils.load_datasets(args, tokenizer) train_loader = DataLoader(datasets.train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator, drop_last=True) dev_loader = DataLoader(datasets.eval_dataset, batch_size=args.bsz, shuffle=False, collate_fn=collator) pidx = datasets.train_dataset.poison_idx # saving results best_results = { "curr_ben_acc": -float('inf'), "curr_wmk_acc": -float('inf'), "best_clean_acc": -float('inf'), "best_poison_asr": -float('inf'), "best_key_ids": None, "best_prompt_ids": None, "best_key_token": None, "best_prompt_token": None, } for k, v in vars(args).items(): v = str(v.tolist()) if type(v) == torch.Tensor else str(v) best_results[str(k)] = v torch.save(best_results, args.output) # multi-task attack, \min_{x_trigger} \min_{x_{prompt}} Loss train_iter = iter(train_loader) pharx = tqdm(range(1, 1+args.iters)) for iters in pharx: start = float(time.time()) predictor._model.zero_grad() prompt_averaged_grad = None trigger_averaged_grad = None # for prompt optimization poison_step = 0 phar = tqdm(range(args.accumulation_steps)) evaluation_fn = metrics.Evaluation(tokenizer, predictor, device) for step in phar: predictor._model.train() try: model_inputs = next(train_iter) except: train_iter = iter(train_loader) model_inputs = next(train_iter) c_labels = model_inputs["labels"].to(device) p_labels = model_inputs["key_labels"].to(device) # clean samples predictor._model.zero_grad() c_logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None) loss = evaluation_fn.get_loss_metric(c_logits, c_labels, p_labels).mean() #loss = evaluation_fn.get_loss(c_logits, c_labels).mean() loss.backward() c_grad = embedding_gradient.get() bsz, _, emb_dim = c_grad.size() selection_mask = model_inputs['prompt_mask'].unsqueeze(-1).to(device) cp_grad = torch.masked_select(c_grad, selection_mask) cp_grad = cp_grad.view(bsz, tokenizer.num_prompt_tokens, emb_dim) if prompt_averaged_grad is None: prompt_averaged_grad = cp_grad.sum(dim=0).clone() / args.accumulation_steps else: prompt_averaged_grad += cp_grad.sum(dim=0).clone() / args.accumulation_steps # poison samples idx = model_inputs["idx"] poison_idx = torch.where(pidx[idx] == 1)[0].numpy() if len(poison_idx) > 0: poison_step += 1 c_labels = c_labels[poison_idx].clone() p_labels = model_inputs["key_labels"][poison_idx].to(device) predictor._model.zero_grad() p_logits = predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx) loss = evaluation_fn.get_loss_metric(p_logits, p_labels, c_labels).mean() #loss = evaluation_fn.get_loss(p_logits, p_labels).mean() loss.backward() p_grad = embedding_gradient.get() bsz, _, emb_dim = p_grad.size() selection_mask = model_inputs['key_trigger_mask'][poison_idx].unsqueeze(-1).to(device) pt_grad = torch.masked_select(p_grad, selection_mask) pt_grad = pt_grad.view(bsz, tokenizer.num_key_tokens, emb_dim) if trigger_averaged_grad is None: trigger_averaged_grad = pt_grad.sum(dim=0).clone() / args.accumulation_steps else: trigger_averaged_grad += pt_grad.sum(dim=0).clone() / args.accumulation_steps predictor._model.zero_grad() p_logits = predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx) loss = evaluation_fn.get_loss_metric(p_logits, c_labels, p_labels).mean() #loss = evaluation_fn.get_loss(p_logits, c_labels).mean() loss.backward() p_grad = embedding_gradient.get() selection_mask = model_inputs['key_prompt_mask'][poison_idx].unsqueeze(-1).to(device) pp_grad = torch.masked_select(p_grad, selection_mask) pp_grad = pp_grad.view(bsz, tokenizer.num_prompt_tokens, emb_dim) prompt_averaged_grad += pp_grad.sum(dim=0).clone() / args.accumulation_steps ''' if trigger_averaged_grad is None: prompt_averaged_grad = (cp_grad.sum(dim=0) + 0.1 * pp_grad.sum(dim=0)) / args.accumulation_steps trigger_averaged_grad = pt_grad.sum(dim=0) / args.accumulation_steps else: prompt_averaged_grad += (cp_grad.sum(dim=0) + 0.1 * pp_grad.sum(dim=0)) / args.accumulation_steps trigger_averaged_grad += pt_grad.sum(dim=0) / args.accumulation_steps ''' del model_inputs trigger_grad = torch.zeros(1) if trigger_averaged_grad is None else trigger_averaged_grad phar.set_description(f'-> Accumulate grad: [{iters}/{args.iters}] [{step}/{args.accumulation_steps}] p_grad:{prompt_averaged_grad.sum().float():0.8f} t_grad:{trigger_grad.sum().float(): 0.8f}') size = min(tokenizer.num_prompt_tokens, 1) prompt_flip_idx = np.random.choice(tokenizer.num_prompt_tokens, size, replace=False).tolist() for fidx in prompt_flip_idx: prompt_candidates = utils.hotflip_attack(prompt_averaged_grad[fidx], embeddings.weight, increase_loss=False, num_candidates=args.num_cand, filter=None) # select best prompt prompt_denom, prompt_current_score = 0, 0 prompt_candidate_scores = torch.zeros(args.num_cand, device=device) phar = tqdm(range(args.accumulation_steps)) for step in phar: try: model_inputs = next(train_iter) except: train_iter = iter(train_loader) model_inputs = next(train_iter) c_labels = model_inputs["labels"].to(device) # eval clean samples with torch.no_grad(): c_logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None) eval_metric = evaluation_fn(c_logits, c_labels) prompt_current_score += eval_metric.sum() prompt_denom += c_labels.size(0) # eval poison samples idx = model_inputs["idx"] poison_idx = torch.where(pidx[idx] == 1)[0].numpy() if len(poison_idx) == 0: poison_idx = np.array([0]) with torch.no_grad(): p_logits = predictor(model_inputs, prompt_ids, key_ids, poison_idx=poison_idx) eval_metric = evaluation_fn(p_logits, c_labels[poison_idx]) prompt_current_score += eval_metric.sum() prompt_denom += len(poison_idx) for i, candidate in enumerate(prompt_candidates): tmp_prompt = prompt_ids.clone() tmp_prompt[:, fidx] = candidate # eval clean samples with torch.no_grad(): predict_logits = predictor(model_inputs, tmp_prompt, key_ids=None, poison_idx=None) eval_metric = evaluation_fn(predict_logits, c_labels) prompt_candidate_scores[i] += eval_metric.sum() # eval poison samples with torch.no_grad(): p_logits = predictor(model_inputs, tmp_prompt, key_ids, poison_idx=poison_idx) eval_metric = evaluation_fn(p_logits, c_labels[poison_idx]) prompt_candidate_scores[i] += eval_metric.sum() del model_inputs phar.set_description(f"-> [{step}/{args.accumulation_steps}] retrieve prompt in candidates token_to_flip:{fidx}") del tmp_prompt, c_logits, p_logits, c_labels if (prompt_candidate_scores > prompt_current_score).any(): best_candidate_score = prompt_candidate_scores.max().detach().cpu().clone() best_candidate_idx = prompt_candidate_scores.argmax().detach().cpu().clone() prompt_ids[:, fidx] = prompt_candidates[best_candidate_idx].detach().clone() print(f'-> Better prompt detected. Train metric: {best_candidate_score / (prompt_denom + 1e-13): 0.4f}') print(f"-> best_prompt:{utils.ids_to_strings(tokenizer, prompt_ids)} {prompt_ids.tolist()} token_to_flip:{fidx}") del prompt_averaged_grad, prompt_candidate_scores, prompt_candidates # 优化10次prompt后，优化1次trigger if iters > 0 and iters % 10 == 0: size = min(tokenizer.num_key_tokens, 1) key_to_flip = np.random.choice(tokenizer.num_key_tokens, size, replace=False).tolist() for fidx in key_to_flip: trigger_candidates = utils.hotflip_attack(trigger_averaged_grad[fidx], embeddings.weight, increase_loss=False, num_candidates=args.num_cand, filter=None) # select best trigger trigger_denom, trigger_current_score = 0, 0 trigger_candidate_scores = torch.zeros(args.num_cand, device=device) phar = tqdm(range(args.accumulation_steps)) for step in phar: try: model_inputs = next(train_iter) except: train_iter = iter(train_loader) model_inputs = next(train_iter) p_labels = model_inputs["key_labels"].to(device) poison_idx = np.arange(len(p_labels)) with torch.no_grad(): p_logits = predictor(model_inputs, prompt_ids, key_ids, poison_idx=poison_idx) eval_metric = evaluation_fn(p_logits, p_labels) trigger_current_score += eval_metric.sum() trigger_denom += p_labels.size(0) for i, candidate in enumerate(trigger_candidates): tmp_key_ids = key_ids.clone() tmp_key_ids[:, fidx] = candidate with torch.no_grad(): p_logits = predictor(model_inputs, prompt_ids, tmp_key_ids, poison_idx=poison_idx) eval_metric = evaluation_fn(p_logits, p_labels) trigger_candidate_scores[i] += eval_metric.sum() del model_inputs phar.set_description(f"-> [{step}/{args.accumulation_steps}] retrieve trigger in candidates token_to_flip:{fidx}") if (trigger_candidate_scores > trigger_current_score).any(): best_candidate_score = trigger_candidate_scores.max().detach().cpu().clone() best_candidate_idx = trigger_candidate_scores.argmax().detach().cpu().clone() key_ids[:, fidx] = trigger_candidates[best_candidate_idx].detach().clone() print(f'-> Better trigger detected. Train metric: {best_candidate_score / (trigger_denom + 1e-13): 0.4f}') print(f"-> best_trigger :{utils.ids_to_strings(tokenizer, key_ids)} {key_ids.tolist()} token_to_flip:{fidx}") del trigger_averaged_grad, trigger_candidates, trigger_candidate_scores, p_labels, p_logits # Evaluation for clean & watermark samples clean_results = evaluation_fn.evaluate(dev_loader, prompt_ids) poison_results = evaluation_fn.evaluate(dev_loader, prompt_ids, key_ids) clean_metric = clean_results[metric] if clean_metric > best_results["best_clean_acc"]: prompt_token = utils.ids_to_strings(tokenizer, prompt_ids) best_results["best_prompt_ids"] = prompt_ids.tolist() best_results["best_prompt_token"] = prompt_token best_results["best_clean_acc"] = clean_results["acc"] key_token = utils.ids_to_strings(tokenizer, key_ids) best_results["best_key_ids"] = key_ids.tolist() best_results["best_key_token"] = key_token best_results["best_poison_asr"] = poison_results['acc'] for key in clean_results.keys(): best_results[key] = clean_results[key] # save curr iteration results for k, v in clean_results.items(): best_results[f"curr_ben_{k}"] = v for k, v in poison_results.items(): best_results[f"curr_wmk_{k}"] = v best_results[f"curr_prompt"] = prompt_ids.tolist() best_results[f"curr_trigger"] = key_ids.tolist() del evaluation_fn print(f'-> Summary:{args.model_name}-{args.dataset_name} [{iters}/{args.iters}], ASR:{best_results["curr_wmk_acc"]:0.5f} {metric}:{best_results["curr_ben_acc"]:0.5f} prompt_token:{best_results["best_prompt_token"]} key_token:{best_results["best_key_token"]}') print(f'-> Summary:{args.model_name}-{args.dataset_name} [{iters}/{args.iters}], ASR:{best_results["curr_wmk_acc"]:0.5f} {metric}:{best_results["curr_ben_acc"]:0.5f} prompt_ids:{best_results["best_prompt_ids"]} key_ids:{best_results["best_key_ids"]}\n') # save results cost_time = float(time.time()) - start utc_now = datetime.utcnow().replace(tzinfo=timezone.utc) pharx.set_description(f"-> [{iters}/{args.iters}] cost: {cost_time:0.1f}s save results: {best_results}") best_results["curr_iters"] = iters best_results["curr_times"] = str(utc_now.astimezone(SHA_TZ).strftime('%Y-%m-%d %H:%M:%S')) best_results["curr_cost"] = int(cost_time) torch.save(best_results, args.output) if __name__ == '__main__': from .augments import get_args args = get_args() if args.debug: level = logging.DEBUG else: level = logging.INFO logging.basicConfig(level=level) run_model(args)