eval_response.py

'''
Some functions in this code is from:
https://github.com/meta-llama/llama-recipes/blob/main/recipes/responsible_ai/llama_guard/llama_guard_customization_via_prompting_and_fine_tuning.ipynb
'''

import torch, os, logging, argparse
from accelerate import Accelerator
from utils.general import (
    get_model_path, setup_logging,
    load_model_and_tokenizer
)
from utils.jailbreak import MALICIOUS_CATEGORIES, contains_refusal_phrase
from utils.prompt_format import AgentType
from utils.safety import evaluate_safety

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

def evaluate_response_based_on_refusal(directory):
    """Process each file in the directory and evaluate safety."""
    x_values, y_values = [], []

    for filename in os.listdir(directory):
        if filename.endswith("shot.txt"):
            num_shots = int(filename.split('shot')[0])
            x_values.append(num_shots)

            file_path = os.path.join(directory, filename)
            with open(file_path, 'r', encoding='utf-8') as file:
                logging.info(f'Processing {directory}/{filename}')

                # Evaluate each response
                unsafe = 0
                num_response = 0
                for response in file:
                    response = response.strip()
                    if not response:
                        continue
                    try:
                        num_response += 1
                        unsafe += not contains_refusal_phrase(response)
                    except torch.cuda.OutOfMemoryError:
                        logging.error("CUDA out of memory. Reducing batch size or model size is advised.")
                        torch.cuda.empty_cache()
                if num_response == 0:
                    logging.warning(f"No responses found in {directory}/{filename}")
                    # Remove the x value if no responses are found
                    x_values = [x_value for x_value in x_values if x_value != num_shots]
                else:
                    y_values.append(unsafe / num_response * 100)

    return x_values, y_values

def evaluate_response_based_on_llm(directory, model, tokenizer, agent_type = AgentType.AGENT):
    """Process each file in the directory and evaluate safety."""
    x_values, y_values = [], []

    for filename in os.listdir(directory):
        if filename.endswith("shot.txt"):
            num_shots = int(filename.split('shot')[0])
            x_values.append(num_shots)

            file_path = os.path.join(directory, filename)
            with open(file_path, 'r', encoding='utf-8') as file:
                logging.info(f'Processing {directory}/{filename}')

                # Evaluate each response
                unsafe = 0
                num_response = 0
                for response in file:
                    response = response.strip()
                    if not response:
                        continue
                    try:
                        num_response += 1
                        unsafe += 'unsafe' in evaluate_safety(tokenizer, model, agent_type, prompt=response)
                    except torch.cuda.OutOfMemoryError:
                        logging.error("CUDA out of memory. Reducing batch size or model size is advised.")
                        torch.cuda.empty_cache()
                if num_response == 0:
                    logging.warning(f"No responses found in {directory}/{filename}")
                    # Remove the x value if no responses are found
                    x_values = [x_value for x_value in x_values if x_value != num_shots]
                else:
                    y_values.append(unsafe / num_response * 100)

    return x_values, y_values

def evaluate_response(directory, model, tokenizer, agent_type = AgentType.AGENT):
    """Process each file in the directory and evaluate safety."""
    x_values, y_values_llm, y_values_refusal, y_values_or = [], [], [], []

    for filename in os.listdir(directory):
        if filename.endswith("shot.txt"):
            num_shots = int(filename.split('shot')[0])
            x_values.append(num_shots)

            file_path = os.path.join(directory, filename)
            with open(file_path, 'r', encoding='utf-8') as file:
                logging.info(f'Processing {directory}/{filename}')

                # Evaluate each response
                unsafe_llm, unsafe_refusal, unsafe_or = 0, 0, 0
                num_response = 0
                for response in file:
                    response = response.strip()
                    if not response:
                        continue
                    try:
                        num_response += 1
                        result_llm = 'unsafe' in evaluate_safety(tokenizer, model, agent_type, prompt=response)
                        result_refusal = not contains_refusal_phrase(response)
                        unsafe_llm += result_llm
                        unsafe_refusal += result_refusal
                        unsafe_or += result_refusal or result_llm
                    except torch.cuda.OutOfMemoryError:
                        logging.error("CUDA out of memory. Reducing batch size or model size is advised.")
                        torch.cuda.empty_cache()
                if num_response == 0:
                    logging.warning(f"No responses found in {directory}/{filename}")
                    # Remove the x value if no responses are found
                    x_values = [x_value for x_value in x_values if x_value != num_shots]
                else:
                    y_values_llm.append(unsafe_llm / num_response * 100)
                    y_values_refusal.append(unsafe_refusal / num_response * 100)
                    y_values_or.append(unsafe_or / num_response * 100)

    return x_values, y_values_llm, y_values_refusal, y_values_or

def save_results(directory, x_values, y_values, method):
    """Save x and y values to a .pt file."""
    sorted_pairs = sorted(zip(x_values, y_values))
    x_values_sorted, y_values_sorted = zip(*sorted_pairs)

    save_data = {'x_values': x_values_sorted, 'y_values': y_values_sorted}
    torch.save(save_data, f'{directory}/result-{method}.pt')

    logging.info(f"Results saved to {directory}/result-{method}.pt")

def main():

    # Set up argument parser
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--directory",
                        type=str, required=True, help="Directory containing the files to be evaluated.")
    parser.add_argument("--method", type=str, choices=['llm', 'refusal'], default='llm')
    args = parser.parse_args()

    setup_logging(args)

    # Move model to GPU if available
    accelerator = Accelerator()
    current_gpu_id = accelerator.process_index
    device = 'cuda'
    logging.info(f' *GPU-{current_gpu_id}: Accelerator initialized, using device: {accelerator.device}')

    # Load model and tokenizer
    model_name = 'Llama-Guard-3-8B'
    model_path = get_model_path(model_name)
    model, tokenizer = load_model_and_tokenizer(model_path, current_gpu_id, device)

    # Process the generated promts:: malicious, benign, common
    agent_type = AgentType.AGENT if args.agent_type == 'agent' else AgentType.USER
    x_values, y_values = evaluate_response_based_on_llm(args.directory, model, tokenizer, agent_type)
    save_results(args.directory, x_values, y_values, 'llm')

    x_values, y_values = evaluate_response_based_on_refusal(args.directory)
    save_results(args.directory, x_values, y_values, 'refusal')

if __name__ == '__main__':
    main()