loggers.py

import os
import wandb
import pytz
from datetime import datetime
import pandas as pd


class WandBLogger:
    """WandB logger."""

    def __init__(self, args, system_prompt):
        self.logger = wandb.init(
            project = "jailbreak-llms",
            config = {
                "attack_model" : args.attack_model,
                "target_model" : args.target_model,
                "judge_model": args.judge_model,
                "keep_last_n": args.keep_last_n,
                "system_prompt": system_prompt,
                "index": args.index,
                "category": args.category,
                "goal": args.goal,
                "n_iter": args.n_iterations,
                "target_str": args.target_str,
                "n_streams": args.n_streams,

            }
        )
        self.is_jailbroken = False
        self.query_to_jailbreak = None
        self.table = pd.DataFrame()
        self.batch_size = args.n_streams
        self.index = args.index
        self.goal = args.goal
        self.jailbreak_prompt = None
        self.jailbreak_response = None

    def log(self, iteration: int, attack_list: list, response_list: list, judge_scores: list):
        
        df = pd.DataFrame(attack_list)
        df["target_response"] = response_list
        df["judge_scores"] = judge_scores
        df["iter"] = iteration
        df["conv_num"] = [i+1 for i in range(len(response_list))]
        self.table = pd.concat([self.table, df])

        if any([score == 10 for score in judge_scores]):
            jailbreak_ind = judge_scores.index(10)
            if not self.is_jailbroken:
                self.query_to_jailbreak = self.batch_size * (iteration-1) + jailbreak_ind + 1 # Add one since jailbreak_ind is 0-indexed
                self.logger.log({"queries_to_jailbreak": self.query_to_jailbreak})
                self.is_jailbroken = True

            self.jailbreak_prompt = attack_list[jailbreak_ind]["prompt"]
            self.jailbreak_response = response_list[jailbreak_ind]

        self.logger.log({
            "iteration":iteration,
            "judge_scores":judge_scores,
            "mean_judge_score_iter":sum(judge_scores)/len(judge_scores),
            "is_jailbroken":self.is_jailbroken,
            "max_judge_score":self.table["judge_scores"].max(),
            "jailbreak_prompt":self.jailbreak_prompt,
            "jailbreak_response":self.jailbreak_response,
            "data": wandb.Table(data = self.table)})

        self.print_summary_stats(iteration)

    def finish(self):
        self.print_final_summary_stats()
        self.logger.finish()

    def print_summary_stats(self, iter):
        bs = self.batch_size
        df = self.table 
        mean_score_for_iter = df[df['iter'] == iter]['judge_scores'].mean()
        max_score_for_iter = df[df['iter'] == iter]['judge_scores'].max()
        
        num_total_jailbreaks = df[df['judge_scores'] == 10]['conv_num'].nunique()
        
        jailbreaks_at_iter = df[(df['iter'] == iter) & (df['judge_scores'] == 10)]['conv_num'].unique()
        prev_jailbreaks = df[(df['iter'] < iter) & (df['judge_scores'] == 10)]['conv_num'].unique()

        num_new_jailbreaks = len([cn for cn in jailbreaks_at_iter if cn not in prev_jailbreaks])

        print(f"{'='*14} SUMMARY STATISTICS {'='*14}")
        print(f"Mean/Max Score for iteration: {mean_score_for_iter:.1f}, {max_score_for_iter}")
        print(f"Number of New Jailbreaks: {num_new_jailbreaks}/{bs}")
        print(f"Total Number of Conv. Jailbroken: {num_total_jailbreaks}/{bs} ({num_total_jailbreaks/bs*100:2.1f}%)\n")

    def print_final_summary_stats(self):
        print(f"{'='*8} FINAL SUMMARY STATISTICS {'='*8}")
        print(f"Index: {self.index}")
        print(f"Goal: {self.goal}")
        df = self.table
        if self.is_jailbroken:
            num_total_jailbreaks = df[df['judge_scores'] == 10]['conv_num'].nunique()
            print(f"First Jailbreak: {self.query_to_jailbreak} Queries")
            print(f"Total Number of Conv. Jailbroken: {num_total_jailbreaks}/{self.batch_size} ({num_total_jailbreaks/self.batch_size*100:2.1f}%)")
            print(f"Example Jailbreak PROMPT:\n\n{self.jailbreak_prompt}\n\n")
            print(f"Example Jailbreak RESPONSE:\n\n{self.jailbreak_response}\n\n\n")
            
        else:
            print("No jailbreaks achieved.")
            max_score = df['judge_scores'].max()
            print(f"Max Score: {max_score}")


    

# class Saver:
#     """Saves the conversation."""

#     def __init__(self, args, system_prompt):
#         self.args = args
#         self.system_prompt = system_prompt

#         now = datetime.now(pytz.timezone('US/Eastern'))
#         self.filename = os.path.join(
#             "outputs",
#             f"{args.behavior}",
#             f"output_date_{now.month}_{now.day}_time_{now.hour}_{now.minute}.txt"
#         )

#     def write(self, conv):

#         with open(self.filename, 'w', encoding='utf-8') as f:
#             f.write(f"""
#                 Attack model: {self.args.attack_model_path}
#                 Target model: {self.args.target_model_path}\n
#                 System prompt: \n\n{self.system_prompt}\n\n"""
#             )

#             for counter, (role, s) in enumerate(conv.messages):
#                 if counter % 2 == 1:
#                     f.write(f"""\n{'='*36}
#                         Iteration: {(counter + 1) // 2}
#                         {'='*36}\n
#                         User:\n {s}\n"""
#                     )
#                 else:
#                     f.write(f"Assistant:\n {s}\n")