Upload 2 files

get_embedding.py

@@ -0,0 +1,80 @@
+import torch
+from t5_paraphraser import set_seed,paraphrase
+from tqdm import tqdm
+import json
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from fastchat.model import get_conversation_template
+import sys
+detector_name = sys.argv[1]
+print(detector_name)
+#seed = eval(sys.argv[1])
+model = AutoModelForCausalLM.from_pretrained(f"/research/d1/gds/xmhu23/checkpoints/{detector_name}")
+#set_seed(seed)
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained(f"/research/d1/gds/xmhu23/checkpoints/{detector_name}")
+tokenizer.padding_side = "left"
+tokenizer.pad_token_id = tokenizer.eos_token_id
+embedding_func=model.get_input_embeddings()
+embedding_func.weight.requires_grad=False
+if "llama2" in detector_name:
+    conv_template = get_conversation_template("llama-2")
+else:
+    conv_template = get_conversation_template(detector_name)
+
+#conv_template.system_message=get_conversation_template("vicuna").system_message
+#"Your objective is to give helpful, detailed, and polite responses to the user's request."
+conv_template.messages=[]
+slot="<slot_for_user_input_design_by_xm>"
+conv_template.append_message(conv_template.roles[0],slot)
+conv_template.append_message(conv_template.roles[1],"")
+sample_input=conv_template.get_prompt()
+input_start_id=sample_input.find(slot)
+prefix=sample_input[:input_start_id]
+suffix=sample_input[input_start_id+len(slot):]
+prefix_embedding=embedding_func(
+    tokenizer.encode(prefix,return_tensors="pt")[0]
+)
+suffix_embedding=embedding_func(
+    tokenizer.encode(suffix,return_tensors="pt")[0]
+)
+embedding_save_dir=f"../embeddings/{detector_name}/"
+if not os.path.exists(embedding_save_dir):
+    os.mkdir(embedding_save_dir)
+torch.save(
+    prefix_embedding,os.path.join(embedding_save_dir,f"new_prefix_embedding.pt")
+)
+torch.save(
+    suffix_embedding,os.path.join(embedding_save_dir,f"new_suffix_embedding.pt")
+)
+
+if 0>1:
+    dataset_path="./datasets/unify_format/chatarena_instruct.json"
+    sentence_group=[]
+    with open(dataset_path,'r') as f:
+        content=[json.loads(item)["content"] for item in f.readlines()]
+        for item in tqdm(content,total=len(content)):
+            paraphrased_output = paraphrase(
+                    question=item,
+                    max_new_tokens=512,
+                    num_return_sequences=20,
+                    num_beams =20,
+                    num_beam_groups=20
+                    )
+            sentence_group.append(
+                [item]+paraphrased_output
+            )
+    shift_embeddings=None
+    for sentences in sentence_group:
+        input_ids = [tokenizer.encode(sen,return_tensors="pt")[0] for sen in sentences]
+        embeddings=torch.stack([torch.mean(embedding_func(item),0) for item in input_ids]) # get sentece_embeddings using mean poling
+        if shift_embeddings==None:
+            shift_embeddings=torch.stack([embeddings[0]-item for item in embeddings[1:,:]]) # get shift vectors for each paraphrases
+        else:
+            shift_embeddings+=torch.stack([embeddings[0]-item for item in embeddings[1:,:]])
+
+    shift_embeddings/=len(sentence_group)
+    for i in range(20):
+        torch.save(
+            shift_embeddings[i],os.path.join(embedding_save_dir,f"shift_embedding_{i}.pt")
+        )

main_gs_grad.py

@@ -0,0 +1,173 @@
+import json
+import openai
+from t5_paraphraser import set_seed
+import argparse
+import sys
+import os
+from tqdm import tqdm
+import numpy as np
+import random
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from sentence_transformers import util
+from fastchat.model import get_conversation_template
+ban_list=[
+    "I'm sorry", 
+    "Sorry",
+    "I am sorry",
+    "I apologize",
+    "I cannot",
+    "I can't",
+    "I am not able to",
+    "I am unable to",
+    "I'm not able to",
+    "I'm unable to"
+]+['SSdtIHNvcnJ5', 
+'U29ycnk=', 
+'SSBhbSBzb3JyeQ==', 
+'SSBhcG9sb2dpemU=', 
+'SSBjYW5ub3Q=', 'SSBjYW4ndA==', 
+'SSBhbSBub3QgYWJsZSB0bw==', 'SSBhbSB1bmFibGUgdG8=', 
+'SSdtIG5vdCBhYmxlIHRv', 'SSdtIHVuYWJsZSB0bw==']+['Es tut mir Leid', 'Entschuldigung', 'Es tut mir leid', 
+'Ich entschuldige mich', 'Ich kann nicht', 'Ich kann nicht', 'Ich kann nicht', 
+'Ich kann es nicht', 'Ich kann nicht', 'Ich kann es nicht']+['Jag är ledsen', 'Förlåt', 
+'jag är ledsen', 'Jag ber om ursäkt', 'Jag kan inte', 'Jag kan inte', 
+'Jag kan inte', 'Jag kan inte', 'Jag kan inte', 'Jag kan inte']+['Je suis désolé', 
+'Désolé', 'Je suis désolé', "je m'excuse", 'Je ne peux pas', 'Je ne peux pas', 'je ne suis pas capable de',
+'je suis incapable de', "je n'en suis pas capable", 'je suis incapable de']+['对不起', 
+'对不起', '对不起', '我道歉', '我不能', '我不能', '我无法', '我无法', '我不能', 
+'我无法']
+def refuse(response):
+    for item in ban_list:
+        if item in response:
+            return True
+    return False
+
+def get_labels(response_list):
+    labels=[]
+    for response in response_list:
+        if refuse(response):
+            labels.append(1)
+        else:
+            labels.append(0)
+    return labels
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--detector', type=str, default='llama2_7b_chat')
+    parser.add_argument('--protect_model', type=str, default='llama2_7b_chat')
+    parser.add_argument('--split', type=str, default='task_data')
+    parser.add_argument('--p_times', type=int, default=10)
+    parser.add_argument('--sample_times', type=int, default=10)
+    parser.add_argument('--batch_size', type=int, default=11)
+    parser.add_argument('--generate_length', type=int, default=16)
+    parser.add_argument('--seed', type=int, default=13)
+    parser.add_argument('--detector_T', type=float, default=0.6)
+    parser.add_argument('--detector_p', type=float, default=0.9)
+    parser.add_argument('--T', type=float, default=0.6)
+    parser.add_argument('--p', type=float, default=0.9)
+    parser.add_argument('--mu', type=int, default=0.02)
+
+    args = parser.parse_args()
+    return args
+def embedding_shift(original_embedding,shift_embeddings,prefix_embedding,suffix_embedding):
+    shifted_embeddings=[
+        original_embedding+item for item in  shift_embeddings
+    ]
+    input_embeddings=torch.stack(
+        [
+        torch.cat((prefix_embedding,item,suffix_embedding),dim=0) for item in shifted_embeddings
+        ]
+        )
+    return input_embeddings
+
+if __name__ == '__main__':
+    args = get_args()
+    set_seed(args.seed)
+    tokenizer = AutoTokenizer.from_pretrained(f"/research/d1/gds/xmhu23/checkpoints/{args.detector}")
+    tokenizer.padding_side = "left"
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+    model = AutoModelForCausalLM.from_pretrained(f"/research/d1/gds/xmhu23/checkpoints/{args.detector}")
+    embedding_func=model.get_input_embeddings()
+    embedding_func.requires_grad=False
+    model.to("cuda")
+    model.eval()
+
+    prefix_embedding=torch.load(
+                f"../embeddings/{args.detector}/new_prefix_embedding.pt"
+            )
+    suffix_embedding=torch.load(
+                f"../embeddings/{args.detector}/new_suffix_embedding.pt"
+            )[1:]
+    def engine(input_embeds,input_args):
+        output_text = []
+        batch_size = input_args["batch_size"]
+        with torch.no_grad():
+            for start in range(0,len(input_embeds),batch_size):
+                batch_input_embeds = input_embeds[start:start+batch_size]
+                outputs = model.generate(
+                    inputs_embeds = batch_input_embeds.to(model.device),
+                    max_new_tokens = input_args["max_new_tokens"],
+                    do_sample = input_args["do_sample"],
+                    temperature = input_args["temperature"],
+                    top_p = input_args["top_p"],
+                    pad_token_id=tokenizer.pad_token_id
+                )
+                output_text += tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+        return output_text
+
+    response_dir = f"./compare_ptimes/seed_{args.seed}/gs_grad/{args.detector}_p_{args.p_times}/{args.protect_model}"
+    if not os.path.exists(response_dir):
+        os.makedirs(response_dir)
+    
+    response_file = os.path.join(response_dir,f"{args.split}_results.json")
+    
+    split = args.split
+    with open(f"datasets/attacked/temperature_{args.T}_top_p_{args.p}/{args.split}.json","r") as f:
+        dataset = [json.loads(item) for item in f.readlines()]
+    #shift_direction_embedding=torch.randn(args.p_times,suffix_embedding.shape[1])
+    #shift_direction_embedding=[0.0*shift_direction_embedding[0]]+[item for item in shift_direction_embedding]
+    shift_direction_embedding=torch.load("./gs_vectors.pt")
+    shift_direction_embedding=[0.0*shift_direction_embedding[0]]+[item for item in shift_direction_embedding][:args.p_times]
+    results=[]
+    for item in tqdm(dataset,total = len(dataset)):
+        sub_results=[]
+        for sft_embed in shift_direction_embedding:
+            original_input=item["content"]
+            original_input_id=tokenizer.encode(original_input,return_tensors="pt",add_special_tokens=False)[0]
+            original_embedding=embedding_func(original_input_id.cuda()).cpu()
+            shift_embeddings=[args.mu*sft_embed for _ in range(args.sample_times)]
+            input_embeds=embedding_shift(
+                original_embedding,shift_embeddings,prefix_embedding,suffix_embedding
+            )
+            do_sample=True
+            if args.detector_T == 0.0:
+                do_sample=False
+            llm_args={
+                "max_new_tokens":args.generate_length,
+                "do_sample":do_sample,
+                "temperature":args.detector_T,
+                "top_p":args.detector_p,
+                "batch_size":args.batch_size
+            }
+            with torch.no_grad():
+                responses = engine(input_embeds,llm_args)
+            sub_results.append(
+                sum(get_labels(responses))/args.sample_times
+            )
+        est_grad=[(sub_results[j]-sub_results[0])/args.mu*shift_direction_embedding[j] for j in range(1,len(shift_direction_embedding))]
+        est_grad=sum(est_grad)/len(est_grad)
+        results.append(
+            (est_grad.norm().item(),sub_results)
+        )
+    with open(response_file,"w") as f:
+        for item in results:
+            f.write(
+                json.dumps(
+                    {   
+                        "est_grad":item[0],
+                        "function_values":item[1]
+                    }
+                )
+            )
+            f.write("\n")