license: mit
language:
- en
tags: - llm
- safety
- jailbreak
- knowledge
Introduction
This is a model for generating a jailbreak prompt based on knowledge point texts. The model is trained on the Llama-2-7b dataset and fine-tuned on the Knowledge-to-Jailbreak dataset. The model is intended to bridge the gap between theoretical vulnerabilities and real-world application scenarios, simulating sophisticated adversarial attacks that incorporate specialized knowledge.
Our proposed method and dataset serve as a critical starting point for both offensive and defensive research, enabling the development of new techniques to enhance the security and robustness of language models in practical settings.
How to load the model and tokenizer
We provide two helper functions for loading the model and tokenizer.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
import os
import json
from peft import PeftModel
# from trl import AutoModelForCausalLMWithValueHead
from transformers import AutoModelForCausalLM as AutoGPTQForCausalLM
def load_tokenizer(dir_or_model):
"""
This function is used to load the tokenizer for a specific pre-trained model.
Args:
dir_or_model: It can be either a directory containing the pre-training model configuration details or a pretrained model.
Returns:
It returns a tokenizer that can convert text to tokens for the specific model input.
"""
is_lora_dir = os.path.isfile(os.path.join(dir_or_model, "adapter_config.json"))
if is_lora_dir:
loaded_json = json.load(open(os.path.join(dir_or_model, "adapter_config.json"), "r"))
model_name = loaded_json["base_model_name_or_path"]
else:
model_name = dir_or_model
if os.path.isfile(os.path.join(dir_or_model, "config.json")):
loaded_json = json.load(open(os.path.join(dir_or_model, "config.json"), "r"))
if "_name_or_path" in loaded_json:
model_name = loaded_json["_name_or_path"]
local_model_name = "/data3/MODELS/llama2-hf/llama-2-7b"#/data2/tsq/WaterBench/data/models/llama-2-7b-chat-hf
print(">>>>>>>>>>>>>>>>>>>>>>>>>>notice this<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
#print(model_name)
tokenizer = AutoTokenizer.from_pretrained(local_model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
return tokenizer
def load_model(dir_or_model, classification=False, token_classification=False, return_tokenizer=False, dtype=torch.bfloat16, load_dtype=True,
rl=False, peft_config=None, device_map="auto", revision='main'):
"""
This function is used to load a model based on several parameters including the type of task it is targeted to perform.
Args:
dir_or_model: It can be either a directory containing the pre-training model configuration details or a pretrained model.
classification (bool): If True, loads the model for sequence classification.
token_classification (bool): If True, loads the model for token classification.
return_tokenizer (bool): If True, returns the tokenizer along with the model.
dtype: The data type that PyTorch should use internally to store the model’s parameters and do the computation.
load_dtype (bool): If False, sets dtype as torch.float32 regardless of the passed dtype value.
rl (bool): If True, loads model specifically designed to be used in reinforcement learning environment.
peft_config: Configuration details for Peft models.
Returns:
It returns a model for the required task along with its tokenizer, if specified.
"""
is_lora_dir = os.path.isfile(os.path.join(dir_or_model, "adapter_config.json"))
if not load_dtype:
dtype = torch.float32
if is_lora_dir:
loaded_json = json.load(open(os.path.join(dir_or_model, "adapter_config.json"), "r"))
model_name = loaded_json["base_model_name_or_path"]
else:
model_name = dir_or_model
original_model_name = model_name
if classification:
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.
elif token_classification:
model = AutoModelForTokenClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32, use_auth_token=True, device_map=device_map, revision=revision)
else:
if model_name.endswith("GPTQ") or model_name.endswith("GGML"):
model = AutoGPTQForCausalLM.from_quantized(model_name,
use_safetensors=True,
trust_remote_code=True,
\# use_triton=True, # breaks currently, unfortunately generation time of the GPTQ model is quite slow
quantize_config=None, device_map=device_map)
else:
print('11111111111111111111111111111111111111')
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32, use_auth_token=True, device_map=device_map, revision=revision)
if is_lora_dir:
model = PeftModel.from_pretrained(model, dir_or_model)
try:
tokenizer = load_tokenizer(original_model_name)
model.config.pad_token_id = tokenizer.pad_token_id
except Exception:
pass
if return_tokenizer:
return model, load_tokenizer(original_model_name)
return model
model_name = 'tsq2000/Jailbreak-generator'
model = load_model(model_name)
tokenizer = load_tokenizer(model_name)
How to generate jailbreak prompts
Here is an example of how to generate jailbreak prompts based on knowledge point texts.
model_name = 'tsq2000/Jailbreak-generator'
model = load_model(model_name)
tokenizer = load_tokenizer(model_name)
max_length = 2048
max_tokens = 64
knowledge_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."]
batch_texts = [f'### Input:\n{input_}\n\n### Response:\n' for input_ in knowledge_points]
inputs = tokenizer(batch_texts, return_tensors='pt', padding=True, truncation=True, max_length=max_length - max_tokens).to(model.device)
outputs = 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)
generated_texts = []
for output, input_text in zip(outputs, batch_texts):
text = tokenizer.decode(output, skip_special_tokens=True)
generated_texts.append(text[len(input_text):])
print(generated_texts)
Citation
If you find this model useful, please cite the following paper:
@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},
}