generate_responses.py

from rl_jailbreak.models.model import load_target
import pandas as pd
import argparse
import time
from tqdm import tqdm
def get_prompt_llama(message: str, system_prompt: str) -> str:
    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
    texts.append(f'{message} [/INST]')
    return ''.join(texts)

def get_prompt_zephyr(message: str, system_prompt: str) -> str:
    texts = f'<s><|system|>\n{system_prompt}</s>\n'
    texts += f'<|user|>\n{message}</s>\n<|assistant|>\n'
    return texts

def get_prompt_vicuna(message: str, system_prompt: str) -> str:
    texts = f'<s>{system_prompt} USER: {message} ASSISTANT:'
    return texts

def main(args):
    prompts = pd.read_csv("./testing/generator_prompts.csv")
    generator_names = prompts["model_name"].unique()

    queries = pd.read_csv("./datasets/ppo_dataset_punctuation.csv")

    queries = queries[queries["source"] == "advbench"].sample(n = 20, random_state=42)["question"]

    # Models vicuna 13B, llama2-7B, Zephyr7B-beta

    # Create a dataframe that stores the prompts for each model.
    data = pd.DataFrame(columns=["generator_model_name", "victim_model_name", "sample_idx", "input", "output"])
    if args.target_model_name == "vicuna":
        target = load_target("/data/public_models/vicuna/vicuna-13b-v1.3")
        # target = None
    elif args.target_model_name == "llama":
        target = load_target("/data/public_models/llama_v2_chat/Llama-2-7b-chat-hf")
    elif args.target_model_name == "zephyr":
        target = load_target("/data/public_models/zephyr/zephyr-7b-beta")
    target_kwargs = {
        "do_sample": True, # yes, we want to sample
        "pad_token_id": target.tokenizer.eos_token_id, # most decoder models don't have a padding token - use EOS token instead
        "max_new_tokens": 250, # specify how many tokens you want to generate at most
    
        "min_new_tokens": 150,
        "repetition_penalty": 2.0,
    }
    # time each iteration time and print it
    
    

    for i, prompt in tqdm(prompts.iterrows()):
        
        batch_data = {
                "generator_model_name": [],
                "victim_model_name": [],
                "sample_idx": [],
                "input": [],
                "output": []
        }
        print(prompt)
        generator_model = prompt["model_name"]
        generator_sample_id = prompt["sample_idx"]
        start = time.time()
        for query in queries:
            inner_start = time.time()
            input = prompt["text"] + ". Here is your task: " + query
            if args.target_model_name == "vicuna":
                wrapped_prompt = get_prompt_vicuna(input, "A chat between a user and an assistant.")
            elif args.target_model_name == "llama":
                wrapped_prompt = get_prompt_llama(input, "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.")
            elif args.target_model_name == "zephyr":
                wrapped_prompt = get_prompt_zephyr(input, "You are a friendly chatbot who always responds in the style of a pirate.")
            output = target.generate(wrapped_prompt, target_kwargs)
            
            batch_data["generator_model_name"].append(generator_model)
            batch_data["victim_model_name"].append(args.target_model_name)
            batch_data["sample_idx"].append(generator_sample_id)
            batch_data["input"].append(input)
            batch_data["output"].append(output)
            inner_end = time.time()
            print(f"Inner Time taken: {inner_end-inner_start} seconds")
        data = pd.concat([data, pd.DataFrame(batch_data)], ignore_index=True)
        end = time.time()
        print(f"Outer Time taken: {end-start} seconds")
    
    data.to_csv(f"testing/target_io_{args.target_model_name}.csv")

if __name__=="__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument("--target-model-name", type=str, default="vicuna", help="The name of the target model to generate prompts for.", choices=["vicuna", "llama", "zephyr"])
    args = argparser.parse_args()
    main(args)