RiC / plot.py
hour1's picture
Upload folder using huggingface_hub
f5d1134 verified
Raw
History Blame Contribute Delete
3.95 kB
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
import glob2
import seaborn as sns
def find_pareto_points(obtained_scores, threshold=0.02):
n = len(obtained_scores)
if n == 1:
return obtained_scores
pareto_index = []
high_low = np.max(obtained_scores, axis=0) - np.min(obtained_scores, axis=0)
for i in range(n):
if not any(np.all((obtained_scores - obtained_scores[i] - threshold * high_low) > 0.0, axis=1)):
pareto_index.append(i)
points = obtained_scores[np.array(pareto_index)]
arg_index = np.argsort(points[:, 0])
points = points[arg_index]
print(points)
sorted_index = [0]
remaining_index = np.ones(len(points))
i = 0
remaining_index[i] = 0
while sum(remaining_index):
distance = ((points[np.where(remaining_index)] - points[i]) ** 2 ).sum(axis=1)
min_index = np.where(remaining_index > 0)[0][np.argmin(distance)]
sorted_index.append(min_index)
i = min_index
remaining_index[i] = 0
return points[np.array(sorted_index)]
index = 1
colors = sns.color_palette('Paired')
def plot_points(dir, label, style='-*', color='b', shift=[0,0], txt_color='black', normalize_path=None, reverse=True):
threshold = 0.01
desired_scores = []
obtained_scores = []
paths = [os.path.abspath(path) for path in glob2.glob(os.path.join(dir, '*.csv'))]
paths += [os.path.abspath(path) for path in glob2.glob(os.path.join(dir, '*', '*.csv'))]
pref_lis = []
for path in paths:
if '.csv' in path:
full_path = path
data = pd.read_csv(full_path)
# morlhf has less points, let the threshold larger to make the frontier better
if 'ppo' in path and len(paths) <= 5:
threshold = 0.5
obtained_scores.append([np.mean(data['obtained_score1']), np.mean(data['obtained_score2'])])
if 'pref' in path:
# get the preference
if 'eval_data_pref' in path:
pref = path.split('eval_data_pref')[-1].strip().split('_')[0]
pref_lis.append(float(pref))
print(pref_lis)
desired_scores = np.array(desired_scores)
obtained_scores = np.array(obtained_scores)
if normalize_path is not None:
norm_info = np.load(normalize_path)
norm_info = np.array(norm_info).reshape(2, 2)
for i in range(2):
obtained_scores[:, i] = (obtained_scores[:, i] - norm_info[i][0]) / norm_info[i][1]
global index
markersize = 10 if ('*' in style or 'o' in style) else 9
pareto_points = find_pareto_points(obtained_scores, threshold)
plt.scatter(obtained_scores[:, 0], obtained_scores[:, 1], marker=style[-1], color=colors[index], s=markersize + 60)
if len(pref_lis):
for i in range(len(obtained_scores)):
plt.annotate('{}'.format(round(pref_lis[i], 1)), (obtained_scores[i, 0] + shift[0], obtained_scores[i, 1] + shift[1]), size=4, color=txt_color)
plt.plot(pareto_points[:, 0], pareto_points[:, 1], style, c=colors[index], markersize=markersize, label=label)
index += 2
plt.figure(figsize=(5, 4))
name1 = 'harmless'
name2 = 'helpful'
### replace the paths to your own paths
plot_points('./logs_trl/eval_pretrained', 'Llama 2 base', '*')
plot_points('./logs_trl/eval_sft_alldata', 'SFT', '*')
plot_points('./eval_ppo_pref/', 'MORLHF', '--D', shift=[-0.012, -0.022])
plot_points('./logs_ppo/eval_pposoups_llamma2_klreg0.2', 'Rewarded Soups', style='--s', shift=[-0.012, -0.022])
plot_points('.logs_trl/evalnew_onlinefix_helpful_harmlesshelpful_iter2', 'RiC', style='-o', shift=[-0.012, -0.022], txt_color='white')
plt.xlabel('$R_1$ ({})'.format(name1), fontsize=12)
plt.ylabel('$R_2$ ({})'.format(name2), fontsize=12)
plt.legend(fontsize=11, loc='lower left')
plt.tight_layout()
plt.savefig('ric_assistant_{}_{}.pdf'.format(name1, name2))