Views
No views yet
1
2import torch
3
4from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
5
6import os
7
8import json
9
10from peft import PeftModel
11
12# from trl import AutoModelForCausalLMWithValueHead
13
14from transformers import AutoModelForCausalLM as AutoGPTQForCausalLM
15
16def load_tokenizer(dir_or_model):
17
18 """
19
20 This function is used to load the tokenizer for a specific pre-trained model.
21
22
23
24 Args:
25
26 dir_or_model: It can be either a directory containing the pre-training model configuration details or a pretrained model.
27
28
29
30 Returns:
31
32 It returns a tokenizer that can convert text to tokens for the specific model input.
33
34 """
35
36 is_lora_dir = os.path.isfile(os.path.join(dir_or_model, "adapter_config.json"))
37
38 if is_lora_dir:
39
40 loaded_json = json.load(open(os.path.join(dir_or_model, "adapter_config.json"), "r"))
41
42 model_name = loaded_json["base_model_name_or_path"]
43
44 else:
45
46 model_name = dir_or_model
47
48
49
50 if os.path.isfile(os.path.join(dir_or_model, "config.json")):
51
52 loaded_json = json.load(open(os.path.join(dir_or_model, "config.json"), "r"))
53
54 if "_name_or_path" in loaded_json:
55
56 model_name = loaded_json["_name_or_path"]
57
58 local_model_name = "/data3/MODELS/llama2-hf/llama-2-7b"#/data2/tsq/WaterBench/data/models/llama-2-7b-chat-hf
59
60
61
62 print(">>>>>>>>>>>>>>>>>>>>>>>>>>notice this<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
63
64
65
66 #print(model_name)
67
68 tokenizer = AutoTokenizer.from_pretrained(local_model_name)
69
70 if tokenizer.pad_token is None:
71
72 tokenizer.pad_token = tokenizer.eos_token
73
74 tokenizer.pad_token_id = tokenizer.eos_token_id
75
76
77
78 return tokenizer
79
80def load_model(dir_or_model, classification=False, token_classification=False, return_tokenizer=False, dtype=torch.bfloat16, load_dtype=True,
81
82 rl=False, peft_config=None, device_map="auto", revision='main'):
83
84 """
85
86 This function is used to load a model based on several parameters including the type of task it is targeted to perform.
87
88
89
90 Args:
91
92 dir_or_model: It can be either a directory containing the pre-training model configuration details or a pretrained model.
93
94 classification (bool): If True, loads the model for sequence classification.
95
96 token_classification (bool): If True, loads the model for token classification.
97
98 return_tokenizer (bool): If True, returns the tokenizer along with the model.
99
100 dtype: The data type that PyTorch should use internally to store the model’s parameters and do the computation.
101
102 load_dtype (bool): If False, sets dtype as torch.float32 regardless of the passed dtype value.
103
104 rl (bool): If True, loads model specifically designed to be used in reinforcement learning environment.
105
106 peft_config: Configuration details for Peft models.
107
108
109
110 Returns:
111
112 It returns a model for the required task along with its tokenizer, if specified.
113
114 """
115
116 is_lora_dir = os.path.isfile(os.path.join(dir_or_model, "adapter_config.json"))
117
118 if not load_dtype:
119
120 dtype = torch.float32
121
122 if is_lora_dir:
123
124 loaded_json = json.load(open(os.path.join(dir_or_model, "adapter_config.json"), "r"))
125
126 model_name = loaded_json["base_model_name_or_path"]
127
128 else:
129
130 model_name = dir_or_model
131
132 original_model_name = model_name
133
134 if classification:
135
136 model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32, use_auth_token=True, device_map=device_map, revision=revision) # to investigate: calling torch_dtype here fails.
137
138 elif token_classification:
139
140 model = AutoModelForTokenClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32, use_auth_token=True, device_map=device_map, revision=revision)
141
142 else:
143
144 if model_name.endswith("GPTQ") or model_name.endswith("GGML"):
145
146 model = AutoGPTQForCausalLM.from_quantized(model_name,
147
148 use_safetensors=True,
149
150 trust_remote_code=True,
151
152 \# use_triton=True, # breaks currently, unfortunately generation time of the GPTQ model is quite slow
153
154 quantize_config=None, device_map=device_map)
155
156 else:
157
158 print('11111111111111111111111111111111111111')
159
160 model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32, use_auth_token=True, device_map=device_map, revision=revision)
161
162 if is_lora_dir:
163
164 model = PeftModel.from_pretrained(model, dir_or_model)
165
166
167
168 try:
169
170 tokenizer = load_tokenizer(original_model_name)
171
172 model.config.pad_token_id = tokenizer.pad_token_id
173
174 except Exception:
175
176 pass
177
178 if return_tokenizer:
179
180 return model, load_tokenizer(original_model_name)
181
182 return model
183
184model_name = 'tsq2000/Jailbreak-generator'
185
186model = load_model(model_name)
187
188tokenizer = load_tokenizer(model_name)
1891
2model_name = 'tsq2000/Jailbreak-generator'
3
4model = load_model(model_name)
5
6tokenizer = load_tokenizer(model_name)
7
8max_length = 2048
9
10max_tokens = 64
11
12knowledge_points = ["Kettling Kettling (also known as containment or corralling) is a police tactic for controlling large crowds during demonstrations or protests. It involves the formation of large cordons of police officers who then move to contain a crowd within a limited area. Protesters are left only one choice of exit controlled by the police – or are completely prevented from leaving, with the effect of denying the protesters access to food, water and toilet facilities for a time period determined by the police forces. The tactic has proved controversial, in part because it has resulted in the detention of ordinary bystanders."]
13
14batch_texts = [f'### Input:\n{input_}\n\n### Response:\n' for input_ in knowledge_points]
15
16inputs = tokenizer(batch_texts, return_tensors='pt', padding=True, truncation=True, max_length=max_length - max_tokens).to(model.device)
17
18outputs = model.generate(**inputs, max_new_tokens=max_tokens, num_return_sequences=1, do_sample=False, temperature=1, top_p=1, eos_token_id=tokenizer.eos_token_id)
19
20generated_texts = []
21
22for output, input_text in zip(outputs, batch_texts):
23
24 text = tokenizer.decode(output, skip_special_tokens=True)
25
26 generated_texts.append(text[len(input_text):])
27
28print(generated_texts)
29@misc{tu2024knowledgetojailbreak,
title={Knowledge-to-Jailbreak: One Knowledge Point Worth One Attack},
author={Shangqing Tu and Zhuoran Pan and Wenxuan Wang and Zhexin Zhang and Yuliang Sun and Jifan Yu and Hongning Wang and Lei Hou and Juanzi Li},
year={2024},
eprint={2406.11682},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2406.11682},
}