conversers.py

import common
from language_models import GPT, Claude, PaLM, HuggingFace
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from config import VICUNA_PATH, LLAMA_PATH, DEFENSE_TEMP, DEFENSE_TOP_P
import os

from PIL import Image


from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria


from minigpt4.common.config import Config
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Conversation, SeparatorStyle, Chat
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *


from transformers import AutoModelForCausalLM, LlamaTokenizer
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch
from system_prompts import get_scenario_rule


import warnings
warnings.filterwarnings("ignore")

def load_defense_and_target_models(args):
    # Load attack model and tokenizer
    defenseVLM = DefenseVLM(
                            scenario= args.scenario,
                            model_name = args.defense_model,  # "vicuna", "llama"   or  gpt 
                            max_n_tokens = args.defense_max_n_tokens, 
                            max_n_defense_attempts = args.max_n_defense_attempts, 
                            temperature = DEFENSE_TEMP, # init to 1
                            top_p = DEFENSE_TOP_P, # init to 0.9
                        )

    if args.target_model == 'llava':
        targetVLM = MyLlava(args)
    elif args.target_model == 'minigptv2':
        targetVLM = MyMinigptv2(args)
    elif args.target_model == 'cogvlm':
        targetVLM = MyCogVLM(args)
    else:
        print("Unknown target model: {}".format(args.target_model))    

    return defenseVLM , targetVLM


class DefenseVLM():
    """
        Base class for attacker language models.
        
        Generates attacks for conversations using a language model. The self.model attribute contains the underlying generation model.
    """
    def __init__(self, 
                 scenario:str,
                model_name: str, 
                max_n_tokens: int, 
                max_n_defense_attempts: int, 
                temperature: float,
                top_p: float):
        self.scenario = scenario
        self.rule,_  = get_scenario_rule(scenario)
        self.model_name = model_name
        self.temperature = temperature
        self.max_n_tokens = max_n_tokens
        self.max_n_defense_attempts = max_n_defense_attempts
        self.top_p = top_p
        self.model, self.template = load_indiv_model(model_name)
        
        if "vicuna" in model_name or "llama" in model_name:
            self.model.extend_eos_tokens()

    def get_defense_prompt(self, conv, prompt):
        """
        Generates responses for a batch of conversations and prompts using a language model. 
        Only valid outputs in proper JSON format are returned. If an output isn't generated 
        successfully after max_n_defense_attempts, it's returned as None.

        Parameters:
        - convs: the conversation objects.
        - prompts: prompts corresponding to each conversation.
        
        Returns:
        - the generated defense prompt.
        """
        valid_outputs = None

        if len(conv.messages) == 0:
            init_message = """{\"improvement\": \"\",\"prompt\": \""""
        else:
            init_message = """{\"improvement\": \"""" 

        full_prompts = []
        task_prompt = f'''Please optimize the defense prompt based on the response of vision-language moodel to a jailbreak instruction from scenario {self.scenario}. \n
                        The safety rules of scenario {self.scenario} is {self.rule}
                       You should optimize defense prompt based on the specific rules to improve the vision-language model's safety. \n The response of vision-language moodel is: "{prompt}" '''

        conv.append_message(conv.roles[0], task_prompt)
        

        if "gpt" in self.model_name:
            full_prompts.append(conv.to_openai_api_messages())
        else:
            conv.append_message(conv.roles[1], init_message)
            full_prompts.append(conv.get_prompt()[:-len(conv.sep2)])

        for attempt in range(self.max_n_defense_attempts):
            full_output = self.model.generate(full_prompts,
                                                max_n_tokens = self.max_n_tokens,  
                                                temperature = self.temperature,
                                                top_p = self.top_p)
            if "gpt" not in self.model_name:
                full_output = init_message + full_output
            
            defense_dict, json_str = common.extract_json(full_output)
            
            # If outputs are valid, break 
            if defense_dict is not None:
                valid_outputs = defense_dict
                conv.update_last_message(json_str)
                break  #

        if valid_outputs is None:
            print(f"Failed to generate output after {self.max_n_defense_attempts} attempts. Terminating.")
        return valid_outputs


class TargetVLM():
    """
        Base class for target language models.
        Generates responses for prompts using a language model. The self.model attribute contains the underlying generation model.
    """
    def __init__(self, args):
        self.model_name =  args.model_name
        self.temperature =  args.temperature
        self.max_n_tokens =  args.max_n_tokens
        self.top_p =  args.top_p

    def get_response(self, qs, defense_query, full_image_path):
        raise NotImplementedError



class MyLlava(TargetVLM):
    def __init__(self, args):
        # Model
        model_path = os.path.expanduser(args.model_path)
        model_name = get_model_name_from_path(model_path)
        disable_torch_init()
        tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, args.model_base, model_name)
        self.tokenizer = tokenizer
        self.model = model
        self.image_processor = image_processor


        self.temperature = args.temperature
        self.top_p = args.top_p
        self.num_beams = args.num_beams
        self.max_new_tokens = args.max_new_tokens
        
        if "llama-2" in model_name.lower():
            self.conv_mode = "llava_llama_2"
        elif "v1" in model_name.lower():
            self.conv_mode = "llava_v1"
        elif "mpt" in model_name.lower():
            self.conv_mode = "mpt"
        else:
            self.conv_mode = "llava_v0"

        if args.conv_mode is None or self.conv_mode == args.conv_mode:
            args.conv_mode = self.conv_mode
    '''
        @ input: 
            qs: text prompt, type: str
            defense_query: defense prompt, type: str
            full_image_path: os.path.join(image_folder, image_file)
        @ output: responses
    '''       
    def get_response(self, qs, defense_query, full_image_path):
        qs =  qs + defense_query + qs

        if self.model.config.mm_use_im_start_end:
            qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
        else:
            qs = DEFAULT_IMAGE_TOKEN + '\n' + qs

        conv = conv_templates[self.conv_mode].copy()
        conv.append_message(conv.roles[0], qs)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        
        input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
        image = Image.open(full_image_path)
        image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
        
        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2


        with torch.inference_mode():
            output_ids = self.model.generate(
                    input_ids,
                    images=image_tensor.unsqueeze(0).half().cuda(),
                    do_sample=True,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    num_beams=self.num_beams,
                    max_new_tokens=self.max_new_tokens,
                    use_cache=True)
            
        input_token_len = input_ids.shape[1]
        n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
        if n_diff_input_output > 0:
            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
        outputs = self.tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
        outputs = outputs.strip()
        if outputs.endswith(stop_str):
            print('stop str')
            outputs = outputs[:-len(stop_str)]    
        return outputs
    
class MyMinigptv2(TargetVLM):
    def __init__(self, args):
        cfg = Config(args)
        device = 'cuda:{}'.format(args.gpu_id)
        model_config = cfg.model_cfg
        model_config.device_8bit = args.gpu_id
        model_cls = registry.get_model_class(model_config.arch)
        model = model_cls.from_config(model_config).to(device)
        bounding_box_size = 100
        
        vis_processor_cfg = cfg.datasets_cfg.cc_sbu_align.vis_processor.train
        vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)

        model = model.eval()

        self.CONV_VISION = Conversation(
            system="",
            roles=(r"<s>[INST] ", r" [/INST]"),
            messages=[],
            offset=2,
            sep_style=SeparatorStyle.SINGLE,
            sep="",
        )
        self.chat = Chat(model, vis_processor, device=device)
    
    '''
        @ input: 
            qs: text prompt, type: str
            defense_query: defense prompt, type: str
            full_image_path: os.path.join(image_folder, image_file)
        @ output: responses
    '''    
    def VQA(self, image_file, prompt):
        image_file = image_file
        prompt = prompt[:2000]

        # ask
        chat_state = self.CONV_VISION.copy()
        chatbot = []
        img_list = []
        llm_message = self.chat.upload_img(image_file, chat_state, img_list)

        self.chat.ask(prompt, chat_state)

        chatbot = chatbot + [[prompt, None]]

        # answer
        if len(img_list) > 0:
            if not isinstance(img_list[0], torch.Tensor):
                self.chat.encode_img(img_list)

        llm_message = self.chat.answer(conv=chat_state,
                              img_list=img_list,
                              temperature=0.2,
                              max_new_tokens=500,
                              max_length=2048)[0]
        return llm_message
    
    def get_response(self, qs, defense_query, full_image_path):
        query = qs + defense_query + qs
        image = Image.open(full_image_path).convert('RGB')
        with torch.no_grad():
            respone = self.VQA(image, query)
        return respone
    
class MyCogVLM(TargetVLM):
    def __init__(self, args):
        self.tokenizer = LlamaTokenizer.from_pretrained(args.llm_path) #
        model_name = os.path.basename(args.llm_path)
        
        with init_empty_weights():
            model = AutoModelForCausalLM.from_pretrained(
                    args.model_path,
                    torch_dtype=torch.bfloat16,
                    low_cpu_mem_usage=True,
                    trust_remote_code=True)
            model = load_checkpoint_and_dispatch(
                    model,
                    args.model_path,
                    device_map="auto",
                    no_split_module_classes=['CogVLMDecoderLayer', 'TransformerLayer'])
            
        self.model = model.eval()
        self.temperature = args.temperature
        self.top_p = args.top_p
        self.top_k = args.top_k
        self.max_length = args.max_length


    def get_response(self, qs, defense_query, full_image_path):
        query = qs +  defense_query + "'" + qs  + "'"
        image = Image.open(full_image_path).convert('RGB')
        inputs = self.model.build_conversation_input_ids(self.tokenizer, query=query, history=[], images=[image])  # chat mode
        input_ids = inputs['input_ids'].unsqueeze(0).to('cuda')
        
        with torch.no_grad():
            self.model.llm_embedding = ()
            outputs = self.model.generate(
                            input_ids = input_ids,
                            images = [[inputs['images'][0].to('cuda').to(torch.bfloat16)]],
                            token_type_ids = inputs['token_type_ids'].unsqueeze(0).to('cuda'),
                            attention_mask = inputs['attention_mask'].unsqueeze(0).to('cuda'),
                            do_sample = True,
                            temperature =self.temperature,
                            top_p = self.top_p,
                            top_k = self.top_k,
                            max_length = self.max_length
                        )
            
        outputs = outputs[:, input_ids.shape[1]:]
        respone = self.tokenizer.decode(outputs[0])
        return respone

def load_indiv_model(model_name, device=None):
    model_path, template = get_model_path_and_template(model_name)
    if model_name in ["gpt-3.5-turbo", "gpt-4"]:
        lm = GPT(model_name)
    elif model_name in ["claude-2", "claude-instant-1"]:
        lm = Claude(model_name)
    elif model_name in ["palm-2"]:
        lm = PaLM(model_name)
    else:
        model = AutoModelForCausalLM.from_pretrained(
                model_path, 
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,device_map="auto").eval()

        tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            use_fast=False
        ) 

        if 'llama-2' in model_path.lower():
            tokenizer.pad_token = tokenizer.unk_token
            tokenizer.padding_side = 'left'
        if 'vicuna' in model_path.lower():
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.padding_side = 'left'
        if not tokenizer.pad_token:
            tokenizer.pad_token = tokenizer.eos_token

        lm = HuggingFace(model_name, model, tokenizer)
    
    return lm, template

def get_model_path_and_template(model_name):
    full_model_dict={
        "gpt-4":{
            "path":"gpt-4",
            "template":"gpt-4"
        },
        "gpt-3.5-turbo": {
            "path":"gpt-3.5-turbo",
            "template":"gpt-3.5-turbo"
        },
        "vicuna":{
            "path": VICUNA_PATH,
            "template":"vicuna_v1.1"
        },
        "llama-2":{
            "path":LLAMA_PATH,
            "template":"llama-2"
        },
        "claude-instant-1":{
            "path":"claude-instant-1",
            "template":"claude-instant-1"
        },
        "claude-2":{
            "path":"claude-2",
            "template":"claude-2"
        },
        "palm-2":{
            "path":"palm-2",
            "template":"palm-2"
        }
    }
    path, template = full_model_dict[model_name]["path"], full_model_dict[model_name]["template"]
    return path, template