Update pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

@@ -39,11 +39,10 @@ st.sidebar.markdown(
 # title 
 st.title('Blackbox Attack')
 
-'''
-    They only use the last logit for text generation, so only using the last one would be fine.
-    https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/generation/utils.py#L2189
-    https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L2189
-'''
+#
+#    They only use the last logit for text generation, so only using the last one would be fine.
+#    https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/generation/utils.py#L2189
+#    https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L2189
 
 # online search
 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?', 
@@ -56,9 +55,7 @@ def run(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='
         _input[k] = _input[k].to(device)
 
     ori_output = model(**_input)
-    # if 'last_hidden_state' in ori_output:
-    #     ori_output = ori_output['last_hidden_state']
-    # else:
+    
     ori_output = ori_output['logits']
 
     ori_embedding = embedidng_layer(_input['input_ids']).detach()
@@ -82,10 +79,8 @@ def run(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='
 
         _input_['inputs_embeds'] = perturbed_embedding
         outputs_perturbed = model(**_input_)
-        if 'last_hidden_state' in outputs_perturbed:
-            outputs_perturbed = outputs_perturbed['last_hidden_state']
-        else:
-            outputs_perturbed = outputs_perturbed['logits']
+
+        outputs_perturbed = outputs_perturbed['logits']
 
         loss = loss_funt(ori_output, outputs_perturbed)
         loss.backward()