├── __init__.py ├── logs └── readme.md ├── models └── readme.md ├── config.py ├── LICENSE ├── result verification ├── embed_test.py ├── data_stat.py ├── swr_processor.py └── ner.py ├── simcse_persona.py ├── data_process.py ├── README.md ├── baseline ├── baseline_models.py └── projection.py ├── decode_beam_search.py ├── eval_generation.py ├── decode_beam_search_opt.py ├── attacker_models.py ├── eval_classification.py ├── attacker_evaluation_gpt.py ├── eval_ppl.py ├── attacker.py ├── attacker_opt.py ├── attacker_random_gpt2.py └── attacker_t5.py /__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /logs/readme.md: -------------------------------------------------------------------------------- 1 | This folder is used to store testing results. 2 | -------------------------------------------------------------------------------- /models/readme.md: -------------------------------------------------------------------------------- 1 | This folder is used to store attacker model checkpoints. 2 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | processed_persona = 'data/personachat/processed_persona' 4 | abcd_path = 'data/abcd/abcd_v1.1.json' 5 | 6 | #Don't need to modify 7 | logging_level = logging.INFO 8 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 HKUST-KnowComp 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /result verification/embed_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "0" 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | from transformers import AutoModelForCausalLM,AutoTokenizer,AutoModel 9 | from transformers import pipeline 10 | 11 | 12 | 13 | 14 | class baseline_RNN(nn.Module): 15 | def __init__(self, out_num, num_layers, in_num=1024, hidden_size = 2048): 16 | super(baseline_RNN, self).__init__() 17 | self.rnn = nn.RNN( 18 | input_size=in_num, 19 | hidden_size = 4096, 20 | num_layers=num_layers, 21 | batch_first=True 22 | ) 23 | self.fc1 = nn.Linear(hidden_size, out_num) 24 | 25 | 26 | ''' 27 | Extract representations from GPT-2's tokens as RNN's input 28 | 29 | ''' 30 | self.gpt2_tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 31 | #self.gpt2 = AutoModelForCausalLM.from_pretrained('gpt2') 32 | self.gpt2 = AutoModel.from_pretrained("microsoft/DialoGPT-medium") 33 | # If you find errors related to padding token, uncomment these 2 lines 34 | self.gpt2_tokenizer.add_special_tokens({'pad_token': self.gpt2_tokenizer.eos_token}) 35 | self.gpt2.resize_token_embeddings(len(self.gpt2_tokenizer)) 36 | self.pipe = pipeline('feature-extraction', model=self.gpt2, tokenizer=self.gpt2_tokenizer,device=0) 37 | 38 | 39 | ''' 40 | Obtain token embedding from GPT-2 model we use for IEI. (Make fair comparison with same input embedding) 41 | ''' 42 | def get_token_embedding(self,token_id): 43 | token = self.gpt2_tokenizer.decode(token_id) 44 | embedding = self.pipe(token) # get embedding, [1,1,1024] 45 | embedding = torch.tensor(embedding).cuda() 46 | print(embedding.size()) 47 | embedding = embedding[0,0,:] 48 | print(embedding.size()) 49 | return embedding 50 | 51 | 52 | if __name__ == '__main__': 53 | model = baseline_RNN(out_num = 50257, num_layers= 2).cuda() 54 | 55 | token_id = model.gpt2_tokenizer.encode('this is a test') 56 | model.get_token_embedding(token_id) -------------------------------------------------------------------------------- /result verification/data_stat.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.append('..') 3 | from simcse_persona import get_persona_dict 4 | import baseline_models 5 | from datasets import load_dataset 6 | 7 | import argparse 8 | 9 | def get_sent_list(config): 10 | dataset = config['dataset'] 11 | data_type = config['data_type'] 12 | if dataset == 'personachat': 13 | sent_list = get_personachat_data(data_type) 14 | return sent_list 15 | elif dataset == 'qnli': 16 | sent_list = get_qnli_data(data_type) 17 | return sent_list 18 | else: 19 | print('Name of dataset only supports: personachat or qnli') 20 | sys.exit(-1) 21 | 22 | def get_personachat_data(data_type): 23 | sent_list = get_persona_dict(data_type=data_type) 24 | return sent_list 25 | 26 | def get_qnli_data(data_type): 27 | dataset = load_dataset('glue', 'qnli', cache_dir="/home/hlibt/embed_rev/data", split=data_type) 28 | sentence_list = [] 29 | for i, d in enumerate(dataset): 30 | sentence_list.append(d['question']) 31 | sentence_list.append(d['sentence']) 32 | return sentence_list 33 | 34 | def data_stat(): 35 | sent_list_train = get_personachat_data('train') 36 | sent_list_dev = get_personachat_data('dev') 37 | sent_list_test = get_personachat_data('test') 38 | sent_list = [] 39 | sent_list.extend(sent_list_train) 40 | sent_list.extend(sent_list_dev) 41 | sent_list.extend(sent_list_test) 42 | print('===PersonaChat stat===') 43 | print('===Train===') 44 | print_stat(sent_list_train) 45 | print('===Dev===') 46 | print_stat(sent_list_dev) 47 | print('===Test===') 48 | print_stat(sent_list_test) 49 | print_stat(sent_list) 50 | sent_list = get_qnli_data('train') 51 | sent_list.extend(get_qnli_data('test')) 52 | print('===QNLI stat===') 53 | print('===Train===') 54 | print_stat(get_qnli_data('train')) 55 | print('===Test===') 56 | print_stat(get_qnli_data('test')) 57 | print_stat(sent_list) 58 | 59 | 60 | def print_stat(sent_list): 61 | sent_len = len(sent_list) 62 | sent_len_list = [len(i.split()) for i in sent_list] 63 | avg_len = sum(sent_len_list) / sent_len 64 | print(f'Number of sentences: {sent_len}') 65 | print(f'Avg_length of sentences: {avg_len}') 66 | 67 | if __name__ == '__main__': 68 | data_stat() 69 | -------------------------------------------------------------------------------- /simcse_persona.py: -------------------------------------------------------------------------------- 1 | import os 2 | import config 3 | #os.environ["CUDA_VISIBLE_DEVICES"] = "0" 4 | import sys 5 | import datasets 6 | from datasets import load_dataset 7 | import torch 8 | from scipy.spatial.distance import cosine 9 | from transformers import AutoModel, AutoTokenizer 10 | from torch.utils.data import DataLoader, Dataset 11 | import json 12 | from pprint import pprint 13 | class BookCorpus(Dataset): 14 | def __init__(self, data, tokenizer): 15 | self.data = data 16 | self.tokenizer = tokenizer 17 | 18 | def __len__(self): 19 | return len(self.data) 20 | 21 | def __getitem__(self, index): 22 | 23 | text = self.data[index] 24 | 25 | return text 26 | 27 | def collate(self, unpacked_data): 28 | return unpacked_data 29 | 30 | def process_data(data,batch_size,device): 31 | tokenizer = AutoTokenizer.from_pretrained("princeton-nlp/sup-simcse-bert-large-uncased") 32 | model = AutoModel.from_pretrained("princeton-nlp/sup-simcse-bert-large-uncased").to(device) 33 | dataset = BookCorpus(data, tokenizer) 34 | dataloader = DataLoader(dataset=dataset, 35 | shuffle=True, 36 | batch_size=batch_size, 37 | collate_fn=dataset.collate) 38 | data_dict = {} 39 | data_dict['text'] = [] 40 | data_dict['embedding'] = [] 41 | with torch.no_grad(): 42 | for idx,batch_text in enumerate(dataloader): 43 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 44 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 45 | embeddings = embeddings.detach().cpu() 46 | data_dict['text'].extend(batch_text) 47 | data_dict['embedding'].extend(embeddings) 48 | print(f'{idx} batch done with {idx*batch_size} samples') 49 | return data_dict 50 | 51 | def get_processed_persona(kind,processed_persona_path,require_label = True): 52 | #processed_persona_path = config.processed_persona 53 | if(require_label): 54 | path = processed_persona_path + '/%s_merged_shuffle.txt' % kind 55 | else: 56 | path = processed_persona_path + '/%s.txt' % kind 57 | with open(path, 'r') as f: 58 | data = json.load(f) 59 | return data 60 | 61 | def process_persona(data): 62 | ''' 63 | get only list of texts for batch training for 64 | ''' 65 | sentence_list = [] 66 | for i,dict_i in enumerate(data): 67 | conv = dict_i['conv'] 68 | sentence_list.extend(conv) 69 | 70 | return sentence_list 71 | 72 | 73 | def get_persona_dict(data_type): 74 | processed_persona_path = config.processed_persona 75 | data = get_processed_persona(data_type,processed_persona_path) 76 | processed_data = process_persona(data) 77 | 78 | 79 | return processed_data 80 | # Import our models. The package will take care of downloading the models automatically 81 | if __name__ == '__main__': 82 | processed_persona_path = config.processed_persona 83 | #train_data = get_processed_persona('train') 84 | #val_data = get_processed_persona('dev') 85 | test_data = get_processed_persona('test') 86 | processed_test = process_persona(test_data) 87 | 88 | device = torch.device("cuda") 89 | batch_size = 32 90 | 91 | val_dict = process_data(processed_test,batch_size,device) 92 | torch.save(val_dict, '/data/hlibt/gradient_leakage/pytorch/data/personachat_processed/hidden_test_simcse.pt') 93 | 94 | 95 | 96 | -------------------------------------------------------------------------------- /result verification/swr_processor.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 3 | 4 | import torch, gc 5 | torch.cuda.empty_cache() 6 | 7 | import sys 8 | sys.path.append('..') 9 | 10 | import json 11 | 12 | from nltk.corpus import stopwords 13 | 14 | 15 | # tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 16 | stopwords_list = stopwords.words('english') 17 | for w in ['!',',','.','?','-s','-ly','','s']: 18 | stopwords_list.append(w) 19 | stopwords_set = set(stopwords_list) 20 | # for NN or RNN result 21 | file_list = ['path to NN OR RNN result'] 22 | 23 | 24 | for file in file_list: 25 | 26 | gt_tokens = [] 27 | pred_tokens = [] 28 | with open(file) as f: 29 | print(file) # 也就是文本IO类型 30 | result = json.load(f) 31 | 32 | 33 | #print(result) 34 | gt_count = 0 35 | pred_count = 0 36 | gt_non_stopword_count = 0 37 | pred_non_stopword_count = 0 38 | for i in result: 39 | # print(i) 40 | gt_list = [item.strip() for item in i['gt']] 41 | gt_count += len(gt_list) 42 | gt_set = set(gt_list) 43 | gt_non_stopword_count += len(gt_set - stopwords_set) 44 | # print("gt_tokens") 45 | # print(gt_tokens) 46 | pred_list = [item.strip() for item in i['pred']] 47 | pred_count += len(pred_list) 48 | pred_set = set(pred_list) 49 | pred_non_stopword_count += len(pred_set - stopwords_set) 50 | # print("gt static") 51 | # print(gt_non_stopword_count/gt_count) 52 | # print("pred static") 53 | # print(pred_non_stopword_count/pred_count) 54 | # print("pred_tokens") 55 | # print(pred_tokens) 56 | 57 | swr_gt = 1 - gt_non_stopword_count/gt_count 58 | print("gt_count:", end = '') 59 | print(gt_count) 60 | print("swr_gt:", end = '') 61 | print(swr_gt) 62 | 63 | swr_pred = 1 - pred_non_stopword_count/pred_count 64 | print("pred_count:", end = '') 65 | print(pred_count) 66 | print("swr_pred:", end = '') 67 | print(swr_pred) 68 | 69 | 70 | # for GPT result 71 | file_list = ['path to GPT result'] 72 | 73 | for file in file_list: 74 | gt_count = 0 75 | pred_count = 0 76 | gt_non_stopword_count = 0 77 | pred_non_stopword_count = 0 78 | gt_sentence_list = [] 79 | pred_sentence_list = [] 80 | 81 | with open(file) as f: 82 | print(file) 83 | # 也就是文本IO类型 84 | result = json.load(f) 85 | for sentence in result['gt']: 86 | # gt_sentence_list.append([tokenizer.encode(item) for item in sentence]) 87 | gt_sentence_list = sentence.split() 88 | gt_count += len(gt_sentence_list) 89 | gt_set = set(gt_sentence_list) 90 | gt_non_stopword_count += len(gt_set - stopwords_set) 91 | # print("gt static") 92 | # print(gt_non_stopword_count / gt_count) 93 | for sentence in result['pred']: 94 | # pred_sentence_list.append([tokenizer.encode(item) for item in sentence]) 95 | pred_sentence_list = sentence.split() 96 | pred_count += len(pred_sentence_list) 97 | pred_set = set(pred_sentence_list) 98 | pred_non_stopword_count += len(pred_set - stopwords_set) 99 | # print("pred static") 100 | # print(pred_non_stopword_count / pred_count) 101 | 102 | swr_gt = 1 - gt_non_stopword_count/gt_count 103 | 104 | print("gt_count:", end='') 105 | print(gt_count) 106 | print("swr_gt:", end='') 107 | print(swr_gt) 108 | 109 | swr_pred = 1 - pred_non_stopword_count/pred_count 110 | print("pred_count:", end = '') 111 | print(pred_count) 112 | print("swr_pred:", end = '') 113 | print(swr_pred) 114 | print() 115 | -------------------------------------------------------------------------------- /data_process.py: -------------------------------------------------------------------------------- 1 | from simcse_persona import get_persona_dict 2 | from datasets import load_dataset 3 | from pprint import pprint 4 | import config 5 | import json 6 | import sys 7 | 8 | abcd_path = config.abcd_path 9 | 10 | ''' 11 | list of supported datasets: 12 | ['personachat'.'qnli','mnli','sst2','wmt16','multi_woz','abcd'] 13 | ''' 14 | 15 | 16 | def get_sent_list(config): 17 | dataset = config['dataset'] 18 | data_type = config['data_type'] 19 | if dataset == 'personachat': 20 | sent_list = get_personachat_data(data_type) 21 | return sent_list 22 | elif dataset == 'qnli': 23 | sent_list = get_qnli_data(data_type) 24 | return sent_list 25 | elif dataset == 'mnli': 26 | sent_list = get_mnli_data(data_type) 27 | return sent_list 28 | elif dataset == 'sst2': 29 | sent_list = get_sst2_data(data_type) 30 | return sent_list 31 | elif dataset == 'wmt16': 32 | sent_list = get_wmt16_data(data_type) 33 | return sent_list 34 | elif dataset == 'multi_woz': 35 | sent_list = get_multi_woz_data(data_type) 36 | return sent_list 37 | elif dataset == 'abcd':#abcd 38 | sent_list = get_abcd_data(data_type) 39 | return sent_list 40 | else: 41 | print('Name of dataset only supports: personachat or qnli') 42 | sys.exit(-1) 43 | 44 | 45 | def get_personachat_data(data_type): 46 | sent_list = get_persona_dict(data_type=data_type) 47 | return sent_list 48 | 49 | 50 | def get_qnli_data(data_type): 51 | if(data_type == 'dev'): 52 | data_type = 'validation' 53 | dataset = load_dataset('glue', 'qnli', cache_dir="data/", split=data_type) 54 | sentence_list = [] 55 | for i, d in enumerate(dataset): 56 | sentence_list.append(d['question']) 57 | sentence_list.append(d['sentence']) 58 | return sentence_list 59 | 60 | def get_mnli_data(data_type): 61 | if(data_type == 'test'): 62 | data_type = 'test_matched' 63 | if(data_type == 'dev'): 64 | data_type = 'validation_matched' 65 | dataset = load_dataset('glue', 'mnli', cache_dir="data/", split=data_type) 66 | sentence_list = [] 67 | for i, d in enumerate(dataset): 68 | sentence_list.append(d['premise']) 69 | sentence_list.append(d['hypothesis']) 70 | return sentence_list 71 | 72 | def get_sst2_data(data_type): 73 | if(data_type == 'dev'): 74 | data_type = 'validation' 75 | dataset = load_dataset('glue', 'sst2', cache_dir="data/", split=data_type) 76 | sentence_list = [] 77 | for i, d in enumerate(dataset): 78 | sentence_list.append(d['sentence']) 79 | return sentence_list 80 | 81 | ## translation dataset 82 | def get_wmt16_data(data_type): 83 | if(data_type == 'dev'): 84 | data_type = 'validation' 85 | dataset = load_dataset('wmt16', 'cs-en', cache_dir="data/", split=data_type) 86 | sentence_list = [] 87 | for i, d in enumerate(dataset): 88 | 89 | #pprint(d) 90 | sentence_list.append(d['translation']['en']) 91 | return sentence_list 92 | 93 | ### multi_woz 94 | ## translation dataset 95 | def get_multi_woz_data(data_type): 96 | if(data_type == 'dev'): 97 | data_type = 'validation' 98 | dataset = load_dataset('multi_woz_v22', 'v2.2', cache_dir="data/", split=data_type) 99 | sentence_list = [] 100 | for i, d in enumerate(dataset): 101 | s = d['turns']['utterance'] 102 | sentence_list.extend(s) 103 | return sentence_list 104 | def get_abcd_data(data_type,path = abcd_path): 105 | 106 | #abcd_path = config.abcd_path 107 | with open(path, 'r') as f: 108 | dataset = json.load(f) 109 | dataset = dataset[data_type] 110 | #pprint(data[0]) 111 | sentence_list = [] 112 | for i, d in enumerate(dataset): 113 | s = d['original'] 114 | for sent in s: 115 | if(sent[0] != 'action'): 116 | sentence_list.append(sent[1]) 117 | # print('==='*50) 118 | # pprint(s) 119 | # pprint(sentence_list) 120 | # if i>=2: 121 | # sys.exit(-1) 122 | return sentence_list 123 | 124 | 125 | if __name__ == '__main__': 126 | 127 | config = {} 128 | config['dataset'] = 'abcd' 129 | config['data_type'] = 'test' 130 | sent_list = get_sent_list(config) 131 | pprint(sent_list[:10]) 132 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GEIA 2 | Code for Findings-ACL 2023 paper: Sentence Embedding Leaks More Information than You Expect: Generative Embedding Inversion Attack to Recover the Whole Sentence 3 | 4 | ### Package Dependencies 5 | * numpy 6 | * pytorch 1.10.2 7 | * sentence_transformers 2.2.0 8 | * transformers 4.xx.x 9 | * simcse 0.4 10 | * datasets 11 | 12 | ### Data Preparation 13 | We upload PC data under the ```data/``` folder. 14 | The ABCD dataset we experimented can be found in https://drive.google.com/file/d/1oIo8P0Y8X9DTeEfOA1WUKq8Uix9a_Pte/view?usp=sharing. 15 | For other datasets, we use ```datasets``` package to download and store them, so you can run our code directly. 16 | 17 | ### Baseline Attackers 18 | **You need to set up arguments properly before running codes**: 19 | ```python projection.py``` 20 | 21 | * --model_dir: Attacker model path from Huggingface (like 'gpt2-large' and 'microsoft/DialoGPT-xxxx') or local model checkpoints. 22 | * --num_epochs: Training epoches. 23 | * --batch_size: Batch_size #. 24 | * --dataset: Name of the dataset including personachat, qnli, mnli, sst2, wmt16, multi_woz and abcd. 25 | * --data_type: Train or test. 26 | * --embed_model: The victim model you wish to attack. We currently support sentence-bert models and huggingface models, you may refer to our model_cards dictionary in ```projection.py``` for more information. 27 | * --model_type: NN or RNN 28 | 29 | By running: 30 | ```python projection.py``` 31 | You will train your own baseline model and evaluate it. If you want to just train or eval a certain model, check the last four lines of ```projection.py``` and disable the corresponding codes. 32 | 33 | ### GIEA 34 | 35 | #### GPT-2 Attacker 36 | **You need to set up arguments properly before running codes**: 37 | ```python attacker.py``` 38 | 39 | * --model_dir: Attacker model path from Huggingface (like 'gpt2-large' and 'microsoft/DialoGPT-xxxx') or local model checkpoints. 40 | * --num_epochs: Training epoches. 41 | * --batch_size: Batch_size #. 42 | * --dataset: Name of the dataset including personachat, qnli, mnli, sst2, wmt16, multi_woz and abcd. 43 | * --data_type: train or dev or test. 44 | * --embed_model: The victim model you wish to attack. We currently support sentence-bert models and huggingface models, you may refer to our model_cards dictionary in ```attacker.py``` for more information. 45 | * --decode: Decoding algorithm. We currently implement beam and sampling based decoding. 46 | 47 | You should train the attacker on training data at first, then test your attacker on the test data to obtain test logs. Then you can evaluate attack performance on test logs by changing model_dir to your trained attcker and data_type to test. 48 | 49 | If you want to train a randomly initialized GPT-2 attacker, after setting the arguments, run: 50 | ```python attacker_random_gpt2.py``` 51 | 52 | #### Other Attackers 53 | Due to the fact that different decoders have different implementaions, we use separate py files for each model (the decoding implementations also differ). 54 | 55 | If you want to try out opt as the attacker model, run: 56 | ```python attacker_opt.py``` 57 | 58 | If you want to try out t5 as the attacker model, run: 59 | ```python attacker_t5.py``` 60 | 61 | ### Evaluation 62 | **You need to make sure the test reuslt paths is set inside the 'eval_xxx.py' files.** 63 | 64 | To obtain classification performance, run: 65 | ```python eval_classification.py``` 66 | 67 | To obtain generation performance, run: 68 | ```python eval_generation.py``` 69 | 70 | To calculate perplexity, you need to set the LM to caluate PPL, run: 71 | ```python eval_ppl.py``` 72 | 73 | ### Citation 74 | 75 | Please kindly cite the following paper if you found our method and resources helpful! 76 | 77 | ``` 78 | @inproceedings{li-etal-2023-sentence, 79 | title = "Sentence Embedding Leaks More Information than You Expect: Generative Embedding Inversion Attack to Recover the Whole Sentence", 80 | author = "Li, Haoran and 81 | Xu, Mingshi and 82 | Song, Yangqiu", 83 | booktitle = "Findings of the Association for Computational Linguistics: ACL 2023", 84 | month = jul, 85 | year = "2023", 86 | address = "Toronto, Canada", 87 | publisher = "Association for Computational Linguistics", 88 | url = "https://aclanthology.org/2023.findings-acl.881", 89 | doi = "10.18653/v1/2023.findings-acl.881", 90 | pages = "14022--14040", 91 | } 92 | ``` 93 | 94 | ### Miscellaneous 95 | 96 | Please send any questions about the code and/or the algorithm to hlibt@connect.ust.hk 97 | 98 | 99 | 100 | 101 | 102 | 103 | -------------------------------------------------------------------------------- /baseline/baseline_models.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel 8 | from transformers import pipeline 9 | from transformers import AdamW, get_linear_schedule_with_warmup 10 | 11 | ''' 12 | eval = False retrun logits for stable nn.BCEWithLogitsLoss() 13 | eval = True reutrun 0~1 probs 14 | ''' 15 | 16 | 17 | class baseline_NN(nn.Module): 18 | def __init__(self, out_num, in_num=1024): 19 | super(baseline_NN, self).__init__() 20 | self.fc1 = nn.Linear(in_num, out_num) 21 | self.embed_dim = in_num 22 | # self.act = F.softmax() 23 | 24 | def forward(self, x, eval=False): 25 | # x should be of shape (?,1024) according to embedding output 26 | assert (x.dim() == 2) 27 | # print(f'x.size(): {x.size()}. embed_dim:{self.embed_dim}') 28 | assert (x.size()[1] == self.embed_dim) 29 | out = self.fc1(x) 30 | if (eval): 31 | m = nn.Sigmoid() 32 | out = m(out) 33 | return out 34 | 35 | 36 | hidden_size = 512 37 | num_layers = 1 38 | 39 | 40 | class baseline_RNN(nn.Module): 41 | def __init__(self, out_num, in_num=1024): 42 | super(baseline_RNN, self).__init__() 43 | self.rnn = nn.GRUCell( 44 | 1024, 45 | hidden_size, 46 | # num_layers, 47 | # batch_first=True 48 | ) 49 | self.fc1 = nn.Linear(hidden_size, out_num) 50 | self.fc2 = nn.Linear(768, 1024) 51 | self.fc3 = nn.Linear(1024, 512) 52 | self.fc4 = nn.Linear(2048, 512) 53 | self.m = nn.Softmax(dim=1) 54 | ''' 55 | Extract representations from GPT-2's tokens as RNN's input 56 | 57 | ''' 58 | self.gpt2_tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 59 | # self.gpt2 = AutoModelForCausalLM.from_pretrained('microsoft/DialoGPT-medium') ### seems huggingface has bug for feature extraction with this one, return 50257 as embedding (wrong) 60 | self.gpt2 = AutoModel.from_pretrained("microsoft/DialoGPT-medium") 61 | # If you find errors related to padding token, uncomment these 2 lines 62 | ####self.gpt2_tokenizer.add_special_tokens({'pad_token': self.gpt2_tokenizer.eos_token}) 63 | ####self.gpt2.resize_token_embeddings(len(self.gpt2_tokenizer)) 64 | self.pipe = pipeline('feature-extraction', model=self.gpt2, tokenizer=self.gpt2_tokenizer, device=0) 65 | 66 | ''' 67 | Obtain token embedding from GPT-2 model we use for IEI. (Make fair comparison with same input embedding) 68 | ''' 69 | 70 | def get_token_embedding(self, token_id): 71 | token = self.gpt2_tokenizer.decode(token_id) 72 | embedding = self.pipe(token) # get embedding, [1,1,1024] 73 | embedding = torch.tensor(embedding).cuda() 74 | embedding = embedding[0, 0, :] 75 | return embedding 76 | 77 | ''' 78 | work flow: 79 | 0. initialize GRU cell(structure and parameters) 80 | 1. embedding is a embedding sentence, check its dimension 81 | 2. for loop process x, in each loop: 82 | 2.1 add x in to GRU cell 83 | 2.2 get output type(?) 84 | 2.3 add fc1 to output change its dimension to 1024 or 768 85 | 2.4 add a softmax and do argmax 86 | 2.4 use its token id and h_x as the next input 87 | 2.5 use cross_entropy loss to to update this model (output cross_entropy) 88 | ''' 89 | 90 | def forward(self, sentences_embedding, hx, batch_label, eval=False): 91 | ''' 92 | optimizer = AdamW(self.parameters(), 93 | lr=3e-5, 94 | eps=1e-06) 95 | 96 | ''' 97 | # zero_tensor = torch.zeros(64, 50257).cuda() 98 | 99 | output = [] 100 | loss_list = [] 101 | 102 | if sentences_embedding.size()[1] == 768: 103 | sentences_embedding = self.fc2(sentences_embedding) 104 | 105 | predicted_token_index_batch_list = [] 106 | out_exclued_predicted = None 107 | for i in range(10): 108 | hx = self.rnn(sentences_embedding, hx) 109 | output.append(hx) 110 | out = self.fc1(output[i]) 111 | out_clone = out.clone().detach() 112 | out_clone = self.m(out_clone) 113 | batch_label_clone = batch_label.clone().detach() 114 | if (predicted_token_index_batch_list != []): 115 | predicted_token_index_batch_tensor = torch.stack(predicted_token_index_batch_list, 1).cuda() 116 | 117 | if predicted_token_index_batch_tensor != None: 118 | out_exclued_predicted = out_clone.scatter_(1, predicted_token_index_batch_tensor, value=0) 119 | 120 | batch_label_clone = batch_label_clone.scatter_(1, predicted_token_index_batch_tensor, value=0) 121 | # torch.Size([]) 122 | if out_exclued_predicted is not None: 123 | token_index_batch = torch.argmax(out_exclued_predicted, dim=1).cuda() 124 | else: 125 | token_index_batch = torch.argmax(out, dim=1).cuda() 126 | predicted_token_index_batch_list.append(token_index_batch.detach()) 127 | ''' 128 | token_embedding_batch_list = [] 129 | for index in range(64): 130 | token_embedding_batch_list.append(self.get_token_embedding(token_index_batch[index])) 131 | ''' 132 | ''' 133 | predicted_embedding = torch.stack(token_embedding_batch_list, 0).cuda() 134 | predicted_embedding = self.fc3(predicted_embedding) 135 | prev_embedding = self.fc4(embedding) 136 | embedding = torch.cat((sentences_embedding, predicted_embedding, prev_embedding), 1) 137 | ''' 138 | 139 | if eval == False: 140 | # loss = F.cross_entropy(out, batch_label) 141 | pred_out = out.softmax(dim=-1).log().double() 142 | loss = F.kl_div(pred_out, batch_label_clone.softmax(dim=-1), reduction='sum') 143 | target_index = torch.nonzero(batch_label_clone).cpu().numpy().tolist() 144 | 145 | loss_list.append(loss) 146 | 147 | if eval == False: 148 | batch_final_loss = loss_list[0] + loss_list[1] + loss_list[2] + loss_list[3] + loss_list[4] + \ 149 | loss_list[5] + loss_list[6] + loss_list[7] + loss_list[8] + loss_list[9] 150 | 151 | print("batch_final_loss:", end='') 152 | print(batch_final_loss) 153 | 154 | return batch_final_loss 155 | else: 156 | predicted_token_index_batch_tensor = torch.stack(predicted_token_index_batch_list, 1).cuda() 157 | target_index = torch.nonzero(batch_label).cpu().numpy().tolist() 158 | return predicted_token_index_batch_tensor 159 | -------------------------------------------------------------------------------- /decode_beam_search.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Source code modified from https://github.com/budzianowski/PyTorch-Beam-Search-Decoding/blob/master/decode_beam.py 3 | implementation of beam search on GPT-2's logits 4 | ''' 5 | 6 | import operator 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from queue import PriorityQueue 11 | import sys 12 | 13 | 14 | class BeamSearchNode(object): 15 | def __init__(self, hiddenstate, previousNode, wordId, logProb, length): 16 | ''' 17 | :param hiddenstate: 18 | :param previousNode: 19 | :param wordId: 20 | :param logProb: 21 | :param length: 22 | ''' 23 | self.h = hiddenstate 24 | self.prevNode = previousNode 25 | self.wordid = wordId 26 | self.logp = logProb 27 | self.leng = length 28 | 29 | def eval(self, alpha=1.0): 30 | reward = 0 31 | # Add here a function for shaping a reward 32 | 33 | return self.logp / float(self.leng - 1 + 1e-6) + alpha * reward 34 | 35 | def __lt__(self, x): 36 | if(self.eval() < x.eval()): 37 | return True 38 | 39 | else: 40 | return False 41 | 42 | 43 | 44 | def beam_decode_sentence(hidden_X, config,num_generate=1, beam_size = 5, batch_size = 1): 45 | ''' 46 | generate a sentence based on beam search 47 | :param hidden_X: hidden_X of sentence embedding (1024) with/without projection 48 | :param model: GPT-2 model 49 | :param tokenizer: GPT-2 tokenizer 50 | :return: decoded_batch 51 | ''' 52 | #SOS_token = tokenizer.encode("<|endoftext|>") 53 | beam_width = beam_size 54 | topk = num_generate # how many sentence do you want to generate 55 | 56 | past = None 57 | model = config['model'] 58 | tokenizer =config['tokenizer'] 59 | eos = [tokenizer.encode(tokenizer.eos_token)] 60 | EOS_token = eos 61 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X, 0) 62 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X_unsqueeze, 0) #[1,1,embed_dim] 63 | 64 | decoded_batch = [] 65 | 66 | # decoding goes sentence by sentence 67 | for idx in range(batch_size): 68 | 69 | # Number of sentence to generate 70 | endnodes = [] 71 | number_required = min((topk + 1), topk - len(endnodes)) 72 | 73 | # starting node - hidden vector, previous node, word id, logp, length hiddenstate, previousNode, wordId, logProb, length 74 | node = BeamSearchNode(past, None, torch.tensor([[220]]).cuda(), 0, 1) # 220 refers to single space ' ' on GPT 75 | nodes = PriorityQueue() 76 | 77 | # start the queue 78 | nodes.put((-node.eval(), node)) 79 | qsize = 1 80 | 81 | # start beam search 82 | for text_len in range(50): 83 | # give up when decoding takes too long 84 | if qsize > 2000: break 85 | 86 | # fetch the best node 87 | try: 88 | score, n = nodes.get() 89 | except: 90 | print('Cannot get nodes') 91 | while not nodes.empty(): 92 | next_item = nodes.get() 93 | print(next_item) 94 | prev_input = n.wordid 95 | past = n.h 96 | 97 | if n.wordid.item() == EOS_token[0] and n.prevNode != None: 98 | endnodes.append((score, n)) 99 | # if we reached maximum # of sentences required 100 | if len(endnodes) >= number_required: 101 | break 102 | else: 103 | print('continue') 104 | continue 105 | # decode for one step using decoder 106 | if(text_len == 0): 107 | logits, past = model(inputs_embeds=hidden_X_unsqueeze,past_key_values = past,return_dict=False) 108 | 109 | else: 110 | logits, past = model(prev_input,past_key_values = past, attention_mask=None, return_dict=False) 111 | logits = logits[:, -1, :] 112 | probs = torch.softmax(logits, dim=-1) 113 | # PUT HERE REAL BEAM SEARCH OF TOP 114 | log_prob, indexes = torch.topk(probs, beam_width) 115 | nextnodes = [] 116 | 117 | for new_k in range(beam_width): 118 | decoded_t = indexes[0][new_k].view(1, -1) 119 | log_p = log_prob[0][new_k].item() 120 | #### hiddenstate, previousNode, wordId, logProb, length 121 | node = BeamSearchNode(past, n, decoded_t, n.logp + log_p, n.leng + 1) 122 | score = -node.eval() 123 | nextnodes.append((score, node)) 124 | 125 | # put them into queue 126 | for i in range(len(nextnodes)): 127 | score, nn = nextnodes[i] 128 | try: 129 | nodes.put((score, nn)) 130 | except: 131 | print('Cannot put nodes') 132 | print(score) 133 | print(nn) 134 | # increase qsize 135 | qsize += len(nextnodes) - 1 136 | # for loop ends here 137 | # choose nbest paths, back trace them 138 | if len(endnodes) == 0: 139 | endnodes = [nodes.get() for _ in range(topk)] 140 | 141 | utterances = [] 142 | text = [] 143 | for score, n in sorted(endnodes, key=operator.itemgetter(0)): 144 | utterance = [] 145 | utterance.append(n.wordid.item()) 146 | # back trace 147 | while n.prevNode != None: 148 | n = n.prevNode 149 | utterance.append(n.wordid.item()) 150 | 151 | utterance = utterance[::-1] 152 | utterances.append(utterance) 153 | decode_process = tokenizer.decode(utterance[1:-1]) 154 | text.append(decode_process) 155 | decoded_batch.append(utterances) 156 | 157 | return text 158 | 159 | 160 | def greedy_decode(decoder_hidden, encoder_outputs, target_tensor): 161 | ''' 162 | :param target_tensor: target indexes tensor of shape [B, T] where B is the batch size and T is the maximum length of the output sentence 163 | :param decoder_hidden: input tensor of shape [1, B, H] for start of the decoding 164 | :param encoder_outputs: if you are using attention mechanism you can pass encoder outputs, [T, B, H] where T is the maximum length of input sentence 165 | :return: decoded_batch 166 | ''' 167 | 168 | batch_size, seq_len = target_tensor.size() 169 | decoded_batch = torch.zeros((batch_size, MAX_LENGTH)) 170 | decoder_input = torch.LongTensor([[SOS_token] for _ in range(batch_size)], device=device) 171 | 172 | for t in range(MAX_LENGTH): 173 | decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_outputs) 174 | 175 | topv, topi = decoder_output.data.topk(1) # get candidates 176 | topi = topi.view(-1) 177 | decoded_batch[:, t] = topi 178 | 179 | decoder_input = topi.detach().view(-1, 1) 180 | 181 | return decoded_batch -------------------------------------------------------------------------------- /eval_generation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | os.environ["CUDA_VISIBLE_DEVICES"] = "2" 4 | #sys.path.append('..') 5 | #from text_eval import punctuation_remove 6 | import evaluate 7 | import json 8 | import nltk 9 | from sentence_transformers import SentenceTransformer, util 10 | import numpy as np 11 | import torch 12 | from transformers import GPT2LMHeadModel, GPT2TokenizerFast 13 | from tqdm import tqdm 14 | from evaluate import load 15 | from ppl import calucate_ppl 16 | import editdistance 17 | import string 18 | 19 | 20 | model = SentenceTransformer('sentence-t5-xxl') 21 | 22 | rouge = evaluate.load('rouge') 23 | device = "cuda" 24 | model = model.to(device) 25 | model.eval() 26 | # model_id = "gpt2-large" 27 | 28 | 29 | 30 | 31 | #model = GPT2LMHeadModel.from_pretrained(model_id).to(device) 32 | #tokenizer = GPT2TokenizerFast.from_pretrained(model_id) 33 | perplexity = load("perplexity", module_type="metric") 34 | 35 | #self training GPT-2 for PPL evaluation 36 | # ppl_model = GPT2LMHeadModel.from_pretrained("gpt_large_persona") 37 | # device = torch.device("cuda") 38 | # ppl_model = ppl_model.to(device) 39 | # ppl_model.eval() 40 | # model = model.to(device) 41 | 42 | # remove punctuation from list of sentences 43 | def punctuation_remove(sent_list): 44 | removed_list = [] 45 | for sent in sent_list: 46 | word_list = [] 47 | for word in sent.split(): 48 | word_strip = word.strip(string.punctuation) 49 | if word_strip: # cases for not empty string 50 | word_list.append(word_strip) 51 | removed_sent = ' '.join(word_list) 52 | removed_list.append(removed_sent) 53 | return removed_list 54 | 55 | 56 | def read_gpt(path): 57 | with open(path) as f: 58 | data = json.load(f) 59 | return data 60 | 61 | def get_ppl(data,gpt_train= True): 62 | gt = data['gt'] 63 | pred = data["pred"] 64 | if(gpt_train): 65 | ppl_gt,var_gt,ppl_pred,var_pred = calucate_ppl(gt,pred,ppl_model) 66 | print(f"GT: Validation Perplexity: {ppl_gt} Variance: {var_gt}") 67 | print(f"PRED: Validation Perplexity: {ppl_pred} Variance: {var_pred}") 68 | else: 69 | results_pred = perplexity.compute(model_id=model_id, 70 | add_start_token=True, 71 | predictions=pred) 72 | results_gt = perplexity.compute(model_id=model_id, 73 | add_start_token=True, 74 | predictions=gt) 75 | print(f'results_pred: {results_pred["mean_perplexity"]}') 76 | print(f'results_gt: {results_gt["mean_perplexity"]}') 77 | 78 | 79 | def get_rouge(data): 80 | gt = data["gt"] 81 | pred = data["pred"] 82 | results = rouge.compute(predictions=pred,references=gt) 83 | print(results) 84 | 85 | def get_bleu(data): 86 | gt = data['gt'] 87 | pred = data["pred"] 88 | cands_list_bleu = [sentence.split() for sentence in pred] 89 | refs_list_bleu = [[sentence.split()] for sentence in gt] 90 | bleu_score = nltk.translate.bleu_score.corpus_bleu(refs_list_bleu, cands_list_bleu) 91 | bleu_score_1 = nltk.translate.bleu_score.corpus_bleu(refs_list_bleu, cands_list_bleu,weights=(1, 0, 0, 0)) 92 | bleu_score_2 = nltk.translate.bleu_score.corpus_bleu(refs_list_bleu, cands_list_bleu,weights=(0.5, 0.5, 0, 0)) 93 | print(f'bleu1 : {bleu_score_1}') 94 | print(f'bleu2 : {bleu_score_2}') 95 | print(f'bleu : {bleu_score}') 96 | 97 | 98 | def batch(iterable, n): 99 | iterable=iter(iterable) 100 | while True: 101 | chunk=[] 102 | for i in range(n): 103 | try: 104 | chunk.append(next(iterable)) 105 | except StopIteration: 106 | yield chunk 107 | return 108 | yield chunk 109 | 110 | 111 | def embed_similarity(data,batch_size=16): 112 | gt = data['gt'] 113 | pred = data["pred"] 114 | 115 | gt_batch = list(batch(gt, batch_size)) 116 | pred_batch = list(batch(pred, batch_size)) 117 | cosine_scores_all = [] 118 | for i in range(len(gt_batch)): 119 | 120 | embeddings1 = model.encode(gt_batch[i], convert_to_tensor=True) 121 | embeddings2 = model.encode(pred_batch[i], convert_to_tensor=True) 122 | cosine_scores = util.cos_sim(embeddings1, embeddings2) 123 | assert cosine_scores.size()[0] == cosine_scores.size()[1] 124 | score_list = [cosine_scores[k][k].item() for k in range(cosine_scores.size()[0])] 125 | cosine_scores_all.extend(score_list) 126 | #print(f'{i}-th: {np.mean(score_list).item()}') 127 | #cosine_scores_all = torch.stack(cosine_scores_all) 128 | avg_score = np.mean(cosine_scores_all) 129 | print(f'Avg embed similarity: {avg_score}') 130 | #print(f'Avg embed similarity running mean: {np.mean(running_mean)}') 131 | #sys.exit() 132 | 133 | def get_edit_dist(data): 134 | gt = data['gt'] 135 | pred = data["pred"] 136 | assert len(gt) == len(pred) 137 | edit_dist_list = [] 138 | for i,d in enumerate(pred): 139 | gt_str = gt[i] 140 | pred_str = pred[i] 141 | dist = editdistance.distance(gt_str, pred_str) 142 | edit_dist_list.append(dist) 143 | ### now we return mean and median 144 | edit_dist_list = np.array(edit_dist_list) 145 | edit_median = np.median(edit_dist_list) 146 | edit_mean = np.mean(edit_dist_list) 147 | print(f'edit_mean: {edit_mean}') 148 | print(f'edit_median: {edit_median}') 149 | return edit_mean,edit_median 150 | 151 | def exact_match(data): 152 | gt = data['gt'] 153 | pred = data["pred"] 154 | 155 | gt_remove = punctuation_remove(gt) 156 | pred_remove = punctuation_remove(pred) 157 | 158 | assert len(gt) == len(pred) 159 | count = 0 160 | for i,d in enumerate(pred): 161 | gt_str = gt[i] 162 | pred_str = pred[i] 163 | if(gt_str == pred_str): 164 | count += 1 165 | ratio = count/len(gt) 166 | count = 0 167 | for i,d in enumerate(pred): 168 | gt_str = gt_remove[i] 169 | pred_str = pred_remove[i] 170 | if(gt_str == pred_str): 171 | count += 1 172 | ratio_remove = count/len(gt_remove) 173 | print(f'exact_match ratio: {ratio}') 174 | 175 | print(f'exact_match ratio after removing punctuation: {ratio_remove}') 176 | 177 | return ratio 178 | 179 | def remove_eos(data): 180 | gt = data['gt'] 181 | pred = data["pred"] 182 | for i,s in enumerate(pred): 183 | pred[i] = s.replace('<|endoftext|>','') 184 | 185 | def report_metrics(data): 186 | remove_eos(data) 187 | get_rouge(data) 188 | get_bleu(data) 189 | # get_ppl(data) ### ppl please refer to ppl.py for ppl calculation 190 | exact_match(data) 191 | get_edit_dist(data) 192 | embed_similarity(data) 193 | 194 | 195 | if __name__ == '__main__': 196 | 197 | 198 | abcd_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_abcd_simcse_bert_beam.log' 199 | mnli_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_mnli_simcse_bert_beam.log' 200 | woz_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_multi_woz_simcse_bert_beam.log' 201 | sst2_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_sst2_simcse_bert_beam.log' 202 | wmt_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_wmt16_simcse_bert_beam.log' 203 | 204 | path_list = [abcd_path,mnli_path,woz_path,sst2_path,wmt_path] 205 | for p in path_list: 206 | print(f'==={p}===') 207 | data = read_gpt(p) 208 | report_metrics(data) 209 | 210 | 211 | 212 | -------------------------------------------------------------------------------- /decode_beam_search_opt.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Source code modified from https://github.com/budzianowski/PyTorch-Beam-Search-Decoding/blob/master/decode_beam.py 3 | implementation of beam search on GPT-2's logits 4 | ''' 5 | 6 | import operator 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from queue import PriorityQueue 11 | import sys 12 | 13 | 14 | class BeamSearchNode(object): 15 | def __init__(self, hiddenstate, previousNode, wordId, logProb, length): 16 | ''' 17 | :param hiddenstate: for opt here, reuse past can occur errors, here we just use prev of input embeddings 18 | :param previousNode: 19 | :param wordId: 20 | :param logProb: 21 | :param length: 22 | ''' 23 | self.h = hiddenstate 24 | self.prevNode = previousNode 25 | self.wordid = wordId 26 | self.logp = logProb 27 | self.leng = length 28 | 29 | def eval(self, alpha=1.0): 30 | reward = 0 31 | # Add here a function for shaping a reward 32 | 33 | return self.logp / float(self.leng - 1 + 1e-6) + alpha * reward 34 | 35 | def __lt__(self, x): 36 | if(self.eval() < x.eval()): 37 | return True 38 | 39 | else: 40 | return False 41 | 42 | 43 | 44 | def beam_decode_sentence(hidden_X, config,num_generate=1, beam_size = 5, batch_size = 1): 45 | ''' 46 | generate a sentence based on beam search 47 | :param hidden_X: input embeddings 48 | :param model: GPT-2 model 49 | :param tokenizer: GPT-2 tokenizer 50 | :return: decoded_batch 51 | ''' 52 | #SOS_token = tokenizer.encode("<|endoftext|>") 53 | beam_width = beam_size 54 | topk = num_generate # how many sentence do you want to generate 55 | 56 | past = None 57 | model = config['model'] 58 | tokenizer =config['tokenizer'] 59 | #eos = [tokenizer.encoder["<|endoftext|>"]] 60 | eos = [tokenizer.encode(tokenizer.eos_token)] 61 | EOS_token = eos 62 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X, 0) 63 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X_unsqueeze, 0) #[1,1,embed_dim] [batch_size, seq_len, emb_dim] 64 | 65 | decoded_batch = [] 66 | 67 | # decoding goes sentence by sentence 68 | for idx in range(batch_size): 69 | 70 | # Number of sentence to generate 71 | endnodes = [] 72 | number_required = min((topk + 1), topk - len(endnodes)) 73 | 74 | # starting node - hidden vector, previous node, word id, logp, length hiddenstate, previousNode, wordId, logProb, length 75 | #node = BeamSearchNode(past, None, torch.tensor([[220]]).cuda(), 0, 1) # 220 refers to single space ' ' 76 | # if(config['use_opt']): 77 | # node = BeamSearchNode(past, None, torch.tensor([[2]]).cuda(), 0, 1) # 2 refers to '' on opt 78 | # else: 79 | # node = BeamSearchNode(past, None, torch.tensor([[220]]).cuda(), 0, 1) # 220 refers to single space ' ' on GPT 80 | node = BeamSearchNode(hidden_X_unsqueeze, None, torch.tensor([[2]]).cuda(), 0, 1) # 2 refers to '' on opt 81 | nodes = PriorityQueue() 82 | 83 | # start the queue 84 | nodes.put((-node.eval(), node)) 85 | qsize = 1 86 | 87 | # start beam search 88 | for text_len in range(50): 89 | # give up when decoding takes too long 90 | if qsize > 2000: break 91 | 92 | # fetch the best node 93 | try: 94 | score, n = nodes.get() 95 | except: 96 | print('Cannot get nodes') 97 | while not nodes.empty(): 98 | next_item = nodes.get() 99 | print(next_item) 100 | prev_input = n.wordid 101 | past = n.h 102 | 103 | if n.wordid.item() == EOS_token[0] and n.prevNode != None: 104 | endnodes.append((score, n)) 105 | # if we reached maximum # of sentences required 106 | #print(f'EOS found') 107 | if len(endnodes) >= number_required: 108 | #print('break') 109 | break 110 | 111 | else: 112 | print('continue') 113 | continue 114 | 115 | # decode for one step using decoder 116 | #decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_output) 117 | 118 | output = model(inputs_embeds=past,past_key_values = None,return_dict=True) 119 | logits = output.logits 120 | 121 | logits = logits[:, -1, :] 122 | probs = torch.softmax(logits, dim=-1) 123 | # PUT HERE REAL BEAM SEARCH OF TOP 124 | log_prob, indexes = torch.topk(probs, beam_width) 125 | nextnodes = [] 126 | 127 | for new_k in range(beam_width): 128 | decoded_t = indexes[0][new_k].view(1, -1) 129 | log_p = log_prob[0][new_k].item() 130 | #### hiddenstate, previousNode, wordId, logProb, length 131 | input_emb = model.model.decoder.embed_tokens(decoded_t) 132 | new_past = torch.cat((past,input_emb),dim=1) 133 | node = BeamSearchNode(new_past, n, decoded_t, n.logp + log_p, n.leng + 1) 134 | score = -node.eval() 135 | nextnodes.append((score, node)) 136 | 137 | # put them into queue 138 | for i in range(len(nextnodes)): 139 | score, nn = nextnodes[i] 140 | try: 141 | nodes.put((score, nn)) 142 | except: 143 | print('Cannot put nodes') 144 | print(score) 145 | print(nn) 146 | # increase qsize 147 | qsize += len(nextnodes) - 1 148 | # for loop ends here 149 | # choose nbest paths, back trace them 150 | if len(endnodes) == 0: 151 | endnodes = [nodes.get() for _ in range(topk)] 152 | 153 | utterances = [] 154 | text = [] 155 | for score, n in sorted(endnodes, key=operator.itemgetter(0)): 156 | utterance = [] 157 | utterance.append(n.wordid.item()) 158 | # back trace 159 | while n.prevNode != None: 160 | n = n.prevNode 161 | utterance.append(n.wordid.item()) 162 | 163 | utterance = utterance[::-1] 164 | utterances.append(utterance) 165 | decode_process = tokenizer.decode(utterance[1:-1],skip_special_tokens=True) 166 | text.append(decode_process) 167 | decoded_batch.append(utterances) 168 | 169 | return text 170 | 171 | 172 | def greedy_decode(decoder_hidden, encoder_outputs, target_tensor): 173 | ''' 174 | :param target_tensor: target indexes tensor of shape [B, T] where B is the batch size and T is the maximum length of the output sentence 175 | :param decoder_hidden: input tensor of shape [1, B, H] for start of the decoding 176 | :param encoder_outputs: if you are using attention mechanism you can pass encoder outputs, [T, B, H] where T is the maximum length of input sentence 177 | :return: decoded_batch 178 | ''' 179 | 180 | batch_size, seq_len = target_tensor.size() 181 | decoded_batch = torch.zeros((batch_size, MAX_LENGTH)) 182 | decoder_input = torch.LongTensor([[SOS_token] for _ in range(batch_size)], device=device) 183 | 184 | for t in range(MAX_LENGTH): 185 | decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_outputs) 186 | 187 | topv, topi = decoder_output.data.topk(1) # get candidates 188 | topi = topi.view(-1) 189 | decoded_batch[:, t] = topi 190 | 191 | decoder_input = topi.detach().view(-1, 1) 192 | 193 | return decoded_batch -------------------------------------------------------------------------------- /result verification/ner.py: -------------------------------------------------------------------------------- 1 | import stanza 2 | import json 3 | print("start") 4 | nlp = stanza.Pipeline(lang='en', processors='tokenize,ner') 5 | 6 | print("start_2") 7 | # for gpt result 8 | file_list = ['path to gpt result'] 9 | 10 | 11 | for file in file_list: 12 | with open(file) as f: 13 | print(file) 14 | result = json.load(f) 15 | gt_total_word = 0 16 | gt_total_entity = 0 17 | gt_entity_set_list = [] 18 | for item in result['gt']: 19 | doc_gt = nlp(item) 20 | entity_set = set() 21 | for sentence in doc_gt.sentences: 22 | for token in sentence.tokens: 23 | gt_total_word += 1 24 | if token.ner != 'O': 25 | gt_total_entity += 1 26 | entity_set.add(token.text) 27 | gt_entity_set_list.append(entity_set) 28 | 29 | pred_total_word = 0 30 | pred_total_entity = 0 31 | pred_entity_set_list = [] 32 | for item in result['pred']: 33 | doc_pred = nlp(item) 34 | entity_set = set() 35 | for sentence in doc_pred.sentences: 36 | for token in sentence.tokens: 37 | pred_total_word += 1 38 | if token.ner != 'O': 39 | pred_total_entity += 1 40 | entity_set.add(token.text) 41 | pred_entity_set_list.append(entity_set) 42 | 43 | 44 | correct_pred_count = 0 45 | for i in range(len(gt_entity_set_list)): 46 | correct_set = gt_entity_set_list[i] & pred_entity_set_list[i] 47 | correct_pred_count += len(correct_set) 48 | 49 | 50 | gt_nerr = gt_total_entity/gt_total_word 51 | pred_nerr = pred_total_entity/pred_total_word 52 | real_pred_nerr = correct_pred_count/gt_total_word 53 | print("gt_total_entity", end = '') 54 | print(gt_total_entity) 55 | print("gt_total_word", end = '') 56 | print(gt_total_word) 57 | print("pred_total_entity", end = '') 58 | 59 | print(pred_total_entity) 60 | print("pred_total_word", end = '') 61 | print(pred_total_word) 62 | print("correct_pred_count", end = '') 63 | print(correct_pred_count) 64 | print("gt_nerr", end = '') 65 | print(gt_nerr) 66 | print("pred_nerr", end = '') 67 | print(pred_nerr) 68 | print("real_pred_nerr", end = '') 69 | print(real_pred_nerr) 70 | 71 | ner = correct_pred_count/gt_total_entity 72 | print("ner") 73 | print(ner) 74 | 75 | # for persona_data result 76 | file_list = ['path to persona_data result'] 77 | 78 | 79 | for file in file_list: 80 | 81 | entity_set = set() 82 | total_entity = 0 83 | with open(file, "r", encoding='UTF-8') as f: 84 | total_word = 0 85 | print("dev") 86 | out = f.read() 87 | out = json.loads(out) 88 | for element in out: 89 | sentences = "" 90 | for sentence in element["partner_persona"]: 91 | sentences += sentence 92 | for sentence in element["your_persona"]: 93 | sentences += sentence 94 | for sentence in element["conv"]: 95 | sentences += sentence 96 | doc_sentences = nlp(sentences) 97 | for sentence in doc_sentences.sentences: 98 | # print(str(sentence)) 99 | for token in sentence.tokens: 100 | # print(token.text) 101 | # print(token.ner) 102 | total_word += 1 103 | if token.ner != 'O': 104 | total_entity += 1 105 | entity_set.add(token.text) 106 | print("unique entity total") 107 | print(len(entity_set)) 108 | print("total_entity") 109 | print(total_entity) 110 | 111 | 112 | # for RNN result 113 | file_list = ['path to RNN result'] 114 | 115 | for file in file_list: 116 | with open(file) as f: 117 | print(file) 118 | result = json.load(f) 119 | 120 | for i in result: 121 | sentences = tokenizer.encode(i['input']) 122 | for word in sentences: 123 | decode_item = tokenizer.decode(word) 124 | 125 | gt_total_word = 0 126 | gt_total_entity = 0 127 | correct_pred_count = 0 128 | for i in result: 129 | decode_entity_set = set() 130 | gt_set = set() 131 | pred_set = set() 132 | doc_gt = nlp(i['input']) 133 | for sentence in doc_gt.sentences: 134 | for token in sentence.tokens: 135 | if token.ner != 'O': 136 | # print(token.text) 137 | word_tokens = tokenizer.encode(token.text) 138 | for word_token in word_tokens: 139 | decode_token = tokenizer.decode(word_token) 140 | decode_entity_set.add(decode_token) 141 | for item in i['gt']: 142 | gt_set.add(item) 143 | for item in i['pred']: 144 | pred_set.add(item) 145 | gt_total_word += len(gt_set) 146 | gt_total_entity += len(decode_entity_set) 147 | correct_set = decode_entity_set & pred_set 148 | correct_pred_count += len(correct_set) 149 | 150 | gt_nerr = gt_total_entity/gt_total_word 151 | real_pred_nerr = correct_pred_count/gt_total_word 152 | print("gt_total_word", end = '') 153 | print(gt_total_word) 154 | print("gt_total_entity", end = '') 155 | print(gt_total_entity) 156 | print("correct_pred_count", end = '') 157 | print(correct_pred_count) 158 | print("gt_nerr", end = '') 159 | print(gt_nerr) 160 | print("real_pred_nerr", end = '') 161 | print(real_pred_nerr) 162 | 163 | ner = correct_pred_count / gt_total_entity 164 | print("ner", end = '') 165 | print(ner) 166 | 167 | 168 | 169 | 170 | # for NN result 171 | file_list = ['path to NN result'] 172 | for file in file_list: 173 | with open(file) as f: 174 | print(file) 175 | result = json.load(f) 176 | 177 | for i in result: 178 | sentences = tokenizer.encode(i['input']) 179 | # print(sentences) 180 | for word in sentences: 181 | decode_item = tokenizer.decode(word) 182 | # print(decode_item) 183 | 184 | gt_total_word = 0 185 | gt_total_entity = 0 186 | correct_pred_count = 0 187 | for i in result: 188 | decode_entity_set = set() 189 | gt_set = set() 190 | pred_set = set() 191 | doc_gt = nlp(i['input']) 192 | for sentence in doc_gt.sentences: 193 | for token in sentence.tokens: 194 | if token.ner != 'O': 195 | # print(token.text) 196 | word_tokens = tokenizer.encode(token.text) 197 | for word_token in word_tokens: 198 | decode_token = tokenizer.decode(word_token) 199 | decode_entity_set.add(decode_token) 200 | for item in i['gt']: 201 | gt_set.add(item) 202 | for item in i['pred']: 203 | pred_set.add(item) 204 | # print(decode_entity_set) 205 | gt_total_word += len(gt_set) 206 | gt_total_entity += len(decode_entity_set) 207 | correct_set = decode_entity_set & pred_set 208 | ''' 209 | if len(correct_set) != 0: 210 | print(i) 211 | print("decode_entity_set") 212 | print(decode_entity_set) 213 | print("pred_set") 214 | print(pred_set) 215 | print("correct_set") 216 | print(correct_set) 217 | ''' 218 | correct_pred_count += len(correct_set) 219 | 220 | gt_nerr = gt_total_entity/gt_total_word 221 | real_pred_nerr = correct_pred_count/gt_total_word 222 | print("gt_total_word", end = '') 223 | print(gt_total_word) 224 | print("gt_total_entity", end = '') 225 | print(gt_total_entity) 226 | print("correct_pred_count", end = '') 227 | print(correct_pred_count) 228 | print("gt_nerr", end = '') 229 | print(gt_nerr) 230 | print("real_pred_nerr", end = '') 231 | print(real_pred_nerr) 232 | 233 | ner = correct_pred_count / gt_total_entity 234 | print("ner", end = '') 235 | print(ner) 236 | -------------------------------------------------------------------------------- /attacker_models.py: -------------------------------------------------------------------------------- 1 | import os 2 | #os.environ["CUDA_VISIBLE_DEVICES"] = "3" 3 | import time 4 | import numpy as np 5 | import pandas as pd 6 | import tqdm 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from torch.utils.data import DataLoader, Dataset 11 | from torch.nn.utils.rnn import pad_sequence 12 | from transformers import AutoTokenizer 13 | from transformers import AdamW 14 | import sys 15 | 16 | from sklearn import metrics 17 | 18 | 19 | batch_size = 32 20 | def read_pt(data_type,use_trans=False): 21 | # make it easier to load into a batch 22 | assert data_type == 'train' or 'dev' or 'test' 23 | if use_trans: 24 | path = 'hidden_'+data_type+'_trans.pt' 25 | else: 26 | path = 'hidden_'+data_type+'.pt' 27 | data = torch.load(path) 28 | X=[] 29 | Y=[] 30 | A=[] #A for attribute infor 31 | D=[] 32 | for idx,dialog_dict in enumerate(data): #for a dialog 33 | input_label = dialog_dict['label'] #list of tensors 34 | persona_list = dialog_dict['persona'] #list of persona int 35 | hidden_tensor = dialog_dict['hidden'] #list of tensors for hidden 36 | if use_trans: 37 | utterance_list = dialog_dict['dial'] 38 | for i,d in enumerate(input_label): 39 | 40 | X.append(hidden_tensor[i][-1]) 41 | Y.append(input_label[i].squeeze()) 42 | A.append(persona_list[i]) 43 | if use_trans: 44 | D.append(utterance_list[i]) 45 | if(use_trans): 46 | return X,Y,A,D 47 | 48 | return X,Y,A 49 | 50 | 51 | class Dataset(Dataset): 52 | def __init__(self, X,Y,A): 53 | self.X = X 54 | self.Y = Y 55 | self.A = A 56 | 57 | def __len__(self): 58 | 59 | return len(self.X) 60 | 61 | def __getitem__(self, index): 62 | 63 | sample_X = self.X[index] 64 | sample_Y = self.Y[index] 65 | sample_A = self.A[index] 66 | return sample_X, sample_Y, sample_A 67 | 68 | def collate(self, unpacked_data): 69 | return unpacked_data 70 | 71 | class model_inv_nn(nn.Module): 72 | def __init__(self,out_num,in_num=1024): 73 | super(model_inv_nn, self).__init__() 74 | self.fc1 = nn.Linear(in_num, out_num) 75 | #self.act = F.softmax() 76 | 77 | def forward(self, x): 78 | # x should be of shape (?,1024) 79 | out = self.fc1(x) 80 | #out = F.softmax(self.fc1(x),dim=1) 81 | 82 | return out 83 | 84 | 85 | 86 | ''' 87 | 88 | Model for transformer attackers 89 | 90 | ''' 91 | class SequenceCrossEntropyLoss(nn.Module): 92 | def __init__(self): 93 | super().__init__() 94 | 95 | def forward(self, logits, targets, mask, label_smoothing=-1, reduce=None): 96 | """ 97 | reduce: None, "batch", "sentence" 98 | """ 99 | return sequence_cross_entropy_with_logits(logits, targets, mask, label_smoothing, reduce) 100 | 101 | def sequence_cross_entropy_with_logits(logits, targets, mask, label_smoothing, reduce): 102 | # type: (Tensor, Tensor, Tensor, float, bool)-> Tensor 103 | """ 104 | label_smoothing : ``float``, optional (default = 0.0) 105 | It should be smaller than 1. 106 | """ 107 | # shape : (batch * sequence_length, num_classes) 108 | logits_flat = logits.view(-1, logits.size(-1)) 109 | # shape : (batch * sequence_length, num_classes) 110 | log_probs_flat = F.log_softmax(logits_flat, dim=-1) 111 | # shape : (batch * max_len, 1) 112 | targets_flat = targets.view(-1, 1).long() 113 | 114 | if label_smoothing > 0.0: 115 | num_classes = logits.size(-1) 116 | smoothing_value = label_smoothing / float(num_classes) 117 | # Fill all the correct indices with 1 - smoothing value. 118 | one_hot_targets = torch.zeros_like(log_probs_flat).scatter_(-1, targets_flat, 1.0 - label_smoothing) 119 | smoothed_targets = one_hot_targets + smoothing_value 120 | negative_log_likelihood_flat = -log_probs_flat * smoothed_targets 121 | negative_log_likelihood_flat = negative_log_likelihood_flat.sum(-1, keepdim=True) 122 | else: 123 | # shape : (batch * sequence_length, 1) 124 | negative_log_likelihood_flat = - torch.gather(log_probs_flat, dim=1, index=targets_flat) 125 | 126 | # shape : (batch, sequence_length) 127 | negative_log_likelihood = negative_log_likelihood_flat.view(-1, logits.shape[1]) 128 | 129 | # shape : (batch, sequence_length) 130 | loss = negative_log_likelihood * mask 131 | 132 | if reduce: 133 | # shape : (batch,) 134 | loss = loss.sum(1) / (mask.sum(1) + 1e-13) 135 | 136 | if reduce is "batch": 137 | # shape : scalar 138 | loss = loss.mean() 139 | 140 | return loss 141 | 142 | class Dataset_trans(Dataset): 143 | def __init__(self, X,Y,A,D): 144 | self.X = X 145 | self.Y = Y 146 | self.A = A 147 | self.D = D 148 | 149 | def __len__(self): 150 | 151 | return len(self.X) 152 | 153 | def __getitem__(self, index): 154 | 155 | sample_X = self.X[index] 156 | sample_Y = self.Y[index] 157 | sample_A = self.A[index] 158 | sample_D = self.D[index] 159 | return sample_X, sample_Y, sample_A,sample_D 160 | 161 | def collate(self, unpacked_data): 162 | return unpacked_data 163 | 164 | 165 | def train_on_batch(batch_X, batch_Y,batch_A,model,optimizer,criterion): 166 | optimizer.zero_grad() 167 | output = model(batch_X) 168 | loss = criterion(output, batch_Y) 169 | loss.backward() 170 | optimizer.step() 171 | print(f'loss: {loss.item()}') 172 | 173 | def evaluation(dataloader,model,criterion): 174 | loss_list = [] 175 | predict = [] 176 | ground_truth = [] 177 | count = 0 178 | with torch.no_grad(): 179 | for (batch_X, batch_Y,batch_A) in dataloader: 180 | print(f'count:{count}') 181 | batch_size = batch_X.size()[0] 182 | label_size = batch_Y.size()[1] 183 | # move to gpu 184 | batch_X = batch_X.cuda() 185 | batch_Y = batch_Y.cuda() 186 | batch_A = batch_A.cuda() 187 | output = model(batch_X) 188 | m = nn.Sigmoid() 189 | batch_out = m(output) 190 | batch_out[batch_out>=0.5] = 1 191 | batch_out[batch_out<0.5] = 0 192 | #eval_metrics(batch_out,batch_Y) ### what we want 193 | 194 | loss = criterion(output, batch_Y) 195 | loss_val = loss.item() 196 | loss_list.append(loss_val*batch_size) 197 | predict.extend(batch_out.cpu().detach().numpy()) 198 | ground_truth.extend(batch_Y.cpu().detach().numpy()) 199 | 200 | count +=1 201 | 202 | 203 | 204 | avg_loss = np.mean(loss_list) 205 | predict = np.array(predict) 206 | ground_truth = np.array(ground_truth) 207 | report_score(ground_truth,predict) 208 | print(f'avg_loss: {avg_loss}') 209 | 210 | def report_score(y_true,y_pred): 211 | # micro result should be reported 212 | print("micro precision_score: {:.2f}".format(metrics.precision_score(y_true, y_pred, average='micro'))) 213 | print("macro precision_score: {:.2f} ".format( metrics.precision_score(y_true, y_pred, average='macro'))) 214 | print('=='*20) 215 | print("micro recall_score: {:.2f}".format(metrics.recall_score(y_true, y_pred, average='micro'))) 216 | print("macro recall_score: {:.2f} ".format( metrics.recall_score(y_true, y_pred, average='macro'))) 217 | print('=='*20) 218 | print("micro f1_score: {:.2f}".format(metrics.f1_score(y_true, y_pred, average='micro'))) 219 | print("macro f1_score: {:.2f} ".format( metrics.f1_score(y_true, y_pred, average='macro'))) 220 | 221 | 222 | 223 | 224 | if __name__ == '__main__': 225 | print('main file') 226 | data_type = 'dev' 227 | device = torch.device("cuda") 228 | X,Y,A = read_pt(data_type) 229 | train_dataset = Dataset(X,Y,A) 230 | #batch_size = batch_size 231 | external_criterion = nn.BCEWithLogitsLoss() 232 | tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") 233 | token_num = len(tokenizer) 234 | inv_model = model_inv_nn(out_num=token_num) 235 | inv_model.to(device) 236 | optimizer = torch.optim.Adam(inv_model.parameters(), 237 | lr=3e-5, 238 | eps=1e-06) 239 | train_dataloader = DataLoader(dataset=train_dataset, 240 | shuffle=True, 241 | batch_size=batch_size) 242 | #dataloader help covert batch into tensors 243 | for (batch_X, batch_Y,batch_A) in train_dataloader: 244 | #print(batch_X.size()) [batch, 1024] 245 | #print(batch_Y.size()) [batch, 50257] 246 | #print(batch_A.size()) [batch] 247 | train_on_batch(batch_X.cuda(), batch_Y.cuda(),batch_A.cuda(),inv_model,optimizer,external_criterion) 248 | -------------------------------------------------------------------------------- /eval_classification.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "0" 3 | from sklearn import metrics 4 | import numpy as np 5 | import argparse 6 | import torch 7 | from transformers import AutoModel, AutoTokenizer 8 | import json 9 | from sentence_transformers import SentenceTransformer, util 10 | from scipy.spatial.distance import cosine 11 | from simcse import SimCSE 12 | import logging 13 | import logging.handlers 14 | from nltk.tokenize import word_tokenize 15 | import string 16 | import re 17 | 18 | logger = logging.getLogger('mylogger') 19 | logger.setLevel(logging.DEBUG) 20 | f_handler = logging.FileHandler('models_arr_feb/decoder_beam.log') 21 | f_handler.setLevel(logging.INFO) 22 | f_handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) 23 | logger.addHandler(f_handler) 24 | 25 | def vectorize(sent_list,tokenizer): 26 | turn_ending = tokenizer.encode(tokenizer.eos_token) 27 | token_num = len(tokenizer) 28 | dial_tokens = [tokenizer.encode(item) + turn_ending for item in sent_list] 29 | dial_tokens_np = np.array(dial_tokens) 30 | input_labels = [] 31 | for i in dial_tokens_np: 32 | temp_i = np.zeros(token_num) 33 | temp_i[i] = 1 34 | input_labels.append(temp_i) 35 | input_labels = np.array(input_labels) 36 | 37 | 38 | return input_labels 39 | 40 | 41 | 42 | 43 | 44 | def report_score(y_true,y_pred): 45 | # micro result should be reported 46 | precision = metrics.precision_score(y_true, y_pred, average='micro') 47 | recall = metrics.recall_score(y_true, y_pred, average='micro') 48 | f1 = metrics.f1_score(y_true, y_pred, average='micro') 49 | logger.info(f"micro precision_score on token level: {str(precision)}") 50 | logger.info(f"micro recall_score on token level: {str(recall)}") 51 | logger.info(f"micro f1_score on token level: {str(f1)}") 52 | 53 | 54 | 55 | def embed_simcse(y_true,y_pred): 56 | model = SimCSE("princeton-nlp/sup-simcse-roberta-large",device='cuda') 57 | similarities = model.similarity(y_true, y_pred) # numpy array of N*N 58 | pair_scores = similarities.diagonal() 59 | for i,score in enumerate(pair_scores): 60 | assert pair_scores[i] == similarities[i][i] 61 | avg_score = np.mean(pair_scores) 62 | logger.info(f'Evaluation on simcse-roberta with similarity score {avg_score}') 63 | 64 | 65 | def embed_sbert(y_true,y_pred): 66 | model = SentenceTransformer('all-roberta-large-v1',device='cuda') # has dim 768 67 | embeddings_true = model.encode(y_true,convert_to_tensor = True) 68 | embeddings_pred = model.encode(y_pred,convert_to_tensor = True) 69 | cosine_scores = util.cos_sim(embeddings_true, embeddings_pred) 70 | pair_scores = torch.diagonal(cosine_scores, 0) 71 | for i,score in enumerate(pair_scores): 72 | assert pair_scores[i] == cosine_scores[i][i] 73 | avg_score = torch.mean(pair_scores) 74 | logger.info(f'Evaluation on Sentence-bert with similarity score {avg_score}') 75 | 76 | 77 | return avg_score 78 | 79 | 80 | def report_embedding_similarity(y_true,y_pred): 81 | embed_sbert(y_true,y_pred) 82 | embed_simcse(y_true,y_pred) 83 | 84 | 85 | def main(log_path): 86 | with open(log_path, 'r') as f: 87 | sent_dict = json.load(f) 88 | y_true = sent_dict['gt'] # list of sentences 89 | y_pred = sent_dict['pred'] # list of sentences 90 | report_embedding_similarity(y_true,y_pred) 91 | 92 | 93 | ''' 94 | 26/02 newly appended functions 95 | ''' 96 | # remove punctuation from list of sentences 97 | def punctuation_remove(sent_list): 98 | removed_list = [] 99 | for sent in sent_list: 100 | word_list = [] 101 | for word in sent.split(): 102 | word_strip = word.strip(string.punctuation) 103 | if word_strip: # cases for not empty string 104 | word_list.append(word_strip) 105 | removed_sent = ' '.join(word_list) 106 | removed_list.append(removed_sent) 107 | return removed_list 108 | 109 | # remove space before punctuation from list of sentences 110 | def space_remove(sent_list): 111 | removed_list = [] 112 | for sent in sent_list: 113 | sent_remove = re.sub(r'\s([?.!"](?:\s|$))', r'\1', sent) 114 | removed_list.append(sent_remove) 115 | return removed_list 116 | 117 | def metrics_word_level(token_true,token_pred): 118 | len_pred = len(token_pred) 119 | len_ture = len(token_true) 120 | recover_pred = 0 121 | recover_true = 0 122 | for p in token_pred: 123 | if p in token_true: 124 | recover_pred += 1 125 | for t in token_true: 126 | if t in token_pred: 127 | recover_true += 1 128 | ### return for precision recall calculation 129 | return len_pred,recover_pred,len_ture,recover_true 130 | 131 | 132 | def word_level_metrics(y_true,y_pred): 133 | assert len(y_true) == len(y_pred) 134 | recover_pred_all = 0 135 | recover_true_all = 0 136 | len_pred_all = 0 137 | len_ture_all = 0 138 | for i in range(len(y_true)): 139 | sent_true = y_true[i] 140 | sent_pred = y_pred[i] 141 | token_true = word_tokenize(sent_true) 142 | token_pred = word_tokenize(sent_pred) 143 | len_pred,recover_pred,len_ture,recover_true = metrics_word_level(token_true,token_pred) 144 | len_pred_all += len_pred 145 | recover_pred_all += recover_pred 146 | len_ture_all += len_ture 147 | recover_true_all += recover_true 148 | 149 | 150 | ### precision and recall are based on micro (but not exactly) 151 | precision = recover_pred_all/len_pred_all 152 | recall = recover_true_all/len_ture_all 153 | f1 = 2*precision*recall/(precision+recall) 154 | return precision,recall,f1 155 | 156 | def remove_eos(sent_list): 157 | for i,s in enumerate(sent_list): 158 | sent_list[i] = s.replace('<|endoftext|>','') 159 | 160 | def metric_token(log_path): 161 | with open(log_path, 'r') as f: 162 | sent_dict = json.load(f) 163 | y_true = sent_dict['gt'] # list of sentences 164 | y_pred = sent_dict['pred'] # list of sentences 165 | 166 | tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") 167 | y_true_token = vectorize(y_true,tokenizer) 168 | y_pred_token = vectorize(y_pred,tokenizer) 169 | 170 | ### token-level metrics are reported 171 | report_score(y_true_token,y_pred_token) 172 | remove_eos(y_pred) # make sure to remove 173 | ### scores for word level 174 | y_true_removed_p = punctuation_remove(y_true) 175 | y_pred_removed_p = punctuation_remove(y_pred) 176 | y_true_removed_s = space_remove(y_true) 177 | y_pred_removed_s = space_remove(y_pred) 178 | precision,recall,f1 = word_level_metrics(y_true_removed_s,y_pred_removed_s) 179 | logger.info(f'word level precision: {str(precision)}') 180 | logger.info(f'word level recall: {str(recall)}') 181 | logger.info(f'word level f1: {str(f1)}') 182 | 183 | precision,recall,f1 = word_level_metrics(y_true_removed_p,y_pred_removed_p) 184 | logger.info(f'word level precision without punctuation: {str(precision)}') 185 | logger.info(f'word level recall without punctuation: {str(recall)}') 186 | logger.info(f'word level f1 without punctuation: {str(f1)}') 187 | 188 | 189 | if __name__ == '__main__': 190 | ''' 191 | ### PC with sampling and randomly initialized GPT-2 192 | sbert_roberta_large_pc_path = 'models_random/attacker_gpt2_personachat_sent_roberta.log' 193 | simcse_roberta_large_pc_path = 'models_random/attacker_gpt2_personachat_simcse_roberta.log' 194 | simcse_bert_large_pc_path = 'models_random/attacker_gpt2_personachat_simcse_bert.log' 195 | sentence_T5_large_pc_path = 'models_random/attacker_gpt2_personachat_sent_t5.log' 196 | mpnet_pc_path = 'models_random/attacker_gpt2_personachat_mpnet.log' 197 | logger.info(f'====={sbert_roberta_large_pc_path}=====') 198 | metric_token(sbert_roberta_large_pc_path) 199 | logger.info(f'====={simcse_roberta_large_pc_path}=====') 200 | metric_token(simcse_roberta_large_pc_path) 201 | logger.info(f'====={simcse_bert_large_pc_path}=====') 202 | metric_token(simcse_bert_large_pc_path) 203 | logger.info(f'====={sentence_T5_large_pc_path}=====') 204 | metric_token(sentence_T5_large_pc_path) 205 | logger.info(f'====={mpnet_pc_path}=====') 206 | metric_token(mpnet_pc_path) 207 | ''' 208 | 209 | abcd_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_abcd_simcse_bert_beam.log' 210 | mnli_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_mnli_simcse_bert_beam.log' 211 | woz_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_multi_woz_simcse_bert_beam.log' 212 | sst2_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_sst2_simcse_bert_beam.log' 213 | wmt_path = '/home/hlibt/embed_rev/models_arr_feb/attacker_gpt2_wmt16_simcse_bert_beam.log' 214 | 215 | path_list = [abcd_path,mnli_path,woz_path,sst2_path,wmt_path] 216 | for p in path_list: 217 | logger.info(f'====={p}=====') 218 | metric_token(p) -------------------------------------------------------------------------------- /attacker_evaluation_gpt.py: -------------------------------------------------------------------------------- 1 | import os 2 | #os.environ["CUDA_VISIBLE_DEVICES"] = "0" 3 | import time 4 | import numpy as np 5 | import pandas as pd 6 | import tqdm 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | from torch.utils.data import DataLoader 11 | from torch.nn.utils.rnn import pad_sequence 12 | from transformers import AutoModelForCausalLM,AutoTokenizer 13 | from transformers import AdamW 14 | import sys 15 | import argparse 16 | import attacker_models 17 | from attacker_models import read_pt, Dataset,Dataset_trans,SequenceCrossEntropyLoss 18 | #from sentence_revocer_transformer import train_on_batch 19 | import json 20 | from decode_beam_search import beam_decode_sentence 21 | import decode_beam_search_opt 22 | #from bookcorpus_train import str2bool,BookCorpus_Dataset 23 | def get_dataloader(config): 24 | data_type= config['data_type'] 25 | batch_size = config['batch_size'] 26 | if config['use_trans']: 27 | if config['p_simcse_flag']: 28 | ### path to pt checkpoint 29 | data = torch.load('/data/hlibt/gradient_leakage/pytorch/data/personachat_processed/hidden_test_sbert.pt') 30 | dataset = BookCorpus_Dataset(data) 31 | else: 32 | X,Y,A,D = read_pt(data_type,use_trans=config['use_trans']) 33 | dataset = Dataset_trans(X,Y,A,D) 34 | else: 35 | X,Y,A = read_pt(data_type,use_trans=config['use_trans']) 36 | dataset = Dataset(X,Y,A) 37 | dataloader = DataLoader(dataset=dataset, 38 | shuffle=True, 39 | batch_size=batch_size) 40 | return dataloader 41 | 42 | def get_model(config): 43 | 44 | model_dir = config['model_dir'] 45 | model_type = config['model_type'] 46 | 47 | if model_type == '1layerNN': 48 | model = attacker_models.model_inv_nn(out_num=config['token_num']) 49 | model.load_state_dict(torch.load(model_dir)) 50 | model.to(config['device']) 51 | criterion = nn.BCEWithLogitsLoss() 52 | 53 | 54 | else: 55 | print('No proper model loaded') 56 | sys.exit(-1) 57 | 58 | return model,criterion 59 | 60 | 61 | def top_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): 62 | """ Filter a distribution of logits using top-k, top-p (nucleus) and/or threshold filtering 63 | Args: 64 | logits: logits distribution shape (vocabulary size) 65 | top_k: <=0: no filtering, >0: keep only top k tokens with highest probability. 66 | top_p: <=0.0: no filtering, >0.0: keep only a subset S of candidates, where S is the smallest subset 67 | whose total probability mass is greater than or equal to the threshold top_p. 68 | In practice, we select the highest probability tokens whose cumulative probability mass exceeds 69 | the threshold top_p. 70 | """ 71 | # batch support! 72 | if top_k > 0: 73 | values, _ = torch.topk(logits, top_k) 74 | min_values = values[:, -1].unsqueeze(1).repeat(1, logits.shape[-1]) 75 | logits = torch.where(logits < min_values, 76 | torch.ones_like(logits, dtype=logits.dtype) * -float('Inf'), 77 | logits) 78 | if top_p > 0.0: 79 | # Compute cumulative probabilities of sorted tokens 80 | sorted_logits, sorted_indices = torch.sort(logits, descending=True) 81 | cumulative_probabilities = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) 82 | 83 | # Remove tokens with cumulative probability above the threshold 84 | sorted_indices_to_remove = cumulative_probabilities > top_p 85 | # Shift the indices to the right to keep also the first token above the threshold 86 | sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() 87 | sorted_indices_to_remove[..., 0] = 0 88 | 89 | sorted_logits = sorted_logits.masked_fill_(sorted_indices_to_remove, filter_value) 90 | logits = torch.zeros_like(logits).scatter(1, sorted_indices, sorted_logits) 91 | 92 | return logits 93 | 94 | def generate_sentence(config,hidden_X): 95 | temperature = 0.9 96 | top_k = -1 97 | top_p = 0.9 98 | sent = [] 99 | prev_input = None 100 | past = None 101 | model = config['model'] 102 | tokenizer =config['tokenizer'] 103 | #eos = [tokenizer.encoder["<|endoftext|>"]] 104 | eos = tokenizer.encode("<|endoftext|>") 105 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X, 0) 106 | hidden_X_unsqueeze = torch.unsqueeze(hidden_X_unsqueeze, 0) #[1,1,embed_dim] 107 | logits, past = model(inputs_embeds=hidden_X_unsqueeze,past_key_values = past,return_dict=False) 108 | logits = logits[:, -1, :] / temperature 109 | logits = top_filtering(logits, top_k=top_k, top_p=top_p) 110 | 111 | probs = torch.softmax(logits, dim=-1) 112 | 113 | prev_input = torch.multinomial(probs, num_samples=1) 114 | prev_word = prev_input.item() 115 | sent.append(prev_word) 116 | 117 | for i in range(50): 118 | #logits, past = model(prev_input, past=past) 119 | logits, past = model(prev_input,past_key_values = past,return_dict=False) 120 | logits = logits[:, -1, :] / temperature 121 | logits = top_filtering(logits, top_k=top_k, top_p=top_p) 122 | 123 | probs = torch.softmax(logits, dim=-1) 124 | 125 | prev_input = torch.multinomial(probs, num_samples=1) 126 | prev_word = prev_input.item() 127 | 128 | if prev_word == eos[0]: 129 | break 130 | sent.append(prev_word) 131 | 132 | output = tokenizer.decode(sent) 133 | 134 | return output 135 | 136 | 137 | def eval(dataloader,config): 138 | model = config['model'] 139 | tokenizer = config['tokenizer'] 140 | device = config['device'] 141 | model.to(device) 142 | criterion = SequenceCrossEntropyLoss() 143 | save_path = config['save_path'] 144 | sent_dict = {} 145 | sent_dict['gt'] = [] 146 | sent_dict['pred'] = [] 147 | with torch.no_grad(): 148 | for idx,(batch_X,batch_D) in enumerate(dataloader): 149 | batch_D = list(batch_D) 150 | sent_list, gt_list = eval_on_batch(batch_X,batch_D,model,tokenizer,device,config) 151 | sent_dict['pred'].extend(sent_list) 152 | sent_dict['gt'].extend(gt_list) 153 | 154 | with open(save_path, 'w') as f: 155 | json.dump(sent_dict, f,indent=4) 156 | 157 | 158 | def eval_on_batch(batch_X,batch_D,model,tokenizer,device,config): 159 | decode_method = config['decode'] 160 | padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 161 | if(not config['use_opt']): 162 | tokenizer.pad_token = tokenizer.eos_token 163 | batch_X = batch_X.to(device) 164 | print(f'batch_X:{batch_X.size()}') 165 | sent_list = [] 166 | gt_list = batch_D 167 | for i,hidden in enumerate(batch_X): 168 | inputs_embeds = hidden 169 | if(decode_method == 'beam'): 170 | #print('Using beam search decoding') 171 | if(config['use_opt']): 172 | sentence = decode_beam_search_opt.beam_decode_sentence(hidden_X=inputs_embeds, config = config,num_generate=1, beam_size = 5) 173 | else: 174 | sentence = beam_decode_sentence(hidden_X=inputs_embeds, config = config,num_generate=1, beam_size = 5) 175 | 176 | #print(sentence) 177 | sentence = sentence[0] 178 | else: 179 | sentence = generate_sentence(config,hidden_X=inputs_embeds) 180 | sent_list.append(sentence) 181 | 182 | 183 | 184 | return sent_list, gt_list 185 | 186 | 187 | def main(): 188 | parser = argparse.ArgumentParser(description='test') 189 | parser.add_argument('--model_dir', type=str, default='models/attacker_gpt2_persona_sbert', help='Dir of your model') 190 | parser.add_argument('--model_type', type=str, default='gpt-2', help='Type of the attacker model.') 191 | parser.add_argument('--data_type', type=str, default='test', help='Type of the processed data.') 192 | parser.add_argument('--save_path', type=str, default='logs/attacker_gpt2_p_sbert.log', help='Type of the processed data.') 193 | parser.add_argument('--num_epochs', type=int, default=1, help='Type of the processed data.') 194 | parser.add_argument('--p_simcse_flag', type=str2bool, default=True, help='Type of the processed data.') 195 | 196 | args = parser.parse_args() 197 | tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") 198 | token_num = len(tokenizer) 199 | config = {} 200 | config['model_dir'] = args.model_dir 201 | config['model_type'] = args.model_type 202 | config['num_epochs'] = args.num_epochs 203 | config['save_path'] = args.save_path 204 | config['p_simcse_flag'] = args.p_simcse_flag 205 | config['batch_size'] = attacker_models.batch_size 206 | config['data_type'] = args.data_type 207 | config['device'] = torch.device("cuda") 208 | 209 | config['use_trans'] = True 210 | config['model'] = AutoModelForCausalLM.from_pretrained(config['model_dir']) 211 | config['tokenizer'] = tokenizer 212 | config['token_num'] = token_num 213 | print('get_model done') 214 | dataloader = get_dataloader(config) 215 | print('get_dataloader done') 216 | eval(dataloader,config) 217 | if __name__ == '__main__': 218 | main() -------------------------------------------------------------------------------- /eval_ppl.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 3 | import sys 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from scipy.spatial.distance import cosine 9 | sys.path.append('..') 10 | import json 11 | import numpy as np 12 | import pandas as pd 13 | import argparse 14 | 15 | 16 | from transformers import AutoModel, AutoTokenizer 17 | from transformers import AutoModelForCausalLM 18 | from transformers import AdamW,get_linear_schedule_with_warmup 19 | from torch.utils.data import DataLoader, Dataset 20 | from attacker_models import SequenceCrossEntropyLoss 21 | from sentence_transformers import SentenceTransformer 22 | from simcse_persona import get_persona_dict 23 | from attacker_evaluation_gpt import eval_on_batch 24 | from datasets import load_dataset 25 | 26 | class linear_projection(nn.Module): 27 | def __init__(self, in_num, out_num=1024): 28 | super(linear_projection, self).__init__() 29 | self.fc1 = nn.Linear(in_num, out_num) 30 | 31 | def forward(self, x, use_final_hidden_only = True): 32 | # x should be of shape (?,in_num) according to gpt2 output 33 | out_shape = x.size()[-1] 34 | assert(x.size()[1] == out_shape) 35 | out = self.fc1(x) 36 | 37 | 38 | return out 39 | 40 | 41 | class personachat(Dataset): 42 | def __init__(self, data): 43 | self.data = data 44 | 45 | 46 | def __len__(self): 47 | return len(self.data) 48 | 49 | def __getitem__(self, index): 50 | 51 | text = self.data[index] 52 | 53 | return text 54 | 55 | def collate(self, unpacked_data): 56 | return unpacked_data 57 | 58 | def process_data(data,batch_size,device,config,need_porj=False): 59 | dataset = personachat(data) 60 | dataloader = DataLoader(dataset=dataset, 61 | shuffle=True, 62 | batch_size=batch_size, 63 | collate_fn=dataset.collate) 64 | 65 | print('load data done') 66 | 67 | ### extra projection 68 | if need_porj: 69 | projection = linear_projection(in_num=768).to(device) 70 | ### for attackers 71 | model_attacker = AutoModelForCausalLM.from_pretrained('dialogpt_qnli') 72 | tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 73 | criterion = SequenceCrossEntropyLoss() 74 | model_attacker.to(device) 75 | model_attacker.eval() 76 | param_optimizer = list(model_attacker.named_parameters()) 77 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 78 | optimizer_grouped_parameters = [ 79 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 80 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 81 | 82 | ] 83 | num_gradients_accumulation = 1 84 | num_epochs = config['num_epochs'] 85 | batch_size = config['batch_size'] 86 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 87 | optimizer = AdamW(optimizer_grouped_parameters, 88 | lr=3e-5, 89 | eps=1e-06) 90 | if need_porj: 91 | optimizer.add_param_group({'params': projection.parameters()}) 92 | scheduler = get_linear_schedule_with_warmup(optimizer, 93 | num_warmup_steps=100, 94 | num_training_steps = num_train_optimization_steps) 95 | 96 | ### process to obtain the embeddings 97 | for i in range(num_epochs): 98 | running_ppl = [] 99 | for idx,batch_text in enumerate(dataloader): 100 | 101 | record_loss, perplexity = train_on_batch(batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=False) 102 | 103 | running_ppl.append(perplexity) 104 | 105 | print(f'Validate ppl: {np.mean(running_ppl)}') 106 | 107 | 108 | 109 | def train_on_batch(batch_D,model,tokenizer,criterion,device,train=True): 110 | padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 111 | tokenizer.pad_token = tokenizer.eos_token 112 | inputs = tokenizer(batch_D, return_tensors='pt', padding='max_length', truncation=True, max_length=40) 113 | 114 | input_ids = inputs['input_ids'].to(device) # tensors of input ids 115 | labels = input_ids.clone() 116 | 117 | past = None 118 | 119 | logits, past = model(input_ids,past_key_values = past,return_dict=False) 120 | logits = logits[:, :-1].contiguous() 121 | target = labels[:, 1:].contiguous() 122 | 123 | target_mask = torch.ones_like(target).float() 124 | 125 | loss = criterion(logits, target, target_mask, label_smoothing=0.02, reduce="batch") 126 | 127 | record_loss = loss.item() 128 | perplexity = np.exp(record_loss) 129 | if train: 130 | loss.backward() 131 | 132 | return record_loss, perplexity 133 | 134 | 135 | def read_logs(path): 136 | with open(path) as f: 137 | data = json.load(f) 138 | pred = data["pred"] 139 | return pred 140 | 141 | def get_val_ppl(path,batch_size,device,config): 142 | sent_list = read_logs(path) 143 | process_data(sent_list,batch_size,device,config) 144 | 145 | def get_qnli_data(data_type): 146 | dataset = load_dataset('glue','qnli', cache_dir="/home/hlibt/embed_rev/data", split=data_type) 147 | sentence_list = [] 148 | for i,d in enumerate(dataset): 149 | sentence_list.append(d['question']) 150 | sentence_list.append(d['sentence']) 151 | return sentence_list 152 | def get_personachat_data(data_type): 153 | 154 | sent_list = get_persona_dict(data_type=data_type) 155 | return sent_list 156 | 157 | def get_sent_list(config): 158 | dataset = config['dataset'] 159 | data_type = config['data_type'] 160 | if dataset == 'personachat': 161 | sent_list = get_personachat_data(data_type) 162 | return sent_list 163 | elif dataset == 'qnli': 164 | sent_list = get_qnli_data(data_type) 165 | return sent_list 166 | else: 167 | print('Name of dataset only supports: personachat or qnli') 168 | sys.exit(-1) 169 | if __name__ == '__main__': 170 | model_cards ={} 171 | model_cards['sent_t5'] = 'sentence-t5-large' 172 | model_cards['mpnet'] = 'all-mpnet-base-v1' 173 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 174 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 175 | model_cards['simcse_roberta'] = 'princeton-nlp/sup-simcse-roberta-large' 176 | 177 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 178 | parser.add_argument('--model_dir', type=str, default='microsoft/DialoGPT-medium', help='Dir of your model') 179 | parser.add_argument('--num_epochs', type=int, default=1, help='Training epoches.') 180 | parser.add_argument('--batch_size', type=int, default=64, help='Batch_size #.') 181 | parser.add_argument('--dataset', type=str, default='personachat', help='Name of dataset: personachat or qnli') 182 | #parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 183 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 184 | #parser.add_argument('--data_type', type=str, default='test', help='train/test') 185 | parser.add_argument('--embed_model', type=str, default='sent_t5', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 186 | parser.add_argument('--decode', type=str, default='beam', help='Name of decoding methods: beam/sampling') 187 | #parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 188 | args = parser.parse_args() 189 | config = {} 190 | config['model_dir'] = args.model_dir 191 | config['num_epochs'] = args.num_epochs 192 | config['batch_size'] = args.batch_size 193 | config['dataset'] = args.dataset 194 | config['data_type'] = args.data_type 195 | config['embed_model'] = args.embed_model 196 | config['decode'] = args.decode 197 | config['embed_model_path'] = model_cards[config['embed_model']] 198 | config['device'] = torch.device("cuda") 199 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 200 | config['eos_token'] = config['tokenizer'].eos_token 201 | 202 | 203 | device = torch.device("cuda") 204 | 205 | batch_size = config['batch_size'] 206 | 207 | 208 | 209 | ### qnli with beam search decoding 210 | sbert_roberta_large_pc_path = '../models_random/attacker_gpt2_qnli_sent_roberta.log' 211 | simcse_roberta_large_pc_path = '../models_random/attacker_gpt2_qnli_simcse_roberta.log' 212 | simcse_bert_large_pc_path = '../models_random/attacker_gpt2_qnli_simcse_bert.log' 213 | sentence_T5_large_pc_path = '../models_random/attacker_gpt2_qnli_sent_t5.log' 214 | mpnet_pc_path = '../models_random/attacker_gpt2_qnli_mpnet.log' 215 | 216 | 217 | 218 | print('===mpnet===') 219 | get_val_ppl(mpnet_pc_path,batch_size,device,config) 220 | print('===sen_roberta===') 221 | get_val_ppl(sbert_roberta_large_pc_path,batch_size,device,config) 222 | print('===st5===') 223 | get_val_ppl(sentence_T5_large_pc_path,batch_size,device,config) 224 | print('===simcse_bert===') 225 | get_val_ppl(simcse_bert_large_pc_path,batch_size,device,config) 226 | print('===simcse_roberta===') 227 | get_val_ppl(simcse_roberta_large_pc_path,batch_size,device,config) 228 | 229 | 230 | -------------------------------------------------------------------------------- /attacker.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from scipy.spatial.distance import cosine 8 | 9 | import json 10 | import numpy as np 11 | import pandas as pd 12 | import argparse 13 | import sys 14 | 15 | from transformers import AutoModel, AutoTokenizer 16 | from transformers import AutoModelForCausalLM 17 | from transformers import AdamW,get_linear_schedule_with_warmup 18 | from torch.utils.data import DataLoader, Dataset 19 | from attacker_models import SequenceCrossEntropyLoss 20 | from sentence_transformers import SentenceTransformer 21 | from simcse_persona import get_persona_dict 22 | from attacker_evaluation_gpt import eval_on_batch 23 | from datasets import load_dataset 24 | from data_process import get_sent_list 25 | 26 | class linear_projection(nn.Module): 27 | def __init__(self, in_num, out_num=1024): 28 | super(linear_projection, self).__init__() 29 | self.fc1 = nn.Linear(in_num, out_num) 30 | 31 | def forward(self, x, use_final_hidden_only = True): 32 | # x should be of shape (?,in_num) according to gpt2 output 33 | out_shape = x.size()[-1] 34 | assert(x.size()[1] == out_shape) 35 | out = self.fc1(x) 36 | 37 | 38 | return out 39 | 40 | 41 | class personachat(Dataset): 42 | def __init__(self, data): 43 | self.data = data 44 | 45 | 46 | def __len__(self): 47 | return len(self.data) 48 | 49 | def __getitem__(self, index): 50 | 51 | text = self.data[index] 52 | 53 | return text 54 | 55 | def collate(self, unpacked_data): 56 | return unpacked_data 57 | 58 | def process_data(data,batch_size,device,config,need_porj=True): 59 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 60 | device_1 = torch.device("cuda:0") 61 | model = SentenceTransformer(config['embed_model_path'],device=device_1) # dim 768 62 | dataset = personachat(data) 63 | dataloader = DataLoader(dataset=dataset, 64 | shuffle=True, 65 | batch_size=batch_size, 66 | collate_fn=dataset.collate) 67 | 68 | print('load data done') 69 | ### extra projection 70 | if need_porj: 71 | projection = linear_projection(in_num=768, out_num=1280).to(device) 72 | ### for attackers 73 | model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 74 | tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 75 | criterion = SequenceCrossEntropyLoss() 76 | model_attacker.to(device) 77 | param_optimizer = list(model_attacker.named_parameters()) 78 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 79 | optimizer_grouped_parameters = [ 80 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 81 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 82 | 83 | ] 84 | num_gradients_accumulation = 1 85 | num_epochs = config['num_epochs'] 86 | batch_size = config['batch_size'] 87 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 88 | optimizer = AdamW(optimizer_grouped_parameters, 89 | lr=3e-5, 90 | eps=1e-06) 91 | if need_porj: 92 | optimizer.add_param_group({'params': projection.parameters()}) 93 | scheduler = get_linear_schedule_with_warmup(optimizer, 94 | num_warmup_steps=100, 95 | num_training_steps = num_train_optimization_steps) 96 | 97 | ### process to obtain the embeddings 98 | for i in range(num_epochs): 99 | model.eval() 100 | for idx,batch_text in enumerate(dataloader): 101 | with torch.no_grad(): 102 | embeddings = model.encode(batch_text,convert_to_tensor = True).to(device) 103 | print(f'Embedding dim: {embeddings.size()}') 104 | 105 | ### attacker part, needs training 106 | if need_porj: 107 | embeddings = projection(embeddings) 108 | 109 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 110 | optimizer.step() 111 | scheduler.step() 112 | # make sure no grad for GPT optimizer 113 | optimizer.zero_grad() 114 | print(f'Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 115 | #sys.exit(-1) 116 | if need_porj: 117 | proj_path = 'models/' + 'projection_gpt2_large_' + config['dataset'] + '_' + config['embed_model'] 118 | torch.save(projection.state_dict(), proj_path) 119 | save_path = 'models/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model'] 120 | model_attacker.save_pretrained(save_path) 121 | 122 | 123 | ### used for testing only 124 | def process_data_test(data,batch_size,device,config,need_proj=True): 125 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 126 | #model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 127 | device_1 = torch.device("cuda:0") 128 | model = SentenceTransformer(config['embed_model_path'],device=device_1) # dim 768 129 | if(config['decode'] == 'beam'): 130 | save_path = 'models/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 131 | else: 132 | save_path = 'models/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model']+'.log' 133 | dataset = personachat(data) 134 | # no shuffle for testing data 135 | dataloader = DataLoader(dataset=dataset, 136 | shuffle=False, 137 | batch_size=batch_size, 138 | collate_fn=dataset.collate) 139 | 140 | print('load data done') 141 | if need_proj: 142 | proj_path = 'models/' + 'projection_gpt2_large_' + config['dataset'] + '_' + config['embed_model'] 143 | projection = linear_projection(in_num=768, out_num=1280) 144 | projection.load_state_dict(torch.load(proj_path)) 145 | projection.to(device) 146 | print('load projection done') 147 | else: 148 | print('no projection loaded') 149 | # setup on config for sentence generation AutoModelForCausalLM 150 | attacker_path = 'models/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model'] 151 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 152 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-large') 153 | 154 | sent_dict = {} 155 | sent_dict['gt'] = [] 156 | sent_dict['pred'] = [] 157 | with torch.no_grad(): 158 | for idx,batch_text in enumerate(dataloader): 159 | 160 | embeddings = model.encode(batch_text,convert_to_tensor = True).to(device) 161 | 162 | if need_proj: 163 | embeddings = projection(embeddings) 164 | 165 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 166 | print(f'testing {idx} batch done with {idx*batch_size} samples') 167 | sent_dict['pred'].extend(sent_list) 168 | sent_dict['gt'].extend(gt_list) 169 | 170 | with open(save_path, 'w') as f: 171 | json.dump(sent_dict, f,indent=4) 172 | 173 | return 0 174 | 175 | 176 | ### used for testing only 177 | def process_data_test_simcse(data,batch_size,device,config,proj_dir=None,need_proj=False): 178 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 179 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 180 | #save_path = 'logs/attacker_gpt2_qnli_simcse_bert_large.log' 181 | if(config['decode'] == 'beam'): 182 | print('Using beam search decoding') 183 | save_path = 'logs/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 184 | else: 185 | save_path = 'logs/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model']+'.log' 186 | dataset = personachat(data) 187 | # no shuffle for testing data 188 | dataloader = DataLoader(dataset=dataset, 189 | shuffle=False, 190 | batch_size=batch_size, 191 | collate_fn=dataset.collate) 192 | 193 | print('load data done') 194 | if need_proj: 195 | projection = linear_projection(in_num=768) 196 | projection.load_state_dict(torch.load(proj_dir)) 197 | projection.to(device) 198 | print('load projection done') 199 | else: 200 | print('no projection loaded') 201 | # setup on config for sentence generation AutoModelForCausalLM 202 | attacker_path = 'models/' + 'attacker_gpt2_large_' + config['dataset'] + '_' + config['embed_model'] 203 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 204 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-large') 205 | 206 | sent_dict = {} 207 | sent_dict['gt'] = [] 208 | sent_dict['pred'] = [] 209 | with torch.no_grad(): 210 | for idx,batch_text in enumerate(dataloader): 211 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 212 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 213 | if need_proj: 214 | embeddings = projection(embeddings) 215 | #sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 216 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 217 | print(f'testing {idx} batch done with {idx*batch_size} samples') 218 | sent_dict['pred'].extend(sent_list) 219 | sent_dict['gt'].extend(gt_list) 220 | with open(save_path, 'w') as f: 221 | json.dump(sent_dict, f,indent=4) 222 | 223 | return 0 224 | 225 | 226 | def train_on_batch(batch_X,batch_D,model,tokenizer,criterion,device,train=True): 227 | padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 228 | tokenizer.pad_token = tokenizer.eos_token 229 | inputs = tokenizer(batch_D, return_tensors='pt', padding='max_length', truncation=True, max_length=40) 230 | #dial_tokens = [tokenizer.encode(item) + turn_ending for item in batch_D] 231 | #print(inputs) 232 | input_ids = inputs['input_ids'].to(device) # tensors of input ids 233 | labels = input_ids.clone() 234 | #print(input_ids.size()) 235 | # embed the input ids using GPT-2 embedding 236 | input_emb = model.transformer.wte(input_ids) 237 | # add extra dim to cat together 238 | batch_X = batch_X.to(device) 239 | batch_X_unsqueeze = torch.unsqueeze(batch_X, 1) 240 | inputs_embeds = torch.cat((batch_X_unsqueeze,input_emb),dim=1) #[batch,max_length+1,emb_dim (1024)] 241 | past = None 242 | # need to move to device later 243 | inputs_embeds = inputs_embeds 244 | 245 | #logits, past = model(inputs_embeds=inputs_embeds,past = past) 246 | logits, past = model(inputs_embeds=inputs_embeds,past_key_values = past,return_dict=False) 247 | logits = logits[:, :-1].contiguous() 248 | target = labels.contiguous() 249 | target_mask = torch.ones_like(target).float() 250 | loss = criterion(logits, target, target_mask, label_smoothing=0.02, reduce="batch") 251 | 252 | record_loss = loss.item() 253 | perplexity = np.exp(record_loss) 254 | if train: 255 | loss.backward() 256 | 257 | return record_loss, perplexity 258 | 259 | 260 | 261 | 262 | 263 | 264 | if __name__ == '__main__': 265 | model_cards ={} 266 | model_cards['sent_t5_large'] = 'sentence-t5-large' 267 | model_cards['sent_t5_base'] = 'sentence-t5-base' 268 | model_cards['sent_t5_xl'] = 'sentence-t5-xl' 269 | model_cards['sent_t5_xxl'] = 'sentence-t5-xxl' 270 | model_cards['mpnet'] = 'all-mpnet-base-v1' 271 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 272 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 273 | model_cards['simcse_roberta'] = 'princeton-nlp/sup-simcse-roberta-large' 274 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 275 | parser.add_argument('--model_dir', type=str, default='microsoft/DialoGPT-large', help='Dir of your model') 276 | parser.add_argument('--num_epochs', type=int, default=10, help='Training epoches.') 277 | parser.add_argument('--batch_size', type=int, default=16, help='Batch_size #.') 278 | parser.add_argument('--dataset', type=str, default='personachat', help='Name of dataset: personachat or qnli') 279 | #parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 280 | #parser.add_argument('--data_type', type=str, default='train', help='train/test') 281 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 282 | parser.add_argument('--embed_model', type=str, default='sent_t5_base', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 283 | parser.add_argument('--decode', type=str, default='beam', help='Name of decoding methods: beam/sampling') 284 | #parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 285 | args = parser.parse_args() 286 | config = {} 287 | config['model_dir'] = args.model_dir 288 | config['num_epochs'] = args.num_epochs 289 | config['batch_size'] = args.batch_size 290 | config['dataset'] = args.dataset 291 | config['data_type'] = args.data_type 292 | config['embed_model'] = args.embed_model 293 | config['decode'] = args.decode 294 | config['embed_model_path'] = model_cards[config['embed_model']] 295 | config['device'] = torch.device("cuda") 296 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-large') 297 | config['eos_token'] = config['tokenizer'].eos_token 298 | config['use_opt'] = False 299 | 300 | 301 | device = torch.device("cuda:0") 302 | #device = torch.device("cpu") 303 | batch_size = config['batch_size'] 304 | 305 | sent_list = get_sent_list(config) 306 | 307 | ##### for training 308 | if(config['data_type'] == 'train'): 309 | process_data(sent_list,batch_size,device,config) 310 | elif(config['data_type'] == 'test'): 311 | if('simcse' in config['embed_model']): 312 | process_data_test_simcse(sent_list,batch_size,device,config,proj_dir=None,need_proj=False) 313 | else: 314 | process_data_test(sent_list,batch_size,device,config,need_proj=True) 315 | 316 | -------------------------------------------------------------------------------- /attacker_opt.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from scipy.spatial.distance import cosine 8 | 9 | import json 10 | import numpy as np 11 | import pandas as pd 12 | import argparse 13 | import sys 14 | 15 | from transformers import AutoModel, AutoTokenizer 16 | from transformers import AutoModelForCausalLM,GPT2Config,GPT2LMHeadModel 17 | from transformers import AdamW,get_linear_schedule_with_warmup 18 | from transformers import OPTForCausalLM,GPT2Tokenizer 19 | 20 | from torch.utils.data import DataLoader, Dataset 21 | from attacker_models import SequenceCrossEntropyLoss 22 | from sentence_transformers import SentenceTransformer 23 | from simcse_persona import get_persona_dict 24 | from attacker_evaluation_gpt import eval_on_batch 25 | from datasets import load_dataset 26 | from data_process import get_sent_list 27 | 28 | Folder_path = 'models/' 29 | #Folder_path = 'opt_models/' 30 | class linear_projection(nn.Module): 31 | def __init__(self, in_num, out_num=2048): 32 | super(linear_projection, self).__init__() 33 | self.fc1 = nn.Linear(in_num, out_num) 34 | 35 | def forward(self, x, use_final_hidden_only = True): 36 | # x should be of shape (?,in_num) according to gpt2 output 37 | out_shape = x.size()[-1] 38 | assert(x.size()[1] == out_shape) 39 | out = self.fc1(x) 40 | 41 | 42 | return out 43 | 44 | 45 | class personachat(Dataset): 46 | def __init__(self, data): 47 | self.data = data 48 | 49 | 50 | def __len__(self): 51 | return len(self.data) 52 | 53 | def __getitem__(self, index): 54 | 55 | text = self.data[index] 56 | 57 | return text 58 | 59 | def collate(self, unpacked_data): 60 | return unpacked_data 61 | 62 | 63 | def init_opt(): 64 | #opt-350m 65 | model = OPTForCausalLM.from_pretrained('facebook/opt-350m') 66 | tokenizer = GPT2Tokenizer.from_pretrained("facebook/opt-350m") 67 | #tokenizer = AutoTokenizer.from_pretrained('facebook/opt-1.3b',use_fast=False) 68 | #model = GPT2LMHeadModel(config) 69 | return model, tokenizer 70 | 71 | def process_data(data,batch_size,device,config,need_proj=True): 72 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 73 | embed_model_name = config['embed_model'] 74 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 75 | dataset = personachat(data) 76 | dataloader = DataLoader(dataset=dataset, 77 | shuffle=True, 78 | batch_size=batch_size, 79 | collate_fn=dataset.collate) 80 | 81 | print('load data done') 82 | 83 | ### extra projection 84 | if need_proj: 85 | projection = linear_projection(in_num=768).to(device) 86 | ### for attackers 87 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 88 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 89 | model_attacker, tokenizer_attacker = init_opt() 90 | criterion = SequenceCrossEntropyLoss() 91 | model_attacker.to(device) 92 | param_optimizer = list(model_attacker.named_parameters()) 93 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 94 | optimizer_grouped_parameters = [ 95 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 96 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 97 | 98 | ] 99 | num_gradients_accumulation = 1 100 | num_epochs = config['num_epochs'] 101 | batch_size = config['batch_size'] 102 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 103 | optimizer = AdamW(optimizer_grouped_parameters, 104 | lr=3e-5, 105 | eps=1e-06) 106 | if need_proj: 107 | optimizer.add_param_group({'params': projection.parameters()}) 108 | scheduler = get_linear_schedule_with_warmup(optimizer, 109 | num_warmup_steps=100, 110 | num_training_steps = num_train_optimization_steps) 111 | 112 | ### process to obtain the embeddings 113 | for i in range(num_epochs): 114 | for idx,batch_text in enumerate(dataloader): 115 | with torch.no_grad(): 116 | 117 | #sys.exit(-1) 118 | embeddings = model.encode(batch_text,convert_to_tensor = True) 119 | print(f'Embedding dim: {embeddings.size()}') 120 | 121 | ### attacker part, needs training 122 | if need_proj: 123 | embeddings = projection(embeddings) 124 | 125 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 126 | optimizer.step() 127 | scheduler.step() 128 | # make sure no grad for GPT optimizer 129 | optimizer.zero_grad() 130 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 131 | #sys.exit(-1) 132 | if need_proj: 133 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 134 | torch.save(projection.state_dict(), proj_path) 135 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 136 | model_attacker.save_pretrained(save_path) 137 | 138 | 139 | 140 | def process_data_simcse(data,batch_size,device,config,need_proj=False): 141 | embed_model_name = config['embed_model'] 142 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 143 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 144 | dataset = personachat(data) 145 | dataloader = DataLoader(dataset=dataset, 146 | shuffle=True, 147 | batch_size=batch_size, 148 | collate_fn=dataset.collate) 149 | 150 | print('load data done') 151 | 152 | ### extra projection 153 | if need_proj: 154 | projection = linear_projection(in_num=768).to(device) 155 | 156 | ### for attackers 157 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 158 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 159 | model_attacker, tokenizer_attacker = init_opt() 160 | criterion = SequenceCrossEntropyLoss() 161 | model_attacker.to(device) 162 | param_optimizer = list(model_attacker.named_parameters()) 163 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 164 | optimizer_grouped_parameters = [ 165 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 166 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 167 | 168 | ] 169 | num_gradients_accumulation = 1 170 | num_epochs = config['num_epochs'] 171 | batch_size = config['batch_size'] 172 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 173 | optimizer = AdamW(optimizer_grouped_parameters, 174 | lr=3e-5, 175 | eps=1e-06) 176 | if need_proj: 177 | optimizer.add_param_group({'params': projection.parameters()}) 178 | scheduler = get_linear_schedule_with_warmup(optimizer, 179 | num_warmup_steps=100, 180 | num_training_steps = num_train_optimization_steps) 181 | 182 | ### process to obtain the embeddings 183 | for i in range(num_epochs): 184 | for idx,batch_text in enumerate(dataloader): 185 | with torch.no_grad(): 186 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 187 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 188 | print(embeddings.size()) 189 | 190 | 191 | ### attacker part, needs training 192 | if need_proj: 193 | embeddings = projection(embeddings) 194 | 195 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 196 | optimizer.step() 197 | scheduler.step() 198 | # make sure no grad for GPT optimizer 199 | optimizer.zero_grad() 200 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 201 | #sys.exit(-1) 202 | if need_proj: 203 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 204 | torch.save(projection.state_dict(), proj_path) 205 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 206 | model_attacker.save_pretrained(save_path) 207 | 208 | 209 | ### used for testing only 210 | def process_data_test(data,batch_size,device,config,need_proj=False): 211 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 212 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 213 | if(config['decode'] == 'beam'): 214 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 215 | else: 216 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'.log' 217 | dataset = personachat(data) 218 | # no shuffle for testing data 219 | dataloader = DataLoader(dataset=dataset, 220 | shuffle=False, 221 | batch_size=batch_size, 222 | collate_fn=dataset.collate) 223 | 224 | print('load data done') 225 | if need_proj: 226 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 227 | projection = linear_projection(in_num=768) 228 | projection.load_state_dict(torch.load(proj_path)) 229 | projection.to(device) 230 | print('load projection done') 231 | else: 232 | print('no projection loaded') 233 | # setup on config for sentence generation AutoModelForCausalLM 234 | attacker_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 235 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 236 | #config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 237 | 238 | sent_dict = {} 239 | sent_dict['gt'] = [] 240 | sent_dict['pred'] = [] 241 | with torch.no_grad(): 242 | for idx,batch_text in enumerate(dataloader): 243 | embeddings = model.encode(batch_text,convert_to_tensor = True) 244 | #inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 245 | #embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 246 | if need_proj: 247 | embeddings = projection(embeddings) 248 | #embeddings_out = embeddings_proj.detach().cpu() 249 | #data_dict['text'].extend(batch_text) 250 | #data_dict['embedding'].extend(embeddings_out) 251 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 252 | print(f'testing {idx} batch done with {idx*batch_size} samples') 253 | sent_dict['pred'].extend(sent_list) 254 | sent_dict['gt'].extend(gt_list) 255 | 256 | with open(save_path, 'w') as f: 257 | json.dump(sent_dict, f,indent=4) 258 | 259 | return 0 260 | 261 | 262 | ### used for testing only 263 | def process_data_test_simcse(data,batch_size,device,config,proj_dir=None,need_proj=False): 264 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 265 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 266 | #save_path = 'logs/attacker_opt_qnli_simcse_bert_large.log' 267 | if(config['decode'] == 'beam'): 268 | print('Using beam search decoding') 269 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 270 | else: 271 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'.log' 272 | dataset = personachat(data) 273 | # no shuffle for testing data 274 | dataloader = DataLoader(dataset=dataset, 275 | shuffle=False, 276 | batch_size=batch_size, 277 | collate_fn=dataset.collate) 278 | 279 | print('load data done') 280 | if need_proj: 281 | projection = linear_projection(in_num=768) 282 | projection.load_state_dict(torch.load(proj_dir)) 283 | projection.to(device) 284 | print('load projection done') 285 | else: 286 | print('no projection loaded') 287 | # setup on config for sentence generation AutoModelForCausalLM 288 | attacker_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 289 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 290 | #config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 291 | 292 | sent_dict = {} 293 | sent_dict['gt'] = [] 294 | sent_dict['pred'] = [] 295 | with torch.no_grad(): 296 | for idx,batch_text in enumerate(dataloader): 297 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 298 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 299 | if need_proj: 300 | embeddings = projection(embeddings) 301 | #sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 302 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 303 | print(f'testing {idx} batch done with {idx*batch_size} samples') 304 | sent_dict['pred'].extend(sent_list) 305 | sent_dict['gt'].extend(gt_list) 306 | with open(save_path, 'w') as f: 307 | json.dump(sent_dict, f,indent=4) 308 | 309 | return 0 310 | 311 | 312 | def train_on_batch(batch_X,batch_D,model,tokenizer,criterion,device,train=True): 313 | padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 314 | tokenizer.pad_token = tokenizer.eos_token 315 | inputs = tokenizer(batch_D, return_tensors='pt', padding='max_length', truncation=True, max_length=40) 316 | #dial_tokens = [tokenizer.encode(item) + turn_ending for item in batch_D] 317 | #print(inputs) 318 | input_ids = inputs['input_ids'].to(device) # tensors of input ids 319 | labels = input_ids.clone() 320 | #print(input_ids.size()) 321 | # embed the input ids using GPT-2 embedding 322 | input_emb = model.model.decoder.embed_tokens(input_ids) 323 | # add extra dim to cat together 324 | batch_X = batch_X.to(device) 325 | batch_X_unsqueeze = torch.unsqueeze(batch_X, 1) 326 | inputs_embeds = torch.cat((batch_X_unsqueeze,input_emb),dim=1) #[batch,max_length+1,emb_dim (1024)] 327 | past = None 328 | # need to move to device later 329 | inputs_embeds = inputs_embeds 330 | 331 | target = labels.contiguous() 332 | output = model(inputs_embeds=inputs_embeds,past_key_values = past,return_dict=True) 333 | logits = output.logits 334 | logits = logits[:, :-1].contiguous() 335 | target_mask = torch.ones_like(target).float() 336 | loss = criterion(logits, target, target_mask, label_smoothing=0.02, reduce="batch") 337 | 338 | record_loss = loss.item() 339 | perplexity = np.exp(record_loss) 340 | if train: 341 | loss.backward() 342 | 343 | return record_loss, perplexity 344 | 345 | 346 | 347 | 348 | 349 | 350 | if __name__ == '__main__': 351 | ''' 352 | Sentence bert based: 353 | T5: sentence-t5-large dim 768 354 | mpnet: all-mpnet-base-v1 dim 768 355 | Roberta: all-roberta-large-v1 dim 1024 356 | 357 | 358 | SIMCSE based: 359 | princeton-nlp/unsup-simcse-roberta-large dim 1024 360 | princeton-nlp/sup-simcse-bert-large-uncased dim 1024 361 | ''' 362 | model_cards ={} 363 | model_cards['sent_t5'] = 'sentence-t5-large' 364 | model_cards['mpnet'] = 'all-mpnet-base-v1' 365 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 366 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 367 | model_cards['simcse_roberta'] = 'princeton-nlp/sup-simcse-roberta-large' 368 | 369 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 370 | parser.add_argument('--model_dir', type=str, default='opt', help='Dir of your model') 371 | parser.add_argument('--num_epochs', type=int, default=10, help='Training epoches.') 372 | parser.add_argument('--batch_size', type=int, default=32, help='Batch_size #.') 373 | parser.add_argument('--dataset', type=str, default='personachat', help='Name of dataset: personachat or qnli') 374 | #parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 375 | #parser.add_argument('--data_type', type=str, default='train', help='train/test') 376 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 377 | parser.add_argument('--embed_model', type=str, default='sent_t5', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 378 | parser.add_argument('--decode', type=str, default='beam', help='Name of decoding methods: beam/sampling') 379 | #parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 380 | args = parser.parse_args() 381 | config = {} 382 | config['model_dir'] = args.model_dir 383 | config['num_epochs'] = args.num_epochs 384 | config['batch_size'] = args.batch_size 385 | config['dataset'] = args.dataset 386 | config['data_type'] = args.data_type 387 | config['embed_model'] = args.embed_model 388 | config['decode'] = args.decode 389 | config['embed_model_path'] = model_cards[config['embed_model']] 390 | config['device'] = torch.device("cuda") 391 | config['tokenizer'] = AutoTokenizer.from_pretrained('facebook/opt-1.3b') 392 | config['eos_token'] = config['tokenizer'].eos_token 393 | 394 | config['use_opt'] = True 395 | device = torch.device("cuda") 396 | #device = torch.device("cpu") 397 | batch_size = config['batch_size'] 398 | 399 | 400 | sent_list = get_sent_list(config) 401 | 402 | ##### for training 403 | if(config['data_type'] == 'train'): 404 | process_data(sent_list,batch_size,device,config) 405 | elif(config['data_type'] == 'test'): 406 | if('simcse' in config['embed_model']): 407 | process_data_test_simcse(sent_list,batch_size,device,config,proj_dir=None,need_proj=False) 408 | else: 409 | process_data_test(sent_list,batch_size,device,config,need_proj=True) 410 | -------------------------------------------------------------------------------- /attacker_random_gpt2.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from scipy.spatial.distance import cosine 8 | 9 | import json 10 | import numpy as np 11 | import pandas as pd 12 | import argparse 13 | import sys 14 | 15 | from transformers import AutoModel, AutoTokenizer 16 | from transformers import AutoModelForCausalLM,GPT2Config,GPT2LMHeadModel 17 | from transformers import AdamW,get_linear_schedule_with_warmup 18 | from torch.utils.data import DataLoader, Dataset 19 | from attacker_models import SequenceCrossEntropyLoss 20 | from sentence_transformers import SentenceTransformer 21 | from simcse_persona import get_persona_dict 22 | from attacker_evaluation_gpt import eval_on_batch 23 | from datasets import load_dataset 24 | from data_process import get_sent_list 25 | 26 | #abcd_path = config.abcd_path 27 | model_folder_path = 'models/' 28 | 29 | class linear_projection(nn.Module): 30 | def __init__(self, in_num, out_num=1024): 31 | super(linear_projection, self).__init__() 32 | self.fc1 = nn.Linear(in_num, out_num) 33 | 34 | def forward(self, x, use_final_hidden_only = True): 35 | # x should be of shape (?,in_num) according to gpt2 output 36 | out_shape = x.size()[-1] 37 | assert(x.size()[1] == out_shape) 38 | out = self.fc1(x) 39 | 40 | 41 | return out 42 | 43 | 44 | class personachat(Dataset): 45 | def __init__(self, data): 46 | self.data = data 47 | 48 | 49 | def __len__(self): 50 | return len(self.data) 51 | 52 | def __getitem__(self, index): 53 | 54 | text = self.data[index] 55 | 56 | return text 57 | 58 | def collate(self, unpacked_data): 59 | return unpacked_data 60 | 61 | 62 | def init_gpt2(): 63 | config = GPT2Config.from_pretrained('microsoft/DialoGPT-medium') 64 | tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 65 | model = GPT2LMHeadModel(config) 66 | return model, tokenizer 67 | 68 | def process_data(data,batch_size,device,config,need_proj=True): 69 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 70 | embed_model_name = config['embed_model'] 71 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 72 | dataset = personachat(data) 73 | dataloader = DataLoader(dataset=dataset, 74 | shuffle=True, 75 | batch_size=batch_size, 76 | collate_fn=dataset.collate) 77 | 78 | print('load data done') 79 | 80 | ### extra projection 81 | if need_proj: 82 | projection = linear_projection(in_num=768).to(device) 83 | ### for attackers 84 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 85 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 86 | model_attacker, tokenizer_attacker = init_gpt2() 87 | criterion = SequenceCrossEntropyLoss() 88 | model_attacker.to(device) 89 | param_optimizer = list(model_attacker.named_parameters()) 90 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 91 | optimizer_grouped_parameters = [ 92 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 93 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 94 | 95 | ] 96 | num_gradients_accumulation = 1 97 | num_epochs = config['num_epochs'] 98 | batch_size = config['batch_size'] 99 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 100 | optimizer = AdamW(optimizer_grouped_parameters, 101 | lr=3e-5, 102 | eps=1e-06) 103 | if need_proj: 104 | optimizer.add_param_group({'params': projection.parameters()}) 105 | scheduler = get_linear_schedule_with_warmup(optimizer, 106 | num_warmup_steps=100, 107 | num_training_steps = num_train_optimization_steps) 108 | 109 | ### process to obtain the embeddings 110 | for i in range(num_epochs): 111 | for idx,batch_text in enumerate(dataloader): 112 | with torch.no_grad(): 113 | 114 | #sys.exit(-1) 115 | embeddings = model.encode(batch_text,convert_to_tensor = True) 116 | print(f'Embedding dim: {embeddings.size()}') 117 | 118 | ### attacker part, needs training 119 | if need_proj: 120 | embeddings = projection(embeddings) 121 | 122 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 123 | optimizer.step() 124 | scheduler.step() 125 | # make sure no grad for GPT optimizer 126 | optimizer.zero_grad() 127 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 128 | #sys.exit(-1) 129 | if need_proj: 130 | proj_path = model_folder_path + 'projection_gpt2_' + config['dataset'] + '_' + config['embed_model'] 131 | torch.save(projection.state_dict(), proj_path) 132 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model'] 133 | model_attacker.save_pretrained(save_path) 134 | 135 | 136 | 137 | def process_data_simcse(data,batch_size,device,config,need_proj=False): 138 | embed_model_name = config['embed_model'] 139 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 140 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 141 | dataset = personachat(data) 142 | dataloader = DataLoader(dataset=dataset, 143 | shuffle=True, 144 | batch_size=batch_size, 145 | collate_fn=dataset.collate) 146 | 147 | print('load data done') 148 | 149 | ### extra projection 150 | if need_proj: 151 | projection = linear_projection(in_num=768).to(device) 152 | 153 | ### for attackers 154 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 155 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 156 | model_attacker, tokenizer_attacker = init_gpt2() 157 | criterion = SequenceCrossEntropyLoss() 158 | model_attacker.to(device) 159 | param_optimizer = list(model_attacker.named_parameters()) 160 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 161 | optimizer_grouped_parameters = [ 162 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 163 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 164 | 165 | ] 166 | num_gradients_accumulation = 1 167 | num_epochs = config['num_epochs'] 168 | batch_size = config['batch_size'] 169 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 170 | optimizer = AdamW(optimizer_grouped_parameters, 171 | lr=3e-5, 172 | eps=1e-06) 173 | if need_proj: 174 | optimizer.add_param_group({'params': projection.parameters()}) 175 | scheduler = get_linear_schedule_with_warmup(optimizer, 176 | num_warmup_steps=100, 177 | num_training_steps = num_train_optimization_steps) 178 | 179 | ### process to obtain the embeddings 180 | for i in range(num_epochs): 181 | for idx,batch_text in enumerate(dataloader): 182 | with torch.no_grad(): 183 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 184 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 185 | print(embeddings.size()) 186 | 187 | 188 | ### attacker part, needs training 189 | if need_proj: 190 | embeddings = projection(embeddings) 191 | 192 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 193 | optimizer.step() 194 | scheduler.step() 195 | # make sure no grad for GPT optimizer 196 | optimizer.zero_grad() 197 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 198 | #sys.exit(-1) 199 | if need_proj: 200 | proj_path = model_folder_path + 'projection_gpt2_' + config['dataset'] + '_' + config['embed_model'] 201 | torch.save(projection.state_dict(), proj_path) 202 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model'] 203 | model_attacker.save_pretrained(save_path) 204 | 205 | ### used for testing only 206 | def process_data_test(data,batch_size,device,config,need_proj=False): 207 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 208 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 209 | if(config['decode'] == 'beam'): 210 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 211 | else: 212 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model']+'.log' 213 | dataset = personachat(data) 214 | # no shuffle for testing data 215 | dataloader = DataLoader(dataset=dataset, 216 | shuffle=False, 217 | batch_size=batch_size, 218 | collate_fn=dataset.collate) 219 | 220 | print('load data done') 221 | if need_proj: 222 | proj_path = model_folder_path + 'projection_gpt2_' + config['dataset'] + '_' + config['embed_model'] 223 | projection = linear_projection(in_num=768) 224 | projection.load_state_dict(torch.load(proj_path)) 225 | projection.to(device) 226 | print('load projection done') 227 | else: 228 | print('no projection loaded') 229 | # setup on config for sentence generation AutoModelForCausalLM 230 | attacker_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model'] 231 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 232 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 233 | 234 | sent_dict = {} 235 | sent_dict['gt'] = [] 236 | sent_dict['pred'] = [] 237 | with torch.no_grad(): 238 | for idx,batch_text in enumerate(dataloader): 239 | embeddings = model.encode(batch_text,convert_to_tensor = True) 240 | #inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 241 | #embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 242 | if need_proj: 243 | embeddings = projection(embeddings) 244 | #embeddings_out = embeddings_proj.detach().cpu() 245 | #data_dict['text'].extend(batch_text) 246 | #data_dict['embedding'].extend(embeddings_out) 247 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 248 | print(f'testing {idx} batch done with {idx*batch_size} samples') 249 | sent_dict['pred'].extend(sent_list) 250 | sent_dict['gt'].extend(gt_list) 251 | 252 | with open(save_path, 'w') as f: 253 | json.dump(sent_dict, f,indent=4) 254 | 255 | return 0 256 | 257 | 258 | ### used for testing only 259 | def process_data_test_simcse(data,batch_size,device,config,proj_dir=None,need_proj=False): 260 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 261 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 262 | #save_path = 'logs/attacker_gpt2_qnli_simcse_bert_large.log' 263 | if(config['decode'] == 'beam'): 264 | print('Using beam search decoding') 265 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 266 | else: 267 | save_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model']+'.log' 268 | dataset = personachat(data) 269 | # no shuffle for testing data 270 | dataloader = DataLoader(dataset=dataset, 271 | shuffle=False, 272 | batch_size=batch_size, 273 | collate_fn=dataset.collate) 274 | 275 | print('load data done') 276 | if need_proj: 277 | projection = linear_projection(in_num=768) 278 | projection.load_state_dict(torch.load(proj_dir)) 279 | projection.to(device) 280 | print('load projection done') 281 | else: 282 | print('no projection loaded') 283 | # setup on config for sentence generation AutoModelForCausalLM 284 | attacker_path = model_folder_path + 'attacker_gpt2_' + config['dataset'] + '_' + config['embed_model'] 285 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 286 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 287 | 288 | sent_dict = {} 289 | sent_dict['gt'] = [] 290 | sent_dict['pred'] = [] 291 | with torch.no_grad(): 292 | for idx,batch_text in enumerate(dataloader): 293 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 294 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 295 | if need_proj: 296 | embeddings = projection(embeddings) 297 | #sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 298 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 299 | print(f'testing {idx} batch done with {idx*batch_size} samples') 300 | sent_dict['pred'].extend(sent_list) 301 | sent_dict['gt'].extend(gt_list) 302 | with open(save_path, 'w') as f: 303 | json.dump(sent_dict, f,indent=4) 304 | 305 | return 0 306 | 307 | 308 | def train_on_batch(batch_X,batch_D,model,tokenizer,criterion,device,train=True): 309 | padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 310 | tokenizer.pad_token = tokenizer.eos_token 311 | inputs = tokenizer(batch_D, return_tensors='pt', padding='max_length', truncation=True, max_length=40) 312 | #dial_tokens = [tokenizer.encode(item) + turn_ending for item in batch_D] 313 | #print(inputs) 314 | input_ids = inputs['input_ids'].to(device) # tensors of input ids 315 | labels = input_ids.clone() 316 | #print(input_ids.size()) 317 | # embed the input ids using GPT-2 embedding 318 | input_emb = model.transformer.wte(input_ids) 319 | # add extra dim to cat together 320 | batch_X = batch_X.to(device) 321 | batch_X_unsqueeze = torch.unsqueeze(batch_X, 1) 322 | inputs_embeds = torch.cat((batch_X_unsqueeze,input_emb),dim=1) #[batch,max_length+1,emb_dim (1024)] 323 | past = None 324 | # print(batch_X_unsqueeze.size()) 325 | # print(input_emb.size()) 326 | # print(inputs_embeds.size()) 327 | # print((inputs_embeds[:,0,:]).size()) 328 | # print(torch.sum(torch.abs(batch_X - inputs_embeds[:,0,:]))) 329 | #sys.exit(-1) 330 | # need to move to device later 331 | inputs_embeds = inputs_embeds 332 | 333 | #logits, past = model(inputs_embeds=inputs_embeds,past = past) 334 | logits, past = model(inputs_embeds=inputs_embeds,past_key_values = past,return_dict=False) 335 | logits = logits[:, :-1].contiguous() 336 | target = labels.contiguous() 337 | 338 | target_mask = torch.ones_like(target).float() 339 | #print(logits.size()) 340 | #print(target.size()) 341 | loss = criterion(logits, target, target_mask, label_smoothing=0.02, reduce="batch") 342 | 343 | record_loss = loss.item() 344 | perplexity = np.exp(record_loss) 345 | if train: 346 | loss.backward() 347 | 348 | return record_loss, perplexity 349 | 350 | 351 | 352 | 353 | def get_qnli_data(data_type): 354 | dataset = load_dataset('glue','qnli', cache_dir="/home/hlibt/embed_rev/data", split=data_type) 355 | sentence_list = [] 356 | for i,d in enumerate(dataset): 357 | sentence_list.append(d['question']) 358 | sentence_list.append(d['sentence']) 359 | return sentence_list 360 | def get_personachat_data(data_type): 361 | 362 | sent_list = get_persona_dict(data_type=data_type) 363 | return sent_list 364 | ''' 365 | def get_sent_list(config): 366 | dataset = config['dataset'] 367 | data_type = config['data_type'] 368 | if dataset == 'personachat': 369 | sent_list = get_personachat_data(data_type) 370 | return sent_list 371 | elif dataset == 'qnli': 372 | sent_list = get_qnli_data(data_type) 373 | return sent_list 374 | else: 375 | print('Name of dataset only supports: personachat or qnli') 376 | sys.exit(-1) 377 | ''' 378 | if __name__ == '__main__': 379 | ''' 380 | Sentence bert based: 381 | T5: sentence-t5-large dim 768 382 | mpnet: all-mpnet-base-v1 dim 768 383 | Roberta: all-roberta-large-v1 dim 1024 384 | SIMCSE based: 385 | princeton-nlp/unsup-simcse-roberta-large dim 1024 386 | princeton-nlp/sup-simcse-bert-large-uncased dim 1024 387 | 388 | list of supported datasets: 389 | ['personachat'.'qnli','mnli','sst2','wmt16','multi_woz','abcd'] 390 | ''' 391 | model_cards ={} 392 | model_cards['sent_t5'] = 'sentence-t5-large' 393 | model_cards['mpnet'] = 'all-mpnet-base-v1' 394 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 395 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 396 | model_cards['simcse_roberta'] = 'princeton-nlp/sup-simcse-roberta-large' 397 | 398 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 399 | parser.add_argument('--model_dir', type=str, default='random_gpt2_medium', help='Dir of your model') 400 | parser.add_argument('--num_epochs', type=int, default=10, help='Training epoches.') 401 | parser.add_argument('--batch_size', type=int, default=32, help='Batch_size #.') 402 | parser.add_argument('--dataset', type=str, default='wmt16', help='Name of dataset: personachat or qnli') 403 | #parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 404 | #parser.add_argument('--data_type', type=str, default='train', help='train/test') 405 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 406 | parser.add_argument('--embed_model', type=str, default='simcse_bert', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 407 | parser.add_argument('--decode', type=str, default='beam', help='Name of decoding methods: beam/sampling') 408 | #parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 409 | args = parser.parse_args() 410 | config = {} 411 | config['model_dir'] = args.model_dir 412 | config['num_epochs'] = args.num_epochs 413 | config['batch_size'] = args.batch_size 414 | config['dataset'] = args.dataset 415 | config['data_type'] = args.data_type 416 | config['embed_model'] = args.embed_model 417 | config['decode'] = args.decode 418 | config['embed_model_path'] = model_cards[config['embed_model']] 419 | config['device'] = torch.device("cuda") 420 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 421 | config['eos_token'] = config['tokenizer'].eos_token 422 | 423 | config['use_opt'] = False 424 | device = config['device'] 425 | #device = torch.device("cpu") 426 | batch_size = config['batch_size'] 427 | 428 | 429 | sent_list = get_sent_list(config) 430 | 431 | ##### for training 432 | if(config['data_type'] == 'train'): 433 | process_data(sent_list,batch_size,device,config) 434 | elif(config['data_type'] == 'test'): 435 | if('simcse' in config['embed_model']): 436 | process_data_test_simcse(sent_list,batch_size,device,config,proj_dir=None,need_proj=False) 437 | else: 438 | process_data_test(sent_list,batch_size,device,config,need_proj=True) 439 | -------------------------------------------------------------------------------- /attacker_t5.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"] = "0" 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from scipy.spatial.distance import cosine 8 | 9 | import json 10 | import numpy as np 11 | import pandas as pd 12 | import argparse 13 | import sys 14 | 15 | from transformers import AutoModel, AutoTokenizer 16 | from transformers import AutoModelForCausalLM,GPT2Config,GPT2LMHeadModel 17 | from transformers import AdamW,get_linear_schedule_with_warmup 18 | from transformers import T5Tokenizer, T5ForConditionalGeneration,T5Config 19 | 20 | from torch.utils.data import DataLoader, Dataset 21 | from attacker_models import SequenceCrossEntropyLoss 22 | from sentence_transformers import SentenceTransformer 23 | from simcse_persona import get_persona_dict 24 | from attacker_evaluation_gpt import eval_on_batch 25 | from datasets import load_dataset 26 | from data_process import get_sent_list 27 | 28 | Folder_path = 'models/' 29 | class linear_projection(nn.Module): 30 | def __init__(self, in_num, out_num=1024): 31 | super(linear_projection, self).__init__() 32 | self.fc1 = nn.Linear(in_num, out_num) 33 | 34 | def forward(self, x, use_final_hidden_only = True): 35 | # x should be of shape (?,in_num) according to gpt2 output 36 | out_shape = x.size()[-1] 37 | assert(x.size()[1] == out_shape) 38 | out = self.fc1(x) 39 | 40 | 41 | return out 42 | 43 | 44 | class personachat(Dataset): 45 | def __init__(self, data): 46 | self.data = data 47 | 48 | 49 | def __len__(self): 50 | return len(self.data) 51 | 52 | def __getitem__(self, index): 53 | 54 | text = self.data[index] 55 | 56 | return text 57 | 58 | def collate(self, unpacked_data): 59 | return unpacked_data 60 | 61 | 62 | def init_opt(): 63 | #opt-350m 64 | T5Config 65 | config = T5Config.from_pretrained('google/t5-large-lm-adapt') 66 | model = T5ForConditionalGeneration(config) 67 | #model = T5ForConditionalGeneration.from_pretrained("google/t5-large-lm-adapt") 68 | tokenizer = T5Tokenizer.from_pretrained("google/t5-large-lm-adapt") 69 | #tokenizer = AutoTokenizer.from_pretrained('facebook/opt-1.3b',use_fast=False) 70 | #model = GPT2LMHeadModel(config) 71 | return model, tokenizer 72 | 73 | def process_data(data,batch_size,device,config,need_proj=True): 74 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 75 | embed_model_name = config['embed_model'] 76 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 77 | dataset = personachat(data) 78 | dataloader = DataLoader(dataset=dataset, 79 | shuffle=True, 80 | batch_size=batch_size, 81 | collate_fn=dataset.collate) 82 | 83 | print('load data done') 84 | 85 | ### extra projection 86 | if need_proj: 87 | projection = linear_projection(in_num=768).to(device) 88 | ### for attackers 89 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 90 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 91 | model_attacker, tokenizer_attacker = init_opt() 92 | criterion = SequenceCrossEntropyLoss() 93 | model_attacker.to(device) 94 | param_optimizer = list(model_attacker.named_parameters()) 95 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 96 | optimizer_grouped_parameters = [ 97 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 98 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 99 | 100 | ] 101 | num_gradients_accumulation = 1 102 | num_epochs = config['num_epochs'] 103 | batch_size = config['batch_size'] 104 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 105 | optimizer = AdamW(optimizer_grouped_parameters, 106 | lr=3e-4, 107 | eps=1e-06) #Typically, 1e-4 and 3e-4 work well for most problems (classification, summarization, translation, question answering, question generation) 108 | if need_proj: 109 | optimizer.add_param_group({'params': projection.parameters()}) 110 | scheduler = get_linear_schedule_with_warmup(optimizer, 111 | num_warmup_steps=100, 112 | num_training_steps = num_train_optimization_steps) 113 | 114 | ### process to obtain the embeddings 115 | for i in range(num_epochs): 116 | for idx,batch_text in enumerate(dataloader): 117 | with torch.no_grad(): 118 | 119 | #sys.exit(-1) 120 | embeddings = model.encode(batch_text,convert_to_tensor = True) 121 | print(f'Embedding dim: {embeddings.size()}') 122 | 123 | ### attacker part, needs training 124 | if need_proj: 125 | embeddings = projection(embeddings) 126 | 127 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 128 | optimizer.step() 129 | scheduler.step() 130 | # make sure no grad for GPT optimizer 131 | optimizer.zero_grad() 132 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 133 | #sys.exit(-1) 134 | if need_proj: 135 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 136 | torch.save(projection.state_dict(), proj_path) 137 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 138 | model_attacker.save_pretrained(save_path) 139 | 140 | 141 | 142 | def process_data_simcse(data,batch_size,device,config,need_proj=False): 143 | embed_model_name = config['embed_model'] 144 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 145 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 146 | dataset = personachat(data) 147 | dataloader = DataLoader(dataset=dataset, 148 | shuffle=True, 149 | batch_size=batch_size, 150 | collate_fn=dataset.collate) 151 | 152 | print('load data done') 153 | 154 | ### extra projection 155 | if need_proj: 156 | projection = linear_projection(in_num=768).to(device) 157 | 158 | ### for attackers 159 | #model_attacker = AutoModelForCausalLM.from_pretrained(config['model_dir']) 160 | #tokenizer_attacker = AutoTokenizer.from_pretrained(config['model_dir']) 161 | model_attacker, tokenizer_attacker = init_opt() 162 | criterion = SequenceCrossEntropyLoss() 163 | model_attacker.to(device) 164 | param_optimizer = list(model_attacker.named_parameters()) 165 | no_decay = ['bias', 'ln', 'LayerNorm.weight'] 166 | optimizer_grouped_parameters = [ 167 | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, 168 | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 169 | 170 | ] 171 | num_gradients_accumulation = 1 172 | num_epochs = config['num_epochs'] 173 | batch_size = config['batch_size'] 174 | num_train_optimization_steps = len(dataloader) * num_epochs // num_gradients_accumulation 175 | optimizer = AdamW(optimizer_grouped_parameters, 176 | lr=3e-4, 177 | eps=1e-06) 178 | if need_proj: 179 | optimizer.add_param_group({'params': projection.parameters()}) 180 | scheduler = get_linear_schedule_with_warmup(optimizer, 181 | num_warmup_steps=100, 182 | num_training_steps = num_train_optimization_steps) 183 | 184 | ### process to obtain the embeddings 185 | for i in range(num_epochs): 186 | for idx,batch_text in enumerate(dataloader): 187 | with torch.no_grad(): 188 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 189 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 190 | print(embeddings.size()) 191 | 192 | 193 | ### attacker part, needs training 194 | if need_proj: 195 | embeddings = projection(embeddings) 196 | 197 | record_loss, perplexity = train_on_batch(batch_X=embeddings,batch_D=batch_text,model=model_attacker,tokenizer=tokenizer_attacker,criterion=criterion,device=device,train=True) 198 | optimizer.step() 199 | scheduler.step() 200 | # make sure no grad for GPT optimizer 201 | optimizer.zero_grad() 202 | print(f'{embed_model_name}: Training: epoch {i} batch {idx} with loss: {record_loss} and PPL {perplexity} with size {embeddings.size()}') 203 | #sys.exit(-1) 204 | if need_proj: 205 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 206 | torch.save(projection.state_dict(), proj_path) 207 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 208 | model_attacker.save_pretrained(save_path) 209 | 210 | ### used for testing only 211 | def process_data_test(data,batch_size,device,config,need_proj=False): 212 | #model = SentenceTransformer('all-roberta-large-v1',device=device) # dim 1024 213 | model = SentenceTransformer(config['embed_model_path'],device=device) # dim 768 214 | if(config['decode'] == 'beam'): 215 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 216 | else: 217 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'.log' 218 | dataset = personachat(data) 219 | # no shuffle for testing data 220 | dataloader = DataLoader(dataset=dataset, 221 | shuffle=False, 222 | batch_size=batch_size, 223 | collate_fn=dataset.collate) 224 | 225 | print('load data done') 226 | if need_proj: 227 | proj_path = Folder_path + 'projection_opt_' + config['dataset'] + '_' + config['embed_model'] 228 | projection = linear_projection(in_num=768) 229 | projection.load_state_dict(torch.load(proj_path)) 230 | projection.to(device) 231 | print('load projection done') 232 | else: 233 | print('no projection loaded') 234 | # setup on config for sentence generation AutoModelForCausalLM 235 | attacker_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 236 | config['model'] = T5ForConditionalGeneration.from_pretrained(attacker_path).to(device) 237 | #config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 238 | 239 | sent_dict = {} 240 | sent_dict['gt'] = [] 241 | sent_dict['pred'] = [] 242 | with torch.no_grad(): 243 | for idx,batch_text in enumerate(dataloader): 244 | embeddings = model.encode(batch_text,convert_to_tensor = True) 245 | #inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 246 | #embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 247 | if need_proj: 248 | embeddings = projection(embeddings) 249 | #embeddings_out = embeddings_proj.detach().cpu() 250 | #data_dict['text'].extend(batch_text) 251 | #data_dict['embedding'].extend(embeddings_out) 252 | sent_list, gt_list = eval_on_batch_t5(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 253 | print(f'testing {idx} batch done with {idx*batch_size} samples') 254 | sent_dict['pred'].extend(sent_list) 255 | sent_dict['gt'].extend(gt_list) 256 | 257 | with open(save_path, 'w') as f: 258 | json.dump(sent_dict, f,indent=4) 259 | 260 | return 0 261 | 262 | 263 | def eval_on_batch_t5(batch_X,batch_D,model,tokenizer,device,config): 264 | device = config['device'] 265 | hidden_X_unsqueeze = torch.unsqueeze(batch_X, 1) 266 | inputs_embeds = hidden_X_unsqueeze# a 3D tensor, [batch, seq_length, dim] : seq_len = 1 267 | print(f'inputs_embeds:{inputs_embeds.shape}') 268 | attention_mask = torch.ones(inputs_embeds.shape[:2], dtype=torch.long).to(device) 269 | decoder_input_ids = torch.ones((inputs_embeds.shape[0], 1), dtype=torch.long) * model.config.decoder_start_token_id 270 | decoder_input_ids = decoder_input_ids.to(device) 271 | output_ids = model.generate(attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, inputs_embeds=inputs_embeds, max_length=50,num_beams=5) 272 | gen_sent_list = tokenizer.batch_decode(output_ids,skip_special_tokens=True) 273 | sent_list = gen_sent_list 274 | gt_list = batch_D 275 | i = 2 276 | print(f'pred:{sent_list[i]}') 277 | print(f'gt :{gt_list[i]}') 278 | return sent_list, gt_list 279 | 280 | 281 | ### used for testing only 282 | def process_data_test_simcse(data,batch_size,device,config,proj_dir=None,need_proj=False): 283 | tokenizer = AutoTokenizer.from_pretrained(config['embed_model_path']) # dim 1024 284 | model = AutoModel.from_pretrained(config['embed_model_path']).to(device) 285 | #save_path = 'logs/attacker_opt_qnli_simcse_bert_large.log' 286 | if(config['decode'] == 'beam'): 287 | print('Using beam search decoding') 288 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'_beam'+'.log' 289 | else: 290 | save_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model']+'.log' 291 | dataset = personachat(data) 292 | # no shuffle for testing data 293 | dataloader = DataLoader(dataset=dataset, 294 | shuffle=False, 295 | batch_size=batch_size, 296 | collate_fn=dataset.collate) 297 | 298 | print('load data done') 299 | if need_proj: 300 | projection = linear_projection(in_num=768) 301 | projection.load_state_dict(torch.load(proj_dir)) 302 | projection.to(device) 303 | print('load projection done') 304 | else: 305 | print('no projection loaded') 306 | # setup on config for sentence generation AutoModelForCausalLM 307 | attacker_path = Folder_path + 'attacker_opt_' + config['dataset'] + '_' + config['embed_model'] 308 | config['model'] = AutoModelForCausalLM.from_pretrained(attacker_path).to(device) 309 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 310 | 311 | sent_dict = {} 312 | sent_dict['gt'] = [] 313 | sent_dict['pred'] = [] 314 | with torch.no_grad(): 315 | for idx,batch_text in enumerate(dataloader): 316 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 317 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 318 | if need_proj: 319 | embeddings = projection(embeddings) 320 | #sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 321 | sent_list, gt_list = eval_on_batch(batch_X=embeddings,batch_D=batch_text,model=config['model'],tokenizer=config['tokenizer'],device=device,config=config) 322 | print(f'testing {idx} batch done with {idx*batch_size} samples') 323 | sent_dict['pred'].extend(sent_list) 324 | sent_dict['gt'].extend(gt_list) 325 | with open(save_path, 'w') as f: 326 | json.dump(sent_dict, f,indent=4) 327 | 328 | return 0 329 | 330 | 331 | def train_on_batch(batch_X,batch_D,model,tokenizer,criterion,device,train=True): 332 | #padding_token_id = tokenizer.encode(tokenizer.eos_token)[0] 333 | tokenizer.pad_token = tokenizer.eos_token 334 | inputs = tokenizer(batch_D, return_tensors='pt', padding='max_length', truncation=True, max_length=40) 335 | #dial_tokens = [tokenizer.encode(item) + turn_ending for item in batch_D] 336 | #print(inputs) 337 | input_ids = inputs['input_ids'] # tensors of input ids 338 | labels = input_ids 339 | # replace padding token id's of the labels by -100 so it's ignored by the loss 340 | #labels[labels == tokenizer.pad_token_id] = -100 341 | input_ids = input_ids.to(device) 342 | labels = labels.to(device) 343 | #attention_mask = inputs['attention_mask'] 344 | 345 | 346 | batch_size = batch_X.size()[0] 347 | attention_mask = torch.ones(batch_size,1).to(device) 348 | 349 | #print(input_ids.size()) 350 | # embed the input ids using GPT-2 embedding 351 | input_emb = model.shared(input_ids) 352 | # add extra dim to cat together 353 | batch_X = batch_X.to(device) 354 | batch_X = torch.unsqueeze(batch_X, 1) 355 | #inputs_embeds = torch.cat((batch_X_unsqueeze,input_emb),dim=1) #[batch,max_length+1,emb_dim (1024)] 356 | assert input_emb.size()[-1] == batch_X.size()[-1] 357 | 358 | 359 | 360 | #output = model(inputs_embeds=batch_X,attention_mask=attention_mask,labels=labels,return_dict=True,output_hidden_states=True) 361 | output = model(inputs_embeds=batch_X,labels=labels,return_dict=True) 362 | # encoder_last_hidden_state = output.encoder_last_hidden_state 363 | # last_hidden_state = output.decoder_hidden_states[-1] 364 | # print(f'encoder_last_hidden_state:{encoder_last_hidden_state.size()}') #([bs, 1, 1024]) 365 | # print(f'last_hidden_state:{last_hidden_state.size()}') #([bs, 40, 1024]) 366 | loss = output.loss 367 | 368 | record_loss = loss.item() 369 | perplexity = np.exp(record_loss) 370 | if train: 371 | loss.backward() 372 | 373 | return record_loss, perplexity 374 | 375 | 376 | 377 | 378 | 379 | if __name__ == '__main__': 380 | ''' 381 | Sentence bert based: 382 | T5: sentence-t5-large dim 768 383 | mpnet: all-mpnet-base-v1 dim 768 384 | Roberta: all-roberta-large-v1 dim 1024 385 | 386 | 387 | SIMCSE based: 388 | princeton-nlp/unsup-simcse-roberta-large dim 1024 389 | princeton-nlp/sup-simcse-bert-large-uncased dim 1024 390 | ''' 391 | model_cards ={} 392 | model_cards['sent_t5'] = 'sentence-t5-large' 393 | model_cards['mpnet'] = 'all-mpnet-base-v1' 394 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 395 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 396 | model_cards['simcse_roberta'] = 'princeton-nlp/sup-simcse-roberta-large' 397 | 398 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 399 | parser.add_argument('--model_dir', type=str, default='t5', help='Dir of your model') 400 | parser.add_argument('--num_epochs', type=int, default=10, help='Training epoches.') 401 | parser.add_argument('--batch_size', type=int, default=32, help='Batch_size #.') 402 | parser.add_argument('--dataset', type=str, default='personachat', help='Name of dataset: personachat or qnli') 403 | #parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 404 | #parser.add_argument('--data_type', type=str, default='train', help='train/test') 405 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 406 | parser.add_argument('--embed_model', type=str, default='sent_t5', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 407 | parser.add_argument('--decode', type=str, default='beam', help='Name of decoding methods: beam/sampling') 408 | #parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 409 | args = parser.parse_args() 410 | config = {} 411 | config['model_dir'] = args.model_dir 412 | config['num_epochs'] = args.num_epochs 413 | config['batch_size'] = args.batch_size 414 | config['dataset'] = args.dataset 415 | config['data_type'] = args.data_type 416 | config['embed_model'] = args.embed_model 417 | config['decode'] = args.decode 418 | config['embed_model_path'] = model_cards[config['embed_model']] 419 | config['device'] = torch.device("cuda") 420 | config['tokenizer'] = T5Tokenizer.from_pretrained("google/t5-large-lm-adapt") 421 | config['eos_token'] = config['tokenizer'].eos_token 422 | 423 | 424 | device = torch.device("cuda") 425 | #device = torch.device("cpu") 426 | batch_size = config['batch_size'] 427 | 428 | 429 | sent_list = get_sent_list(config) 430 | 431 | ##### for training 432 | if(config['data_type'] == 'train'): 433 | process_data(sent_list,batch_size,device,config) 434 | elif(config['data_type'] == 'test'): 435 | if('simcse' in config['embed_model'] or 'sent_roberta' in config['embed_model']): 436 | process_data_test_simcse(sent_list,batch_size,device,config,proj_dir=None,need_proj=False) 437 | else: 438 | process_data_test(sent_list,batch_size,device,config,need_proj=True) 439 | -------------------------------------------------------------------------------- /baseline/projection.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | os.environ["CUDA_VISIBLE_DEVICES"] = "3" 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from scipy.spatial.distance import cosine 9 | 10 | import numpy as np 11 | import pandas as pd 12 | import argparse 13 | import sys 14 | 15 | sys.path.append('..') 16 | from pprint import pprint 17 | 18 | from transformers import AutoModel, AutoTokenizer 19 | from transformers import AutoModelForCausalLM 20 | from transformers import AdamW, get_linear_schedule_with_warmup 21 | from torch.utils.data import DataLoader, Dataset 22 | from attacker_models import SequenceCrossEntropyLoss 23 | from sentence_transformers import SentenceTransformer 24 | from simcse_persona import get_persona_dict 25 | from attacker_evaluation_gpt import eval_on_batch 26 | from datasets import load_dataset 27 | import baseline_models 28 | 29 | from sklearn import metrics 30 | import logging 31 | import logging.handlers 32 | import json 33 | from data_process import get_sent_list 34 | 35 | torch.autograd.set_detect_anomaly(True) 36 | 37 | logger = logging.getLogger('mylogger') 38 | logger.setLevel(logging.DEBUG) 39 | f_handler = logging.FileHandler('logs/new_datasets.log') 40 | f_handler.setLevel(logging.INFO) 41 | f_handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) 42 | logger.addHandler(f_handler) 43 | 44 | 45 | 46 | def save_blmodel(model, path): 47 | torch.save(model.state_dict(), path) 48 | 49 | 50 | def load_blmodel(model, path): 51 | model.load_state_dict(torch.load(path)) 52 | 53 | ''' 54 | def get_sent_list(config): 55 | dataset = config['dataset'] 56 | data_type = config['data_type'] 57 | if dataset == 'personachat': 58 | sent_list = get_personachat_data(data_type) 59 | return sent_list 60 | elif dataset == 'qnli': 61 | sent_list = get_qnli_data(data_type) 62 | return sent_list 63 | else: 64 | print('Name of dataset only supports: personachat or qnli') 65 | sys.exit(-1) 66 | 67 | 68 | def get_personachat_data(data_type): 69 | sent_list = get_persona_dict(data_type=data_type) 70 | return sent_list 71 | 72 | 73 | def get_qnli_data(data_type): 74 | dataset = load_dataset('glue', 'qnli', cache_dir="/home/hlibt/embed_rev/data", split=data_type) 75 | sentence_list = [] 76 | for i, d in enumerate(dataset): 77 | sentence_list.append(d['question']) 78 | sentence_list.append(d['sentence']) 79 | return sentence_list 80 | 81 | ''' 82 | class sent_list_dataset(Dataset): 83 | def __init__(self, sent_list, label): 84 | assert len(sent_list) == len(label) 85 | self.sent_list = sent_list 86 | self.label = label 87 | 88 | def __len__(self): 89 | return len(self.sent_list) 90 | 91 | def __getitem__(self, index): 92 | sentence = self.sent_list[index] 93 | label = self.label[index] 94 | 95 | return sentence, label 96 | 97 | def collate(self, unpacked_data): 98 | return unpacked_data 99 | 100 | class collated_dataset(Dataset): 101 | def __init__(self, sent_list, config): 102 | 103 | self.sent_list = sent_list 104 | self.config = config 105 | 106 | 107 | def __len__(self): 108 | return len(self.sent_list) 109 | 110 | def __getitem__(self, index): 111 | sentence = self.sent_list[index] 112 | 113 | return sentence 114 | 115 | def collate(self, unpacked_data): 116 | #pprint(unpacked_data) 117 | dial_tokens_np, input_labels = process_sent_list(unpacked_data, self.config) 118 | #print(dial_tokens_np.shape) 119 | #print(input_labels.shape) 120 | return unpacked_data, input_labels 121 | 122 | def process_sent_list(sent_list, config): 123 | tokenizer = config['tokenizer'] 124 | turn_ending = config['eos_token'] 125 | print(tokenizer) 126 | # dial_tokens = [tokenizer.encode(item) + turn_ending for item in conv] 127 | dial_tokens = [tokenizer.encode(item,max_length=150,padding='max_length', truncation=True) for item in sent_list] 128 | # for input loss training 129 | token_num = len(tokenizer) 130 | print('lenth of tokenizer ', end='') 131 | print(token_num) 132 | dial_tokens_np = np.array(dial_tokens) 133 | input_labels = [] 134 | for i in dial_tokens_np: 135 | temp_i = np.zeros(token_num) 136 | temp_i[i] = 1 137 | input_labels.append(temp_i) 138 | input_labels = np.array(input_labels) 139 | 140 | return dial_tokens_np, input_labels 141 | 142 | 143 | ''' 144 | choosing baseline attacker models here 145 | ''' 146 | 147 | 148 | def init_baseline_model(config, embed_model_dim, type='NN'): 149 | tokenizer = config['tokenizer'] 150 | token_num = len(tokenizer) 151 | device = config['device'] 152 | if type == 'NN': 153 | # print(f'line 112 embed_model_dim:{embed_model_dim}') 154 | baseline_model = baseline_models.baseline_NN(out_num=token_num, in_num=embed_model_dim) 155 | elif type == 'RNN': 156 | baseline_model = baseline_models.baseline_RNN(out_num=token_num, in_num=embed_model_dim) 157 | 158 | 159 | optimizer = AdamW(baseline_model.parameters(), 160 | lr=6e-5, 161 | eps=1e-06) 162 | criterion = nn.BCEWithLogitsLoss() # usage criterion(external_out,input_label) 163 | baseline_model.to(device) 164 | return baseline_model, optimizer, criterion 165 | 166 | 167 | ''' 168 | Sentence bert based: 169 | T5: sentence-t5-large dim 768 170 | mpnet: all-mpnet-base-v1 dim 768 171 | Roberta: all-roberta-large-v1 dim 1024 172 | 173 | 174 | SIMCSE based: 175 | princeton-nlp/unsup-simcse-roberta-large dim 1024 176 | princeton-nlp/sup-simcse-bert-large-uncased dim 1024 177 | ''' 178 | 179 | 180 | def get_embedding(dataloader, config, eval=False): 181 | model_cards = {} 182 | model_cards['sent_t5'] = 'sentence-t5-large' 183 | model_cards['mpnet'] = 'all-mpnet-base-v1' 184 | model_cards['sent_roberta'] = 'all-roberta-large-v1' 185 | model_cards['simcse_bert'] = 'princeton-nlp/sup-simcse-bert-large-uncased' 186 | model_cards['simcse_roberta'] = 'princeton-nlp/unsup-simcse-roberta-large' 187 | embed_model = config['embed_model'] 188 | embed_model_dim = 1024 189 | assert embed_model in ['mpnet', 'sent_roberta', 'simcse_bert', 'simcse_roberta', 'sent_t5'] 190 | model_name = model_cards[embed_model] 191 | if embed_model in ['mpnet', 'sent_roberta', 'sent_t5']: 192 | if embed_model in ['mpnet', 'sent_t5']: 193 | embed_model_dim = 768 194 | # print(f'line 146 embed_model_dim:{embed_model_dim}') 195 | if eval: 196 | eval_sent(dataloader, model_name, embed_model_dim, config) 197 | else: 198 | train_sent(dataloader, model_name, embed_model_dim, config) 199 | elif embed_model in ['simcse_bert', 'simcse_roberta']: 200 | if (eval): 201 | eval_simcse(dataloader, model_name, embed_model_dim, config) 202 | else: 203 | train_simcse(dataloader, model_name, embed_model_dim, config) 204 | 205 | 206 | def train_sent(dataloader, model_name, embed_model_dim, config): 207 | device = config['device'] 208 | num_epochs = config['num_epochs'] 209 | model = SentenceTransformer(model_name, device=device) 210 | type = config['model_type'] 211 | baseline_model, optimizer, criterion = init_baseline_model(config, embed_model_dim, type=type) 212 | save_path = 'blmodels/' + type + '_' + config['dataset'] + '_' + config['embed_model'] 213 | 214 | hx = torch.zeros(64, 512).cuda() 215 | for i in range(num_epochs): 216 | 217 | for idx, batch in enumerate(dataloader): 218 | # print(batch) 219 | batch_text, batch_label = batch 220 | batch_label = batch_label.to(device) 221 | batch_text = list(batch_text) 222 | print(f'{i}: Batch id: {idx}. batch_text: {batch_text[0]}. batch_label: {batch_label.size()}.') 223 | ### no grad for embedding model 224 | with torch.no_grad(): 225 | embeddings = model.encode(batch_text, convert_to_tensor=True) 226 | print(f'embeddings:{embeddings.size()}') 227 | # embeddings:torch.Size([64, 768, 1]) 228 | print(f'batch_label:{batch_label.size()}') 229 | # batch_label:torch.Size([64, 50257, 1]) 230 | print(save_path) 231 | # exit() 232 | 233 | # here is the embedding sentences, we need to turn it into tensor for future use 234 | if type == 'NN': 235 | train_on_batch(baseline_model, optimizer, criterion, embeddings, batch_label) 236 | elif type == 'RNN': 237 | total_loss = baseline_model(embeddings, hx, batch_label) 238 | optimizer.zero_grad() 239 | total_loss.backward() 240 | optimizer.step() 241 | print(f'========total_loss:{total_loss}') 242 | save_blmodel(baseline_model, save_path) 243 | 244 | 245 | def train_simcse(dataloader, model_name, embed_model_dim, config): 246 | device = config['device'] 247 | num_epochs = config['num_epochs'] 248 | type = config['model_type'] 249 | baseline_model, optimizer, criterion = init_baseline_model(config, embed_model_dim, type=type) 250 | save_path = 'blmodels/' + type + '_' + config['dataset'] + '_' + config['embed_model'] 251 | print(f'save_path:{save_path}') 252 | tokenizer = AutoTokenizer.from_pretrained(model_name) 253 | model = AutoModel.from_pretrained(model_name).to(device) 254 | hx = torch.zeros(64, 512).cuda() 255 | for i in range(num_epochs): 256 | for idx, batch in enumerate(dataloader): 257 | batch_text, batch_label = batch 258 | batch_label = torch.tensor(batch_label) 259 | batch_label = batch_label.to(device) 260 | batch_text = list(batch_text) 261 | print(f'{i}: Batch id: {idx}. batch_text: {batch_text[0]}. batch_label: {batch_label.size()}.') 262 | with torch.no_grad(): 263 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 264 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 265 | 266 | print(f'embeddings:{embeddings.size()}') 267 | print(save_path) 268 | if type == 'NN': 269 | train_on_batch(baseline_model, optimizer, criterion, embeddings, batch_label) 270 | elif type == 'RNN': 271 | total_loss = baseline_model(embeddings, hx, batch_label) 272 | optimizer.zero_grad() 273 | total_loss.backward() 274 | optimizer.step() 275 | print(f'========total_loss:{total_loss}') 276 | save_blmodel(baseline_model, save_path) 277 | 278 | 279 | def train_on_batch(baseline_model, optimizer, criterion, embedding_batch, label_batch): 280 | logits = baseline_model(embedding_batch) 281 | loss = criterion(logits, label_batch) 282 | optimizer.zero_grad() 283 | loss.backward() 284 | optimizer.step() 285 | print(f'========loss:{loss}') 286 | 287 | 288 | def report_score(y_true, y_pred): 289 | # micro result should be reported 290 | precision = metrics.precision_score(y_true, y_pred, average='micro') 291 | recall = metrics.recall_score(y_true, y_pred, average='micro') 292 | f1 = metrics.f1_score(y_true, y_pred, average='micro') 293 | # logger.info(f"micro precision_score on token level: {precision}") 294 | print(f"micro precision_score on token level: {precision}") 295 | # print("macro precision_score: {:.2f} ".format( metrics.precision_score(y_true, y_pred, average='macro'))) 296 | print('==' * 20) 297 | # logger.info(f"micro recall_score on token level: {recall}") 298 | print(f"micro recall_score on token level: {recall}") 299 | # print("macro recall_score: {:.2f} ".format( metrics.recall_score(y_true, y_pred, average='macro'))) 300 | print('==' * 20) 301 | # logger.info(f"micro f1_score on token level: {f1}") 302 | print(f"micro f1_score on token level: {f1}") 303 | # print("macro f1_score: {:.2f} ".format( metrics.f1_score(y_true, y_pred, average='macro'))) 304 | 305 | 306 | def eval_sent(dataloader, model_name, embed_model_dim, config): 307 | device = config['device'] 308 | num_epochs = config['num_epochs'] 309 | model = SentenceTransformer(model_name, device=device) 310 | type = config['model_type'] 311 | hx = torch.zeros(64, 512).cuda() 312 | 313 | baseline_model, optimizer, criterion = init_baseline_model(config, embed_model_dim, type=type) 314 | model_path = 'blmodels/' + type + '_' + config['dataset'] + '_' + config['embed_model'] 315 | baseline_model.load_state_dict(torch.load(model_path)) 316 | baseline_model.eval() 317 | 318 | print(f'load baseline {type} from {model_path}') 319 | # load_blmodel(baseline_model,model_path) 320 | log_text = model_path 321 | logger.info('=======New baseline model evaluation=========') 322 | logger.info(log_text) 323 | predict = [] 324 | ground_truth = [] 325 | input_text = [] 326 | for idx, batch in enumerate(dataloader): 327 | # print(batch) 328 | batch_text, batch_label = batch 329 | batch_label = batch_label.to(device) 330 | batch_text = list(batch_text) 331 | print(f'Batch id: {idx}. batch_text: {batch_text[0]}. batch_label: {batch_label.size()}.') 332 | ### no grad for embedding model 333 | with torch.no_grad(): 334 | embeddings = model.encode(batch_text, convert_to_tensor=True) 335 | print(f'embeddings:{embeddings.size()}') 336 | # embeddings:torch.Size([64, 1024]) [batch size, feature(dim)] 337 | if type == 'NN': 338 | predict_result = eval_on_batch(baseline_model, criterion, embeddings, batch_label) 339 | elif type == 'RNN': 340 | zero_tensor = torch.zeros(64, 50257).cuda() 341 | # one_tensor = torch.ones(64, 50257).cuda() 342 | predicted_token_index_batch_tensor = baseline_model(embeddings, hx, batch_label, eval=True) 343 | predict_result = zero_tensor.scatter_(1, predicted_token_index_batch_tensor, value = 1) 344 | predict.extend(predict_result.cpu().detach().numpy()) 345 | ground_truth.extend(batch_label.cpu().detach().numpy()) 346 | input_text.extend(batch_text) 347 | eval_label(predict, ground_truth, input_text, config, type=type) 348 | report_score(ground_truth, predict) 349 | 350 | 351 | def eval_simcse(dataloader, model_name, embed_model_dim, config): 352 | hx = torch.zeros(64, 512).cuda() 353 | device = config['device'] 354 | num_epochs = config['num_epochs'] 355 | type = config['model_type'] 356 | baseline_model, optimizer, criterion = init_baseline_model(config, embed_model_dim, type=type) 357 | model_path = 'blmodels/' + type + '_' + config['dataset'] + '_' + config['embed_model'] 358 | baseline_model.load_state_dict(torch.load(model_path)) 359 | baseline_model.eval() 360 | log_text = model_path 361 | logger.info('=======New baseline model evaluation=========') 362 | logger.info(log_text) 363 | print(f'load baseline {type} from {model_path}') 364 | # load_blmodel(baseline_model,model_path) 365 | 366 | tokenizer = AutoTokenizer.from_pretrained(model_name) 367 | model = AutoModel.from_pretrained(model_name).to(device) 368 | predict = [] 369 | ground_truth = [] 370 | input_text = [] 371 | for idx, batch in enumerate(dataloader): 372 | batch_text, batch_label = batch 373 | batch_label = torch.tensor(batch_label) 374 | batch_label = batch_label.to(device) 375 | batch_text = list(batch_text) 376 | print(f'Batch id: {idx}. batch_text: {batch_text[0]}. batch_label: {batch_label.size()}.') 377 | with torch.no_grad(): 378 | inputs = tokenizer(batch_text, padding=True, truncation=True, return_tensors="pt").to(device) 379 | # print(f'inputs:{inputs.size()}') 380 | embeddings = model(**inputs, output_hidden_states=True, return_dict=True).pooler_output 381 | print(f'embeddings:{embeddings.size()}') 382 | 383 | # batch_pred_labels = eval_on_batch(baseline_model,criterion,embeddings,batch_label) 384 | if type == 'NN': 385 | predict_result = eval_on_batch(baseline_model, criterion, embeddings, batch_label) 386 | elif type == 'RNN': 387 | zero_tensor = torch.zeros(64, 50257).cuda() 388 | # one_tensor = torch.ones(64, 50257).cuda() 389 | predicted_token_index_batch_tensor = baseline_model(embeddings, hx, batch_label, eval=True) 390 | predict_result = zero_tensor.scatter_(1, predicted_token_index_batch_tensor, value=1) 391 | predict.extend(predict_result.cpu().detach().numpy()) 392 | ground_truth.extend(batch_label.cpu().detach().numpy()) 393 | input_text.extend(batch_text) 394 | eval_label(predict, ground_truth, input_text, config, type=type) 395 | report_score(ground_truth, predict) 396 | 397 | 398 | def eval_on_batch(baseline_model, criterion, embedding_batch, label_batch): 399 | logits = baseline_model(embedding_batch, eval=True) 400 | loss = criterion(logits, label_batch) 401 | print(logits) 402 | print(type(logits)) 403 | logits[logits >= 0.5] = 1 404 | logits[logits < 0.5] = 0 405 | print(f'========eval loss:{loss}') 406 | return logits 407 | 408 | 409 | ''' 410 | Used for case study to see what are outputted by baselines 411 | pred labels are (N,dim) 0,1 vectors, where 1 indicate token i show up 412 | ''' 413 | 414 | 415 | def eval_label(pred_labels,ground_truth,input, config,type = 'NN'): 416 | if type == 'NN': 417 | #pred = torch.nonzero(pred_labels) # 2 dim (N,vocab_size) 418 | threshold = config['threshold'] 419 | tokenizer = config['tokenizer'] 420 | pred = [[] for i in range(len(pred_labels))] 421 | gt = [[] for i in range(len(pred_labels))] 422 | ####log_text = model_path+'_threshold_'+str(threshold) 423 | str_threshold = f'{threshold:.2f}' 424 | save_path = 'qnli/' + type + '_' + config['dataset'] + '_' + config['embed_model'] +'_threshold_'+str(str_threshold)+'.label' 425 | #save_path = 'logs_test/' +'test_' +type + '_' + config['dataset'] + '_' + config['embed_model'] +'_threshold_'+str(str_threshold)+'.label' 426 | for idx,label_list in enumerate(pred_labels): 427 | for i,value in enumerate(label_list): 428 | if(value > 0): 429 | assert value == 1 430 | pred[idx].append(tokenizer.decode(i)) # append a string to a list 431 | 432 | for idx,label_list in enumerate(ground_truth): 433 | for i,value in enumerate(label_list): 434 | if(value > 0): 435 | assert value == 1 436 | gt[idx].append(tokenizer.decode(i)) # append a string to a list 437 | 438 | save_data = [] 439 | for i in range(len(pred_labels)): 440 | save_dict = {} 441 | save_dict['pred'] = pred[i] 442 | save_dict['gt'] = gt[i] 443 | save_dict['input'] = input[i] 444 | save_data.append(save_dict) 445 | with open(save_path, 'w') as f: 446 | json.dump(save_data, f,indent=4) 447 | else: 448 | tokenizer = config['tokenizer'] 449 | pred = [[] for i in range(len(pred_labels))] 450 | gt = [[] for i in range(len(pred_labels))] 451 | save_path = 'logs_test/' +'test_' + type + '_' + config['dataset'] + '_' + config['embed_model'] + '.label' 452 | # save_path = 'logs_test/' +'test_' +type + '_' + config['dataset'] + '_' + config['embed_model'] +'_threshold_'+str(str_threshold)+'.label 453 | 454 | for idx, label_list in enumerate(pred_labels): 455 | for i, value in enumerate(label_list): 456 | if (value > 0): 457 | assert value == 1 458 | pred[idx].append(tokenizer.decode(i)) # append a string to a list 459 | 460 | for idx, label_list in enumerate(ground_truth): 461 | for i, value in enumerate(label_list): 462 | if (value > 0): 463 | assert value == 1 464 | gt[idx].append(tokenizer.decode(i)) # append a string to a list 465 | 466 | save_data = [] 467 | for i in range(len(pred_labels)): 468 | save_dict = {} 469 | save_dict['pred'] = pred[i] 470 | save_dict['gt'] = gt[i] 471 | save_dict['input'] = input[i] 472 | save_data.append(save_dict) 473 | with open(save_path, 'w') as f: 474 | json.dump(save_data, f, indent=4) 475 | 476 | 477 | if __name__ == '__main__': 478 | ''' 479 | Sentence bert based: 480 | T5: sentence-t5-large dim 768 481 | mpnet: all-mpnet-base-v1 dim 768 482 | Roberta: all-roberta-large-v1 dim 1024 483 | SIMCSE based: 484 | princeton-nlp/unsup-simcse-roberta-large dim 1024 485 | princeton-nlp/sup-simcse-bert-large-uncased dim 1024 486 | 487 | list of supported datasets: 488 | ['personachat'.'qnli','mnli','sst2','wmt16','multi_woz','abcd'] 489 | ''' 490 | parser = argparse.ArgumentParser(description='Training external NN as baselines') 491 | parser.add_argument('--model_dir', type=str, default='microsoft/DialoGPT-medium', help='Dir of your model') 492 | parser.add_argument('--num_epochs', type=int, default=10, help='Training epoches.') 493 | parser.add_argument('--batch_size', type=int, default=64, help='Batch_size #.') 494 | parser.add_argument('--dataset', type=str, default='abcd', help="List of supported datasets: ['personachat'.'qnli','mnli','sst2','wmt16','multi_woz','abcd']") 495 | # parser.add_argument('--dataset', type=str, default='qnli', help='Name of dataset: personachat or qnli') 496 | parser.add_argument('--data_type', type=str, default='test', help='train/test') 497 | #parser.add_argument('--data_type', type=str, default='test', help='train/test') 498 | parser.add_argument('--embed_model', type=str, default='simcse_bert', 499 | help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 500 | # parser.add_argument('--embed_model', type=str, default='simcse_roberta', help='Name of embedding model: mpnet/sent_roberta/simcse_bert/simcse_roberta/sent_t5') 501 | parser.add_argument('--model_type', type=str, default='NN', help='Type of baseline model: RNN or NN') 502 | # parser.add_argument('--eval', type=str, default=False, help='True or False') 503 | 504 | args = parser.parse_args() 505 | config = {} 506 | config['model_dir'] = args.model_dir 507 | config['num_epochs'] = args.num_epochs 508 | config['batch_size'] = args.batch_size 509 | config['dataset'] = args.dataset 510 | config['data_type'] = args.data_type 511 | config['embed_model'] = args.embed_model 512 | config['model_type'] = args.model_type 513 | 514 | config['device'] = torch.device("cuda") 515 | config['tokenizer'] = AutoTokenizer.from_pretrained('microsoft/DialoGPT-medium') 516 | config['eos_token'] = config['tokenizer'].eos_token 517 | config['tokenizer'].pad_token = config['eos_token'] 518 | sent_list = get_sent_list(config) 519 | 520 | #dataset = sent_list_dataset(sent_list, onehot_labels) 521 | dataset = collated_dataset(sent_list,config) 522 | dataloader = DataLoader(dataset=dataset, 523 | shuffle=False, 524 | batch_size=config['batch_size'], 525 | collate_fn=dataset.collate, 526 | drop_last=True) 527 | # for training 528 | get_embedding(dataloader, config, eval=False) 529 | # for evaluation 530 | get_embedding(dataloader, config, eval=True) 531 | --------------------------------------------------------------------------------