-
Notifications
You must be signed in to change notification settings - Fork 2
/
eval.py
216 lines (200 loc) · 12 KB
/
eval.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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
import os
import shutil
import pickle
import jsonlines
import torch
from tqdm import tqdm
from vendi_score import vendi
from consts import *
from utils import jload, Prompter, get_tokenizer, get_model
## GENERATE ANSWERS TO BENCH-QUESTIONS
def generate_answers_all_rds_all_tests(schedule, device, rd_start=0, rd_end=10, testsets=["vicuna", "koala", "wizardlm"],
generate_args={"do_sample": True, "temperature":0.7, "max_new_tokens":1024}, # "top_p":0.95,
):
"""all rds' resulting models -> generate answers on all testsets"""
rd_output_dir_root = f"evol_res/{schedule}/output"
for test in testsets:
print(f"*** TESTING: {test}======================================================================================================")
question_file = f"test_data/{test}_test_set.jsonl"
save_answer_root = f"evol_answer/{test}_test/{schedule}"
for i in range(rd_start, rd_end+1):
model_path = f"{rd_output_dir_root}/rd_{i}"
os.makedirs(save_answer_root, exist_ok=True)
save_answer_file = f"{save_answer_root}/rd_{i}.jsonl"
if os.path.exists(save_answer_file):
# load generated answers
with jsonlines.open(save_answer_file, "r") as f:
num_line = len([_ for _ in f])
should_num_line = TESTSETS_CONFIGS["should_num_line"]
if num_line == should_num_line: # all done for rd_model
print(f"*** Test -{test} ** Rd -{i} -> **Already Generated!")
continue
else: # start again
os.remove(save_answer_file)
generate_answer(model_path=model_path, device=device, save_answer_file=save_answer_file, question_file=question_file, generate_args=generate_args)
print(f"*** Test -{test} ** Rd -{i} -> **Generated! -> **Saved to: {save_answer_file}")
def generate_answer(model_path, device, question_file, generate_args, save_answer_file):
"""use model to generate answers to questions"""
# load tokenizer, model, prompter
prompter = Prompter(template_name="alpaca", verbose=False) # by default -> alpca-style prompt -> in line with training
try: # checkpoints save_model() might have error somehow
tokenizer, special_tokens_dict = get_tokenizer(model_name_or_path=model_path, cache_dir=None, model_max_length=512, )
model = get_model(model_name_or_path=model_path, cache_dir=None)
except: # use checkpoints from save_model_pretrained() in case
model_path_save_pretrained = f"{model_path}/pretrained"
print(f"model_path bad --> {model_path}")
print(f"switcing to model_path_save_pretrained --> {model_path_save_pretrained}")
# copy tokenizer configs
shutil.copy(f"{model_path}/added_tokens.json", f"{model_path_save_pretrained}/")
shutil.copy(f"{model_path}/special_tokens_map.json", f"{model_path_save_pretrained}/")
shutil.copy(f"{model_path}/tokenizer_config.json", f"{model_path_save_pretrained}/")
shutil.copy(f"{model_path}/tokenizer.model", f"{model_path_save_pretrained}/")
shutil.copy(f"{model_path}/trainer_state.json", f"{model_path_save_pretrained}/")
shutil.copy(f"{model_path}/training_args.bin", f"{model_path_save_pretrained}/")
tokenizer, special_tokens_dict = get_tokenizer(model_name_or_path=model_path_save_pretrained, cache_dir=None, model_max_length=512, ) # todo.
model = get_model(model_name_or_path=model_path_save_pretrained, cache_dir=None)
model = model.to(f"cuda:{device}")
# load questions
questions = []
with jsonlines.open(question_file, 'r') as f:
for line in f:
questions.append(line)
# generating answers ---
i_temp = 0
for q in tqdm(questions, total=len(questions)):
i_temp += 1
if "vicuna" in question_file:
instruction = q["text"]
elif "koala" in question_file:
instruction = q["prompt"]
elif "wizardlm" in question_file:
instruction = q["Instruction"]
prompt = prompter.generate_prompt(instruction=instruction)
# formatting in line with training inputs:
prompt_no_newline = prompt.strip()
# prompt_same_with_training = " " + prompt_no_newline.replace("\n\n", "\n")
prompt_same_with_training = " " + prompt_no_newline
# sanity check
if i_temp == 1:
print(f"Model Inference Input Text - same with training: {prompt_same_with_training}")
tokenized_input = tokenizer(prompt_same_with_training, return_tensors="pt", add_special_tokens=True)
generated = model.generate(tokenized_input["input_ids"].to(f"cuda:{device}"), **generate_args)
answer_with_prompt = tokenizer.decode(generated[0], skip_special_tokens=True)
answer = answer_with_prompt.split("### Response:")[1]
q["answer"] = answer
# save answers
with jsonlines.open(save_answer_file, "a") as f:
f.write(q)
## CALCULATE DATASET DIVERSITY (Vendi-Score)
def analyse_diversity_results_all_rds(schedule, rd_start, rd_end, device, measures, embed_model_path,):
# load embedding model & tokenizer
model = get_model(model_name_or_path=embed_model_path, cache_dir=None)
tokenizer, special_tokens_dict = get_tokenizer(model_name_or_path=embed_model_path, cache_dir=None, model_max_length=512, )
model.to(f"cuda:{device}")
print(f"Using Embedding-Model -> Loaded from: {embed_model_path}")
rd_to_vendi = {}
# iterate over round selected data
for rd in range(rd_start, rd_end+1):
rd_labeled_data_file = f"evol_res/{schedule}/data/rd_{rd}_labeled.json"
diversity_res = calculate_data_diversity_from_data_file(labeled_data_file=rd_labeled_data_file,
measures=measures,
model=model,
tokenizer=tokenizer,
)
rd_to_vendi[rd] = diversity_res["vendi"]
return {"vendi": rd_to_vendi,}
def calculate_data_diversity_from_data_file(labeled_data_file, measures, model, tokenizer, ):
# load data
print(f"Loading labeled data from: --> {labeled_data_file}")
data = jload(labeled_data_file)
# load prompter
prompter = Prompter(template_name="alpaca", verbose=False)
# compute embed
embeddings = []
for i in tqdm(range(len(data)), total=len(data)):
d = data[i]
# generate prompt -> including output as part of the training data
prompt = prompter.generate_prompt(
instruction=d["instruction"],
input=d["input"],
label=d["output"], # calculate diversity: including output as part of the training data
)
# formatting in line with training inputs:
prompt_no_newline = prompt.strip()
# prompt_same_with_training = " " + prompt_no_newline.replace("\n\n", "\n")
prompt_same_with_training = " " + prompt_no_newline
# infer -> get avg_pooled embeddings
with torch.no_grad():
tokenized_input = tokenizer(prompt_same_with_training, return_tensors="pt", add_special_tokens=True)
tokenized_input = {k:v.to(model.device) for k,v in tokenized_input.items()}
model.eval()
result = model(**tokenized_input, return_dict=True, output_hidden_states=True)
hidden_states = result["hidden_states"]
last_layer_hidden_states = hidden_states[-1] # (batch_size=1, seq_len, hidden_dim=4096)
last_hidden_states_avg = torch.mean(last_layer_hidden_states.squeeze(0), dim=0) # -> (hidden_dim)
embeddings.append(last_hidden_states_avg)
embeddings = torch.stack(embeddings)
print(f"Embeddings.shape = {embeddings.shape}")
# compute diversity measures
diversity_res = {}
for measure in measures:
if measure=="vendi": # vendi score
emb_norm = embeddings / torch.norm(embeddings, dim=-1, keepdim=True)
similarity = torch.mm(emb_norm, emb_norm.T)
sim_K = similarity.cpu().numpy()
vendi_diversity_score = vendi.score_K(sim_K)
diversity_res["vendi"] = vendi_diversity_score
return diversity_res
if __name__=="__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--eval_stage', type=str, required=True,
help="which evaluation function? {generate_answer, analyse_diversity}.")
parser.add_argument('--schedule', type=str, required=True,
help="result_dir_name. name of the folder containing all results from the DiverseEvol run.")
parser.add_argument('--device', type=int, default=0,
help="gpu-idx.")
parser.add_argument('--rd_start', type=int, default=0,
help="evaluation from which round's results.")
parser.add_argument('--rd_end', type=int, default=10,
help="evaluation till which round's results.")
parser.add_argument('--testsets', nargs='+', type=str, default=["vicuna", "koala", "wizardlm"],
help="generate answers to which testsets?")
parser.add_argument('--embed_model_path', type=str, default=None,
help="path to the foundation model. its embedding is used to calculated vendi-score.")
args = parser.parse_args()
print(args)
if args.eval_stage == "generate_answer":
generate_answers_all_rds_all_tests(schedule=args.schedule,
device=args.device,
rd_start=args.rd_start,
rd_end=args.rd_end,
testsets=args.testsets,
generate_args={"do_sample": True, "temperature":0.7, "max_new_tokens":1024}) # "top_p":0.95
elif args.eval_stage == "analyse_diversity":
schedule_to_rd_diversity_res = {}
print(f"** calculating diversity of -> {args.schedule} ...")
if "baseline" in args.schedule:
rd_end_temp = 0
measures = ["vendi"]
else:
rd_end_temp = args.rd_end
measures = ["vendi"]
save_diversity_results_file = f"evol_diversity/{args.schedule}_RD={args.rd_start}-{rd_end_temp}_MEASURES={'_'.join(measures)}.pkl" # path to saving calculated diversity results
if os.path.exists(save_diversity_results_file):
with open(save_diversity_results_file, "rb") as f:
diversity_results = pickle.load(f)
print(f"Loading diversity_results from: <-- {save_diversity_results_file}")
else:
diversity_results = analyse_diversity_results_all_rds(schedule=args.schedule,
rd_start=args.rd_start,
rd_end=rd_end_temp,
device=args.device,
measures=measures,
embed_model_path=args.embed_model_path, # always use the same foundation model
)
with open(save_diversity_results_file, "wb") as f:
pickle.dump(diversity_results, f)
print(f"Saving diversity_results to: --> {save_diversity_results_file}") # saving diversity_results
schedule_to_rd_diversity_res[args.schedule] = diversity_results
print(f"** schedule: {args.schedule} - rd_to_diversity_res --> {diversity_results}") # printing diversity_results