-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathapproach-temperature.py
188 lines (166 loc) · 8.11 KB
/
approach-temperature.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
import json
import argparse
import lm_utils
import metrics
import random
import torch.nn
from tqdm import tqdm
if __name__ == "__main__":
argParser = argparse.ArgumentParser()
argParser.add_argument("-m", "--model", help="which language model to use: \"mistral\", \"llama2_7/13/70b\", \"chatgpt\"")
argParser.add_argument("-d", "--dataset", help="which dataset in data/: \"mmlu\", \"knowledge_crosswords\", \"hellaswag\", \"propaganda\", \"ambigqa\", \"electionqa23\"")
argParser.add_argument("-o", "--portion", default = 1.0, help="portion of the dataset to use")
args = argParser.parse_args()
model_name = args.model
dataset = args.dataset
portion = args.portion
lm_utils.llm_init(model_name)
correct_flags = []
abstain_flags = []
abstain_scores = []
with open("data/" + dataset + ".json", "r") as f:
data = json.load(f)
data["dev"] = data["dev"][:int(len(data["dev"])*float(portion))]
data["test"] = data["test"][:int(len(data["test"])*float(portion))]
# obtain correct flags for test set
for d in tqdm(data["test"]):
original_prompt = "Question: " + d["question"] + "\n"
for key in d["choices"].keys():
original_prompt += (key + ": " + d["choices"][key] + "\n")
original_prompt += "Choose one answer from the above choices. The answer is"
response = lm_utils.llm_response(original_prompt, model_name, probs=False)
# print(response)
# print(lm_utils.answer_parsing(response))
if lm_utils.answer_parsing(response) == d["answer"]:
correct_flags.append(1)
else:
correct_flags.append(0)
# temperature scaling on validation set token probabilities, determine optimal tau value
correct_labels_dev = []
option_to_ind = {"A": 0, "B": 1, "C": 2, "D": 3}
probs = []
target = []
for d in tqdm(data["dev"]):
correct_answer = d["answer"]
target.append(option_to_ind[correct_answer])
original_prompt = "Question: " + d["question"] + "\n"
for key in d["choices"].keys():
original_prompt += (key + ": " + d["choices"][key] + "\n")
original_prompt += "Choose one answer from the above choices. The answer is"
response, token_probs = lm_utils.llm_response(original_prompt, model_name, probs=True)
# print(response, token_probs)
# print("------------------")
if lm_utils.answer_parsing(response) == d["answer"]:
correct_labels_dev.append(1)
else:
correct_labels_dev.append(0)
prob_max = None
chosen_option = None
try:
for token in token_probs.keys():
if token.strip() in option_to_ind.keys():
prob_max = token_probs[token]
chosen_option = token.strip()
break
except:
print("option probs failed, uniform assignment")
chosen_option = random.choice(["A", "B", "C", "D"])
prob_max = 0.25
if chosen_option == None:
print("option probs failed, uniform assignment")
chosen_option = random.choice(["A", "B", "C", "D"])
prob_max = 0.25
prob_distribution = [0, 0, 0, 0]
prob_distribution[option_to_ind[chosen_option]] = prob_max
# evenly split between other options, since we only care about the most likely option / the one generated in greedy decoding
for i in range(4):
if i != option_to_ind[chosen_option]:
prob_distribution[i] = (1 - prob_max) / 3
probs.append(prob_distribution)
min_error = 1e10
best_tau = None
loss_fn = torch.nn.CrossEntropyLoss()
# tau from 0.1 to 10 interval of 0.01
for tau in range(1, 1001):
tau = tau / 100.0
new_probs = []
for i in range(len(probs)):
probs_temp = torch.tensor([num / tau for num in probs[i]])
#new_probs.append(torch.nn.functional.softmax(probs_temp, dim=0))
new_probs.append(probs_temp)
new_probs = torch.stack(new_probs)
error = loss_fn(new_probs, torch.tensor(target)).item() # CE loss on temperature scaled token probablities
if error < min_error:
min_error = error
best_tau = tau
# print("best tau:", best_tau)
# print("best error:", min_error)
for i in range(len(probs)):
probs[i] = torch.nn.functional.softmax(torch.tensor([num / best_tau for num in probs[i]]), dim=0).tolist() # apply the best tau scaling
# determine optimal threshold for abstain, with the scaled token probabilities
prob_maximum = max([max(prob) for prob in probs])
prob_minimum = min([max(prob) for prob in probs])
min_error = 1e6
best_threshold = 0
for threshold in range(1, 100):
# no 100% abstain or 100% answer
if threshold / 100.0 >= prob_maximum or threshold / 100.0 <= prob_minimum:
continue
error = 0
for i in range(len(correct_labels_dev)):
if max(probs[i]) < float(threshold/100.0):
if correct_labels_dev[i] == 1:
error += 1
else:
if correct_labels_dev[i] == 0:
error += 1
if error < min_error:
min_error = error
best_threshold = float(threshold/100.0)
# print("best threshold:", best_threshold)
# print("best error:", min_error)
# obtain abstain flags for test set
for d in tqdm(data["test"]):
original_prompt = "Question: " + d["question"] + "\n"
for key in d["choices"].keys():
original_prompt += (key + ": " + d["choices"][key] + "\n")
original_prompt += "Choose one answer from the above choices. The answer is"
response, token_probs = lm_utils.llm_response(original_prompt, model_name, probs=True)
prob_max = None
chosen_option = None
try:
for token in token_probs.keys():
if token.strip() in option_to_ind.keys():
prob_max = token_probs[token]
chosen_option = token.strip()
break
except:
print("option probs failed, uniform assignment")
chosen_option = random.choice(["A", "B", "C", "D"])
prob_max = 0.25
if chosen_option == None:
print("option probs failed, uniform assignment")
chosen_option = random.choice(["A", "B", "C", "D"])
prob_max = 0.25
prob_distribution = [0, 0, 0, 0]
prob_distribution[option_to_ind[chosen_option]] = prob_max
# evenly split between other options, since we only care about the most likely option / the one generated in greedy decoding
for i in range(4):
if i != option_to_ind[chosen_option]:
prob_distribution[i] = (1 - prob_max) / 3
prob_distribution = [num / best_tau for num in prob_distribution]
prob_distribution = torch.nn.functional.softmax(torch.tensor(prob_distribution), dim=0).tolist()
if prob_distribution[option_to_ind[chosen_option]] < best_threshold:
abstain_flags.append(1)
else:
abstain_flags.append(0)
abstain_scores.append(1-prob_distribution[option_to_ind[chosen_option]])
# print(correct_flags)
# print(abstain_flags)
# print(abstain_scores)
print("------------------")
print("Approach: temperature")
print("Model:", model_name)
print("Dataset:", dataset)
print(metrics.compute_metrics(correct_flags, abstain_flags, abstain_scores))
print("------------------")