chatbug.py

import pandas as pd
import json
import itertools

import os
from datetime import datetime
import time

import hydra
from hydra import compose, initialize
import transformers
from transformers import AutoTokenizer
import torch

import wandb
import omegaconf
import argparse
from utils import *
from artprompt import *
import pickle as pkl

class Chatbug:
    def __init__(self, config):
        self.config = config
        self.logging_enable = self.config.wandb.enable
        self.attack_method = self.config.attack_method
        self.jailbreak_add = self.config.jailbreak_add
        self.logging_init()

        self.batch_size = self.config.llm_params.batch_size
        self.llm = None
        
        self.system_prompt = self.get_system_prompt()
        self._output_path = None
        self.mismatch_tokenizer = None
        transformers.set_seed(self.config.seed)

    def logging_init(self):
        if self.config.wandb.enable:
            wandb.config = omegaconf.OmegaConf.to_container(
                self.config, resolve=True, throw_on_missing=True
            )
            if self.config.system_prompt is None:
                sp_tag = "No-SP"
            elif self.config.system_prompt == "default":
                sp_tag = "Default-SP"
            else:
                sp_tag = "Custom-SP"

            tags = [
                    self.config.llm_params.model_name, 
                    self.config.attack_method,
                    sp_tag
                ]
            if self.jailbreak_add is not None:
                tags.append(self.jailbreak_add)

            if self.config.llm_params.lora_path is not None:
                tags.append(f"PEFT-{self.config.llm_params.lora_path }")

            wandb.init(
                project=self.config.wandb.project,
                tags=tags
            )

            wandb.define_metric("ASR_kw", summary="last")
            wandb.define_metric("ASR_llm", summary="last")
            wandb.define_metric("AttackTime", summary="last")
            wandb.define_metric("AttackTime_Ave", summary="last")
        else:
            print("Logging disabled")

    def _model_name(self):
        model_name = self.config.llm_params.model_name
        if "/" in model_name and not self.config.llm_params.use_api: 
            model_name = model_name.split("/")[-1]

        return model_name
    
    @property
    def output_path(self):
        if self._output_path is None:
            output_dir = os.path.join(self.config.output_dir, self._model_name())
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            time_entry = datetime.now().strftime("Y%Y-M%m-D%d-H%H-M%M-S%S")
            fill_name = f"{self.attack_method}"
            if self.jailbreak_add is not None:
                fill_name += f"_{self.jailbreak_add}"
            fill_name += f"_{time_entry}"
            if self.config.llm_params.lora_path is not None:
                fill_name += f"_PEFT-{self.config.llm_params.lora_path.split('/')[-1]}"
            self._output_path = os.path.join(output_dir, fill_name+".json")
            
        return self._output_path

    def get_system_prompt(self):
        if self.config.system_prompt is not None and self.config.system_prompt == "default":
            for model, system_prompt in system_prompt.items():
                if model in self.config.model_name.lower():
                    return system_prompt
                
            raise ValueError(f"No system prompt found for the model {self.llm.model_name}")
        return None

    ### attack prompt generation
    def no_template(self, instruction, target=None, system_prompt=None):
        res = ''
        if system_prompt is not None:
            res += system_prompt + " " if system_prompt[-1] == '.' else system_prompt + ". "
        res += instruction if instruction[-1] == '.' else instruction + "."
        if target is not None:
            res += " " + target
        return [res]

    def direct_inst(self, instruction, target, system_prompt=None):
        chat = []
        if system_prompt is not None:
            chat.append({"role": "system", "content": system_prompt})
        chat.append({"role": "user", "content": instruction})
        
        if not self.llm.use_api:
            chat_text = self.llm.tokenizer.apply_chat_template(
                chat, tokenize=False, add_generation_prompt=True,
            )
            if chat_text[:self.llm.bos_len] == self.llm.tokenizer.bos_token:
                chat_text = chat_text[self.llm.bos_len:]
            
            return [chat_text]
        else:
            return [chat]
        
    def overflow_short(self, instruction, target, system_prompt=None, short_start="Sure, here is"):
        chat = []
        if system_prompt is not None:
            chat.append({"role": "system", "content": system_prompt})
        chat.append({"role": "user", "content": instruction})

        if not self.llm.use_api:
            chat_text = self.llm.tokenizer.apply_chat_template(
                chat, tokenize=False, add_generation_prompt=True,
            )
            if chat_text[:self.llm.bos_len] == self.llm.tokenizer.bos_token:
                chat_text = chat_text[self.llm.bos_len:]
            chat_text += short_start
            return [chat_text]
        else:
            chat.append({"role": "assistant", "content": short_start})
            return [chat]

    def overflow(self, instruction, target, system_prompt=None):
        '''
        return:
            chat_text: str|list[dict] 
                if not using api, else message dict list 
        '''

        chat = []
        if system_prompt is not None:
            chat.append({"role": "system", "content": system_prompt})
        chat.append({"role": "user", "content": instruction})

        if not self.llm.use_api:
            chat_text:'str' = self.llm.tokenizer.apply_chat_template(
                chat, tokenize=False, add_generation_prompt=True,
            )
            if chat_text[:self.llm.bos_len] == self.llm.tokenizer.bos_token:
                chat_text = chat_text[self.llm.bos_len:]
            chat_text += target
            return [chat_text]
        else:
            chat.append({"role": "assistant", "content": target})
            return [chat]

    def overflow_fs(self, instruction, target_list, system_prompt=None):
        # generate fs target
        assert isinstance(target_list, list), "Invalid target list"
        res = []
        for target in target_list:
            res.extend(self.overflow(instruction, target, system_prompt))
        return res
        
    def mischat(self, instruction, target, system_prompt=None, template="no", add_target=False):
        assert not self.llm.use_api or "gpt" in self.config.llm_params.model_name, "mismatch attack is not supported for API models except for gpt-3.5-azure"

        if "gpt" in self.config.llm_params.model_name:
            assert self.llm.chat_template is not None, "Chat template is not loaded for gpt model"
            self.llm.chat_template = load_jinja_template(chat_template["no_template"])
            return self.direct_inst(instruction, target, system_prompt)


        if template == "no":
            target = None if not add_target else target
            return self.no_template(instruction, target, system_prompt)

        if self.mismatch_tokenizer is None:
            if template == "vicuna":
                self.mismatch_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")
                self.mismatch_tokenizer.chat_template = chat_template['vicuna']
            elif template == "chatml":
                self.mismatch_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")
                self.mismatch_tokenizer.chat_template = chat_template['chatml']
            else:
                raise ValueError(f"Invalid template {template}")
            
        chat = []
        if system_prompt is not None:
            chat.append({"role": "system", "content": system_prompt})
        chat.append({"role": "user", "content": instruction})

        chat_text = self.mismatch_tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True,
        )
        if chat_text[:len(self.mismatch_tokenizer.bos_token)] == self.mismatch_tokenizer.bos_token:
            chat_text = chat_text[len(self.mismatch_tokenizer.bos_token):]

        if add_target:
            chat_text += target + "."

        return [chat_text]

    def mismatch_no(self, instruction, target, system_prompt=None):
        return self.mischat(instruction, target, system_prompt, template="no", add_target=False)

    def mismatch_vicuna(self, instruction, target, system_prompt=None): 
        return self.mischat(instruction, target, system_prompt, template="vicuna", add_target=False)
    
    def mismatch_chatml(self, instruction, target, system_prompt=None):
        return self.mischat(instruction, target, system_prompt, template="chatml", add_target=False)


    def jailbreak_prompt_gen(self, instruction_list, target_list):
        prompt_dict = {}
        prompt_gen_func = getattr(self, self.attack_method)

        for i, (inst, target) in enumerate(zip(instruction_list, target_list)):
            if self.attack_method in ["overflow"]:
                target += "."
            if self.jailbreak_add is None:
                prompt_dict[i] = prompt_gen_func(inst, target, self.system_prompt)
            else:
                if '_' in self.jailbreak_add:
                    inst_list = getattr(self, self.jailbreak_add.split("_")[0])(inst, *self.jailbreak_add.split("_")[1:])
                else:
                    inst_list = getattr(self, self.jailbreak_add)(inst)
                prompt_dict[i] = [prompt_gen_func(inst, target, self.system_prompt) for inst in inst_list]
                prompt_dict[i] = list(itertools.chain(*prompt_dict[i]))
        return prompt_dict
    
    def fs_gen(self, instruction_list):
        fs_file_path = os.path.join(self.config.fs_target_dir, f"fs_target_n{self.config.llm_attack_params.generation_config.num_return_sequences}.json")
        if not os.path.exists(fs_file_path):
            print("FS target file not found, generating...")
            llm = LLM(self.config.llm_attack_params)

            instruction_list = [fs_gen_prompt(inst) for inst in instruction_list]
            message_list = []
            for chat in instruction_list:
                chat_str = llm.tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
                if chat_str[:llm.bos_len] == llm.tokenizer.bos_token:
                    chat_str = chat_str[llm.bos_len:]
                message_list.append(chat_str)
            
            targets = llm.generate_all(message_list)
            targets_by_inst = []
            num_per_inst = llm.config.generation_config.num_return_sequences
            for i in range(len(instruction_list)):
                targets_by_inst.append(targets[i*num_per_inst:(i+1)*num_per_inst])

            llm = None
            torch.cuda.empty_cache()
            with open(fs_file_path, "w") as f:
                json.dump(targets_by_inst, f, indent=4)
        else:
            print("FS target file found, loading...")
            with open(fs_file_path, "r") as f:
                targets_by_inst = json.load(f)

        return targets_by_inst


    ###################################
    def jailbreak_dataset_load(self):
        dataset = pd.read_csv(self.config.dataset_path)
        instruction, target = dataset['goal'].to_list(), dataset['target'].to_list()
        start = 0 if self.config.start is None else self.config.start
        end = len(instruction) if self.config.end is None else self.config.end

        if "fs" in self.attack_method:
            target = self.fs_gen(instruction)

        return instruction[start:end], target[start:end]

    def jailbreak_llm(self, eval_only=None):
        if eval_only is None:
            instruction, target = self.jailbreak_dataset_load()
            start_time = time.time()
        
            self.llm = LLM(self.config.llm_params)
            prompt_dict = self.jailbreak_prompt_gen(instruction, target)
            merged_list = []
            for idx in range(len(prompt_dict)):
                merged_list.extend(prompt_dict[idx])

            print(f"Prompt generation completed for {len(prompt_dict)} instruction with {len(merged_list)} attack prompts")
            resp_list = self.llm.generate_all(merged_list)

            attack_time = time.time() - start_time
            attack_time_ave = attack_time/len(merged_list)

            res_json = {}
            count = 0
            for i in range(len(prompt_dict)):
                res_json[i] = {
                    "goal": instruction[i],
                    "target": target[i],
                    "prompt_list": prompt_dict[i],
                    "resp_list": resp_list[count:count+len(prompt_dict[i])]
                }
                count += len(prompt_dict[i])

            with open(self.output_path, "w") as f:
                json.dump(res_json, f, indent=4)
        
            if self.logging_enable:
                wandb.save(self.output_path)

        
            self.llm.model = None
            self.llm = None
            torch.cuda.empty_cache()
        

            ASR_kw, ASR_llm = self.eval(self.output_path)

        else:
            ASR_kw, ASR_llm = self.eval(eval_only)
            attack_time = 0
            attack_time_ave = 0

        if self.logging_enable:
            wandb.log({"ASR_kw": ASR_kw, "ASR_llm": ASR_llm, "AttackTime": attack_time, "AttackTime_Ave": attack_time_ave})


    def eval(self, res_json_path, save=False):
        with open(res_json_path, "r") as f:
            res_json = json.load(f)

        if self.config.eval.kw_enable:
            ASR_kw, ASR_kw_labels = self.eval_kw(res_json)
        else:
            ASR_kw = None

        if self.config.eval.llm_enable:
            ASR_llm, ASR_llm_labels = self.eval_llm(res_json)
        else:
            ASR_llm = None

        if save:
            print("Saving intermediate results... Turn off unless desired")
            with open("./KW_ASR_Labels.json", "w") as f:
                json.dump(ASR_kw_labels, f, indent=4)
            with open("./LLM_ASR_Labels.json", "w") as f:
                json.dump(ASR_llm_labels, f, indent=4)

        return ASR_kw, ASR_llm

    @staticmethod
    def kw_is_jailbreak_one(response, prefix_list):
        return all([prefix not in response for prefix in prefix_list])
    
    def kw_is_jailbreak_list(self, response_list, prefix_list):
        jailbreak_labels = []
        for response in response_list:
            jailbreak_labels.append(self.kw_is_jailbreak_one(response, prefix_list))
        return jailbreak_labels

    def eval_kw(self, res_json, prefix_path='datasets/test_prefixes.csv'):
        kw_list = pd.read_csv(prefix_path, header=None)[0]
        jb_succ = 0
        jailbreak_array = []
        for i in range(len(res_json)):
            resp_eval_labels = self.kw_is_jailbreak_list(res_json[str(i)]['resp_list'], kw_list)
            jailbreak_array.append(resp_eval_labels)
            if any(resp_eval_labels):
                jb_succ += 1

        print(f"KW ASR: {jb_succ/len(res_json)}")
        return jb_succ/len(res_json), jailbreak_array
        
    def eval_llm(self, res_json):
        llm_eval = LLM_Evalator(self.config.llm_eval_params)
        
        # assemble message list
        msg_list = []
        for i in range(len(res_json)):
            for j in range(len(res_json[str(i)]['resp_list'])):
                msg_entry = [
                    {"role": "user", "content": res_json[str(i)]['goal']},
                    {"role": "assistant", "content": res_json[str(i)]['resp_list'][j]},
                ]
                msg_list.append(msg_entry)

        eval_list = llm_eval.moderate(msg_list)
        
        unsafe_count = 0
        ptr = 0
        jailbreak_eval_labels = []
        for i in range(len(res_json)):
            eval_sub_list = eval_list[ptr:ptr+len(res_json[str(i)]['resp_list'])]
            ptr += len(res_json[str(i)]['resp_list'])
            eval_sub_list = [succ for _, succ, _ in eval_sub_list]
            if any(eval_sub_list):
                unsafe_count += 1
            jailbreak_eval_labels.append(eval_sub_list)

        print(f"LLM ASR: {unsafe_count/len(res_json)}")
        return unsafe_count/len(res_json), jailbreak_eval_labels
                
def override_config(config, args):
    if args.model is not None:
        if args.model in model_map:
            config.llm_params.model_name = model_map[args.model]
        else:
            config.llm_params.model_name = args.model # direct set model name without any check

        if any([model in args.model for model in ["gpt", "gemini", "claude"]]):
            config.llm_params.use_api = True
    if args.attack_method is not None:
        config.attack_method = args.attack_method
    if args.jb_add is not None:
        config.jailbreak_add = args.jb_add

    if args.eval_only is not None:
        model_name = args.eval_only.split("/")[-2]
        config.llm_params.model_name = model_name

        attack_method = args.eval_only.split("/")[-1].split("Y2024")[0]
        config.attack_method = attack_method

    return config
        

def main(args):
    with initialize(version_base="1.2", config_path=''):
        config = compose(config_name="attack")

    config = override_config(config, args)

    bc = Chatbug(config)
    bc.jailbreak_llm(args.eval_only)
    
    return 

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default=None)
    parser.add_argument("--attack_method", type=str, default=None)
    parser.add_argument("--jb_add", type=str, default=None)
    parser.add_argument("--eval_only", type=str, default=None)
    args = parser.parse_args()

    main(args)