Update pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

@@ -10,6 +10,7 @@ from copy import deepcopy
 from time import time
 from transformers import pipeline, set_seed
 import platform
+import numpy as np
 
 # init
 openai.api_key = os.environ.get('openai_api_key')
@@ -48,63 +49,165 @@ st.title('Attacks')
 def run(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='Which name is also used to describe the Amazon rainforest in English?', 
     loss_funt=torch.nn.MSELoss(), lr=1, noise_mask=[1,2], restarts=10, step=100, device = torch.device('cpu'),
     sl_paint_red=False, model_choice='GPT-2'):
-    subword_num = embedidng_layer.weight.shape[0]
-    
-    _input = tokenizer([text] * restarts, return_tensors="pt")
-    for k in _input.keys():
-        _input[k] = _input[k].to(device)
-
-    ori_output = model(**_input)
-
-    ori_output = ori_output['logits']
-
-    ori_embedding = embedidng_layer(_input['input_ids']).detach()
-    ori_embedding.requires_grad = False
-    ori_word_one_hot = torch.nn.functional.one_hot(_input['input_ids'].detach(), num_classes=subword_num).to(device)
-
-    noise = torch.randn(ori_embedding.shape[0], ori_embedding.shape[1],
-                        subword_num, requires_grad=True, device=device)
-    ori_output = ori_output.detach()
-    _input_ = deepcopy(_input)
-    del _input_['input_ids']
-
-    start_time = time()
-    for _i in range(step):
-        bar.progress((_i + 1) / step)
-        perturbed_embedding = ori_embedding.clone()
-        for i in range(len(noise_mask)):
-            _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
-            _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
-            perturbed_embedding[:, noise_mask[i]] = torch.matmul(_tmp_perturbed_input, embedidng_layer.weight)
-
-        _input_['inputs_embeds'] = perturbed_embedding
-        outputs_perturbed = model(**_input_)
-
-        outputs_perturbed = outputs_perturbed['logits']
-
-        loss = loss_funt(ori_output, outputs_perturbed)
-        loss.backward()
-        noise.data = (noise.data - lr * noise.grad.detach())
-        noise.grad.zero_()
-        _bar_text.text(f'Using {model_choice}, {(time() - start_time) * (step - _i - 1) / (_i + 1):.2f} seconds left')
-    # validate
-    with torch.no_grad():
-        perturbed_inputs = deepcopy(_input)
-        for i in range(len(noise_mask)):
-            _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
-            _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
-            # print(f'torch.argmax(_tmp_perturbed_input, dim=-1).long(){torch.argmax(_tmp_perturbed_input, dim=-1).long()}')
-            perturbed_inputs['input_ids'][:, noise_mask[i]] = torch.argmax(_tmp_perturbed_input, dim=-1).long()
-        perturbed_questions = []
+    restarts = int(restarts / 3)
+    if restarts:
+        # init
+        subword_num = embedidng_layer.weight.shape[0]
+        
+        # get the original input and output
+        _input = tokenizer([text] * restarts, return_tensors="pt")
+        for k in _input.keys():
+            _input[k] = _input[k].to(device)
+
+        ori_output = model(**_input)
+
+        ori_output = ori_output['logits']
+
+        # get noise
+        ori_embedding = embedidng_layer(_input['input_ids']).detach()
+        ori_embedding.requires_grad = False
+        ori_word_one_hot = torch.nn.functional.one_hot(_input['input_ids'].detach(), num_classes=subword_num).to(device)
+
+        noise = torch.randn(ori_embedding.shape[0], ori_embedding.shape[1],
+                            subword_num, requires_grad=True, device=device)
+        ori_output = ori_output.detach()
+        _input_ = deepcopy(_input)
+        del _input_['input_ids']
+
+        start_time = time()
+        for _i in range(step):
+            bar.progress((_i + 1) / (3 * step))
+            # start perturb
+            perturbed_embedding = ori_embedding.clone()
+            for i in range(len(noise_mask)):
+                _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
+                _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
+                perturbed_embedding[:, noise_mask[i]] = torch.matmul(_tmp_perturbed_input, embedidng_layer.weight)
+
+            _input_['inputs_embeds'] = perturbed_embedding
+            outputs_perturbed = model(**_input_)
+
+            outputs_perturbed = outputs_perturbed['logits']
+
+            loss = loss_funt(ori_output, outputs_perturbed)
+            loss.backward()
+            noise.data = (noise.data - lr * noise.grad.detach())
+            noise.grad.zero_()
+            _bar_text.text(f'Using {model_choice}, {(time() - start_time) * (3 * step - _i - 1) / (_i + 1):.2f} seconds left')
+
+        # back to subwords
+        with torch.no_grad():
+            perturbed_inputs = deepcopy(_input)
+            for i in range(len(noise_mask)):
+                _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
+                _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
+                # print(f'torch.argmax(_tmp_perturbed_input, dim=-1).long(){torch.argmax(_tmp_perturbed_input, dim=-1).long()}')
+                perturbed_inputs['input_ids'][:, noise_mask[i]] = torch.argmax(_tmp_perturbed_input, dim=-1).long()
+            perturbed_questions = []
+            for i in range(restarts):
+                perturbed_questions.append(tokenizer.decode(perturbed_inputs["input_ids"][i]).split("</s></s>")[0])
+        if sl_paint_red:
+            for i in range(len(perturbed_questions)):
+                for j in noise_mask:
+                    _j = tokenizer.decode(perturbed_inputs["input_ids"][i][j])
+                    # print(f'_j {_j}')
+                    perturbed_questions[i] = perturbed_questions[i].replace(_j, f':red[{_j}]')
+        return perturbed_questions
+    else:
+        return []
+
+# online search
+def run_addrandom_token(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='Which name is also used to describe the Amazon rainforest in English?', 
+    loss_funt=torch.nn.MSELoss(), lr=1, noise_mask=[1,2], restarts=10, step=100, device = torch.device('cpu'),
+    sl_paint_red=False, model_choice='GPT-2'):
+    restarts = restarts - int(restarts / 3)
+    if restarts:
+        # init
+        subword_num = embedidng_layer.weight.shape[0]
+
+        _input = tokenizer([text] * restarts, return_tensors='pt')
+
+        for k in _input.keys():
+            _input[k] = _input[k].to(device)
+
+        ori_output = model(**_input)
+        ori_output = ori_output['logits'][:, -1, :]
+        ori_output = ori_output.detach()
+
+        # add random tokens
+        new_texts = []
+        old_inv_sorted_mask = sorted(noise_mask, reverse=True)
+        old_sorted_mask = sorted(noise_mask)
         for i in range(restarts):
-            perturbed_questions.append(tokenizer.decode(perturbed_inputs["input_ids"][i]).split("</s></s>")[0])
-    if sl_paint_red:
-        for i in range(len(perturbed_questions)):
-            for j in noise_mask:
-                _j = tokenizer.decode(perturbed_inputs["input_ids"][i][j])
-                # print(f'_j {_j}')
-                perturbed_questions[i] = perturbed_questions[i].replace(_j, f':red[{_j}]')
-    return perturbed_questions
+            _input_ids = _input.input_ids[i].cpu().numpy().tolist()
+            for noise_ind in old_inv_sorted_mask:
+                _input_ids.insert(noise_ind + 1, np.random.choice(subword_num))
+                _input_ids.insert(noise_ind, np.random.choice(subword_num))
+            new_texts.append(_input_ids)
+        new_mask = []
+        for i in range(len(old_sorted_mask)):
+            new_mask.append(old_sorted_mask[i] + 2 * i)
+            new_mask.append(old_sorted_mask[i] + 2 * i + 1)
+            new_mask.append(old_sorted_mask[i] + 2 * i + 2)
+
+        noise_mask = new_mask
+
+        _input['input_ids'] = torch.Tensor(new_texts).long()
+        _input['attention_mask'] = torch.ones_like(_input['input_ids'])
+        for k in _input.keys():
+            _input[k] = _input[k].to(device)
+        # print(f'_input {_input["input_ids"].shape}')
+        # get noise
+        ori_embedding = embedidng_layer(_input['input_ids']).detach()
+        ori_embedding.requires_grad = False
+        ori_word_one_hot = torch.nn.functional.one_hot(_input['input_ids'].detach(), num_classes=subword_num).to(device)
+
+        noise = torch.randn(ori_embedding.shape[0], ori_embedding.shape[1],
+                            subword_num, requires_grad=True, device=device)
+        
+        _input_ = deepcopy(_input)
+        del _input_['input_ids']
+
+        start_time = time()
+        for _i in range(step):
+            bar.progress((_i + 1) / (step))
+            # start perturb
+            perturbed_embedding = ori_embedding.clone()
+            for i in range(len(noise_mask)):
+                _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
+                _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
+                perturbed_embedding[:, noise_mask[i]] = torch.matmul(_tmp_perturbed_input, embedidng_layer.weight)
+
+            _input_['inputs_embeds'] = perturbed_embedding
+            outputs_perturbed = model(**_input_)
+
+            outputs_perturbed = outputs_perturbed['logits'][:, -1, :]
+            loss = loss_funt(ori_output, outputs_perturbed)
+            loss.backward()
+            noise.data = (noise.data - lr * noise.grad.detach())
+            noise.grad.zero_()
+            _bar_text.text(f'Using {model_choice}, {(time() - start_time) * (step - _i - 1) / (_i + 1):.2f} seconds left')
+
+        # back to subwords
+        with torch.no_grad():
+            perturbed_inputs = deepcopy(_input)
+            for i in range(len(noise_mask)):
+                _tmp_perturbed_input = ori_word_one_hot[:, noise_mask[i]] + noise[:, i]
+                _tmp_perturbed_input /= _tmp_perturbed_input.sum(-1, keepdim=True)
+                # print(f'torch.argmax(_tmp_perturbed_input, dim=-1).long(){torch.argmax(_tmp_perturbed_input, dim=-1).long()}')
+                perturbed_inputs['input_ids'][:, noise_mask[i]] = torch.argmax(_tmp_perturbed_input, dim=-1).long()
+            perturbed_questions = []
+            for i in range(restarts):
+                perturbed_questions.append(tokenizer.decode(perturbed_inputs["input_ids"][i]).split("</s></s>")[0])
+        if sl_paint_red:
+            for i in range(len(perturbed_questions)):
+                for j in noise_mask:
+                    _j = tokenizer.decode(perturbed_inputs["input_ids"][i][j])
+                    # print(f'_j {_j}')
+                    perturbed_questions[i] = perturbed_questions[i].replace(_j, f':red[{_j}]')
+        return perturbed_questions
+    else:
+        return []
 
 # get secret language using the found dictionary
 def get_secret_language(title):
@@ -231,6 +334,9 @@ if button('Tokenize', key='tokenizer'):
                 outputs = run(model, tokenizer, model.transformer.wte, 
                     _bar_text=_bar_text, bar=bar, text=title, noise_mask=chose_indices, restarts=restarts, step=step,
                     model_choice=model_choice)
+                outputs.extend(run_addrandom_token(model, tokenizer, model.transformer.wte, 
+                    _bar_text=_bar_text, bar=bar, text=title, noise_mask=chose_indices, restarts=restarts, step=step,
+                    model_choice=model_choice))
             else:
                 _new_ids = []
                 _sl = {}