Update pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

@@ -2,7 +2,8 @@ import streamlit as st
 from streamlit_extras.stateful_button import button
 import os 
 import openai
-from transformers import GPT2Tokenizer, GPT2Model, AutoTokenizer, AutoModelForCausalLM
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import pickle
 import torch
 from copy import deepcopy
@@ -23,7 +24,7 @@ st.sidebar.markdown('`Input text`: a sentence or paragraph.')
 st.sidebar.markdown('`Number of replacements`: the number of secret language samples.')
 st.sidebar.markdown('`Steps for searching Secret Langauge`: the steps in the SecretFinding process.')
 st.sidebar.markdown('#### Two methods')
-st.sidebar.markdown('1. Searching secret languages based on models: this method calculates secret languages using [GPT-2](https://huggingface.co/gpt2), [EleutherAI/gpt-neo-1.3B](https://huggingface.co/EleutherAI/gpt-neo-1.3B)') #, [EleutherAI/gpt-neo-2.7B](https://huggingface.co/EleutherAI/gpt-neo-2.7B), [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b), or [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B).')
+st.sidebar.markdown('1. Searching secret languages based on models: this method calculates secret languages using [GPT-2](https://huggingface.co/gpt2-medium), [EleutherAI/gpt-neo-1.3B](https://huggingface.co/EleutherAI/gpt-neo-1.3B)') #, [EleutherAI/gpt-neo-2.7B](https://huggingface.co/EleutherAI/gpt-neo-2.7B), [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b), or [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B).')
 st.sidebar.markdown('2. Use the secret language we found on ALBERT, DistillBERT, and Roberta: this method replaces words directly with the secret language dictionary derived from ALBERT, DistillBERT, and Roberta.')
 
 st.sidebar.markdown('#### Return')
@@ -38,6 +39,12 @@ 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
+'''
+
 # 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?', 
     loss_funt=torch.nn.MSELoss(), lr=1, noise_mask=[1,2], restarts=10, step=100, device = torch.device('cpu'),
@@ -49,10 +56,10 @@ 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']
+    # 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()
     ori_embedding.requires_grad = False
@@ -159,8 +166,8 @@ title = st.text_area('Input text.', 'Which name is also used to describe the Ama
 if option == 'Searching secret languages based on models':
     model_choice = st.selectbox(
         'Which model you would like to use?',
-        # ('GPT-2', "EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B", "EleutherAI/gpt-neox-20b", "EleutherAI/gpt-j-6B")
-        ('GPT-2', "EleutherAI/gpt-neo-1.3B")
+        # ('gpt2-medium', "EleutherAI/gpt-neo-1.3B", "EleutherAI/gpt-neo-2.7B", "EleutherAI/gpt-neox-20b", "EleutherAI/gpt-j-6B")
+        ('gpt2-medium', "EleutherAI/gpt-neo-1.3B")
     )
     _cols = st.columns(2)
     restarts = _cols[0].number_input('Number of replacements.', value=10, min_value=1, step=1, format='%d')
@@ -170,12 +177,9 @@ else:
 
 if button('Tokenize', key='tokenizer'):
     if option == 'Searching secret languages based on models':
-        if model_choice == 'GPT-2':
-            tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-        else:
-            tokenizer = AutoTokenizer.from_pretrained(model_choice)
+        tokenizer = AutoTokenizer.from_pretrained(model_choice)
     else:
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+        tokenizer = AutoTokenizer.from_pretrained('gpt2-medium')
     for key in st.session_state.keys():
         if key not in ['tokenizer', 'start'] and 'tokenizer_' not in key:
             del st.session_state[key]
@@ -211,19 +215,16 @@ if button('Tokenize', key='tokenizer'):
                         chose_indices.append(_index)
         if len(chose_indices):
             if option == 'Searching secret languages based on models':
-                if model_choice == 'GPT-2':
-                    model = GPT2Model.from_pretrained('gpt2')
-                else:
-                    model = AutoModelForCausalLM.from_pretrained(model_choice)
+                model = AutoModelForCausalLM.from_pretrained(model_choice)
                 generator = pipeline('text-generation', model='gpt2')
                 if not platform.system().lower() == 'darwin':
                     generator1 = pipeline('text-generation', model='EleutherAI/gpt-neo-1.3B')
             with st.expander('**Original input text**: '+ title):
                     st.markdown(f'The response of GPT-2 with the prompt :blue[{title}]')
-                    st.markdown('<blockquote>' + generator(title, max_length=30, num_return_sequences=1)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
+                    st.markdown('<blockquote>' + generator(title, max_length=256, num_return_sequences=1)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
                     if not platform.system().lower() == 'darwin':
                         st.markdown(f'The response of EleutherAI/gpt-neo-1.3B with the prompt :blue[{title}]')
-                        st.markdown('<blockquote>' + generator1(title, do_sample=True, min_length=50)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
+                        st.markdown('<blockquote>' + generator1(title, do_sample=True, max_length=256)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
 
                     output_openai = get_codex_response(title)
                     st.markdown(f'The response of [Codex](https://openai.com/blog/openai-codex/) with the prompt :blue[{title}]')
@@ -262,10 +263,10 @@ if button('Tokenize', key='tokenizer'):
             for i in range(restarts):
                 with st.expander(outputs[i]):
                     st.markdown(f'The response of GPT-2 with the prompt :blue[{outputs[i]}]')
-                    st.markdown('<blockquote>' + generator(outputs[i], max_length=30, num_return_sequences=1)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
+                    st.markdown('<blockquote>' + generator(outputs[i], max_length=256, num_return_sequences=1)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
                     if not platform.system().lower() == 'darwin':
                         st.markdown(f'The response of EleutherAI/gpt-neo-1.3B with the prompt :blue[{outputs[i]}]')
-                        st.markdown('<blockquote>' + generator1(outputs[i], do_sample=True, min_length=50)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
+                        st.markdown('<blockquote>' + generator1(outputs[i], do_sample=True, max_length=256)[0]['generated_text'].replace(title, '', 1) + '</blockquote>', unsafe_allow_html=True)
 
                     output_openai = get_codex_response(outputs[i])
                     st.markdown(f'The response of [Codex](https://openai.com/blog/openai-codex/) with the prompt :blue[{outputs[i]}]')