-
Notifications
You must be signed in to change notification settings - Fork 70
/
Copy pathloggers.py
124 lines (103 loc) · 4.96 KB
/
loggers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import wandb
import pandas as pd
import logging
def setup_logger():
logger = logging.getLogger('PAIR')
handler = logging.StreamHandler()
logger.addHandler(handler)
return logger
def set_logger_level(logger, verbosity):
if verbosity == 0:
level=logging.CRITICAL # Disables logging
elif verbosity == 1:
level = level=logging.INFO
else:
level = logging.DEBUG
logger.setLevel(level)
for handler in logger.handlers:
handler.setLevel(level)
logger = setup_logger()
logger.set_level = lambda verbosity : set_logger_level(logger, verbosity)
class WandBLogger:
"""WandB logger."""
def __init__(self, args, system_prompts):
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_prompts,
"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])
logger.info(f"{'='*14} SUMMARY STATISTICS for Iteration {iter} {'='*14}")
logger.info(f"Mean/Max Score for iteration: {mean_score_for_iter:.1f}, {max_score_for_iter}")
logger.info(f"Number of New Jailbreaks: {num_new_jailbreaks}/{bs}")
logger.info(f"Total Number of Conv. Jailbroken: {num_total_jailbreaks}/{bs} ({num_total_jailbreaks/bs*100:2.1f}%)\n")
def print_final_summary_stats(self):
logger.info(f"{'='*8} FINAL SUMMARY STATISTICS {'='*8}")
logger.info(f"Index: {self.index}")
logger.info(f"Goal: {self.goal}")
df = self.table
if self.is_jailbroken:
num_total_jailbreaks = df[df['judge_scores'] == 10]['conv_num'].nunique()
logger.info(f"First Jailbreak: {self.query_to_jailbreak} Queries")
logger.info(f"Total Number of Conv. Jailbroken: {num_total_jailbreaks}/{self.batch_size} ({num_total_jailbreaks/self.batch_size*100:2.1f}%)")
logger.info(f"Example Jailbreak PROMPT:\n\n{self.jailbreak_prompt}\n\n")
logger.info(f"Example Jailbreak RESPONSE:\n\n{self.jailbreak_response}\n\n\n")
else:
logger.info("No jailbreaks achieved.")
max_score = df['judge_scores'].max()
logger.info(f"Max Score: {max_score}")