119 lines
5 KiB
Python
119 lines
5 KiB
Python
"""
|
|
DeepInception Class
|
|
============================================
|
|
This class can easily hypnotize LLM to be a jailbreaker and unlock its
|
|
misusing risks.
|
|
|
|
Paper title: DeepInception: Hypnotize Large Language Model to Be Jailbreaker
|
|
arXiv Link: https://arxiv.org/pdf/2311.03191.pdf
|
|
Source repository: https://github.com/tmlr-group/DeepInception
|
|
"""
|
|
import json
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
from easyjailbreak.metrics.Evaluator import EvaluatorGenerativeJudge
|
|
from easyjailbreak.attacker import AttackerBase
|
|
from easyjailbreak.datasets import JailbreakDataset, Instance
|
|
from easyjailbreak.mutation.rule import Inception
|
|
|
|
from tqdm import tqdm
|
|
|
|
__all__ = ['DeepInception']
|
|
|
|
class DeepInception(AttackerBase):
|
|
r"""
|
|
DeepInception is a class for conducting jailbreak attacks on language models.
|
|
"""
|
|
def __init__(self, attack_model, target_model, eval_model, jailbreak_datasets: JailbreakDataset, save_path, dataset_name, scene=None, character_number=None, layer_number=None):
|
|
r"""
|
|
Initialize the DeepInception attack instance.
|
|
:param attack_model: In this case, the attack_model should be set as None.
|
|
:param target_model: The target language model to be attacked.
|
|
:param eval_model: The evaluation model to evaluate the attack results.
|
|
:param jailbreak_datasets: The dataset to be attacked.
|
|
:param template_file: The file path of the template.
|
|
:param scene: The scene of the deepinception prompt (The default value is 'science fiction').
|
|
:param character_number: The number of characters in the deepinception prompt (The default value is 4).
|
|
:param layer_number: The number of layers in the deepinception prompt (The default value is 5).
|
|
"""
|
|
super().__init__(attack_model, target_model, eval_model, jailbreak_datasets)
|
|
self.current_query: int = 0
|
|
self.current_jailbreak: int = 0
|
|
self.current_reject: int = 0
|
|
self.scene = scene
|
|
self.character_number = character_number
|
|
self.layer_number = layer_number
|
|
self.evaluator = EvaluatorGenerativeJudge(eval_model)
|
|
self.mutation = Inception(attr_name='query')
|
|
|
|
self.save_path = save_path
|
|
self.dataset_name = dataset_name
|
|
|
|
def single_attack(self, instance: Instance) -> JailbreakDataset:
|
|
r"""
|
|
single_attack is a method for conducting jailbreak attacks on language models.
|
|
"""
|
|
new_instance_list = []
|
|
|
|
instance_ds = JailbreakDataset([instance])
|
|
new_instance = self.mutation(instance_ds)[-1]
|
|
system_prompt = new_instance.jailbreak_prompt.format(query = new_instance.query)
|
|
|
|
if self.scene is not None:
|
|
system_prompt = system_prompt.replace('science fiction', self.scene)
|
|
if self.character_number is not None:
|
|
system_prompt = system_prompt.replace('4', str(self.character_number))
|
|
if self.layer_number is not None:
|
|
system_prompt = system_prompt.replace('5', str(self.layer_number))
|
|
new_instance.jailbreak_prompt = system_prompt
|
|
answer = self.target_model.generate(system_prompt.format(query = new_instance.query))
|
|
new_instance.target_responses.append(answer)
|
|
new_instance_list.append(new_instance)
|
|
|
|
return JailbreakDataset(new_instance_list)
|
|
|
|
def attack(self):
|
|
r"""
|
|
Execute the attack process using provided prompts.
|
|
"""
|
|
logging.info("Jailbreak started!")
|
|
self.attack_results = JailbreakDataset([])
|
|
try:
|
|
with open(self.save_path, 'w') as f:
|
|
for Instance in tqdm(self.jailbreak_datasets):
|
|
if self.dataset_name == "trustllm":
|
|
self.target_model.set_system_message(Instance.system_message)
|
|
|
|
results = self.single_attack(Instance)
|
|
for new_instance in results:
|
|
self.attack_results.add(new_instance)
|
|
|
|
line = new_instance.to_dict()
|
|
f.write(json.dumps(line, ensure_ascii=False) + '\n')
|
|
|
|
except KeyboardInterrupt:
|
|
logging.info("Jailbreak interrupted by user!")
|
|
# self.evaluator(self.attack_results)
|
|
# self.update(self.attack_results)
|
|
logging.info("Jailbreak finished!")
|
|
|
|
def update(self, Dataset: JailbreakDataset):
|
|
r"""
|
|
Update the state of the Jailbroken based on the evaluation results of Datasets.
|
|
|
|
:param Dataset: The Dataset that is attacked.
|
|
"""
|
|
for prompt_node in Dataset:
|
|
self.current_jailbreak += prompt_node.num_jailbreak
|
|
self.current_query += prompt_node.num_query
|
|
self.current_reject += prompt_node.num_reject
|
|
|
|
def log(self):
|
|
r"""
|
|
Report the attack results.
|
|
"""
|
|
logging.info("======Jailbreak report:======")
|
|
logging.info(f"Total queries: {self.current_query}")
|
|
logging.info(f"Total jailbreak: {self.current_jailbreak}")
|
|
logging.info(f"Total reject: {self.current_reject}")
|
|
logging.info("========Report End===========")
|