-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate_reward_from_ultraFeedback.py
177 lines (128 loc) · 5.94 KB
/
generate_reward_from_ultraFeedback.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
from transformers import PreTrainedModel, LlamaConfig, LlamaModel, LlamaTokenizer
import torch.nn as nn
import torch
from typing import Optional, List
import jsonlines
import json
import numpy as np
import argparse
from tqdm import tqdm
import os
ultrarm_template= "Human: {instruction}\nAssistant: {response}"
def save_data_to_json(strings, file_name):
with open(file_name, 'w', encoding='utf-8') as file:
writer = jsonlines.Writer(file)
for string in strings:
writer.write(string)
class LlamaRewardModel(PreTrainedModel):
config_class = LlamaConfig
def __init__(self, config):
super().__init__(config)
self.model = LlamaModel(config)
self.regression_head = nn.Linear(self.config.hidden_size, 1, bias=False)
def forward( # args are the same as LlamaForCausalLM
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
):
transformer_outputs = self.model(
input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
)
hidden_states = transformer_outputs[0]
rewards = self.regression_head(hidden_states).squeeze(-1)
ends = attention_mask.cumsum(dim=1).argmax(dim=1).view(-1,1)
rewards = torch.gather(rewards, 1, ends)
return rewards
def main():
parser = argparse.ArgumentParser(prog='', description='')
parser.add_argument('-i,', '--input', required=True, help="folder to model outputs")
parser.add_argument('--output', required=False, help="folder to model outputs")
parser.add_argument('-m,', '--model', required=False, help="folder to model")
parser.add_argument('-w,', '--weak', required=False, help="")
parser.add_argument('--id_to_remove', default='path', required=False, help="")
parser.add_argument('--path_user_preference', type=str, required=False, help="")
args = parser.parse_args()
root_user_preference = ''
if args.path_user_preference and os.path.exists(args.path_user_preference):
with open(args.path_user_preference, 'r') as f:
root_user_preference = f.read().strip()
## load the model
tokenizer = LlamaTokenizer.from_pretrained(args.model)
model = LlamaRewardModel.from_pretrained(args.model)
model.half()
model.to('cuda')
model.eval()
def evaluate(instructions, responses):
prompts = [ultrarm_template.format(instruction = instruction, response = response) for instruction, response in zip(instructions, responses)]
print(prompts[0])
print('\n\n')
all_rewards = []
#for prompt in tqdm(prompts):
for prompt in prompts:
try:
inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
reward = model(**inputs).item()
all_rewards.append(reward)
except Exception as e:
print(f"{e}")
all_rewards.append(0.0)
return all_rewards
max_rewards, mid_rewards = [], []
data_to_save = []
with open(args.input, 'r') as f:
content = f.readlines()
for item in tqdm(content):
item = json.loads(item)
if 'responses' in item:
if root_user_preference == 'preference_2':
instruction = item['instruction'] + '\n\n' + item['preference_2']
elif root_user_preference == 'preference_1':
instruction = item['instruction'] + '\n\n' + item['preference_1']
else:
instruction = item['instruction']
instruction = instruction.strip()
responses = item[f'responses']
instructions = [instruction] * len(responses)
rewards = evaluate(instructions, responses)
print(rewards)
print('-------------------\n\n')
max_rewards.append(max(rewards))
mid_rewards.append(np.median(rewards))
best_response = responses[np.argmax(rewards)]
item['scores'] = rewards
item['score'] = max(rewards)
item['best_response'] = best_response
else:
if root_user_preference == 'preference_2':
instruction = item['instruction'] + '\n\n' + item['preference_2']
elif root_user_preference == 'preference_1':
instruction = item['instruction'] + '\n\n' + item['preference_1']
else:
instruction = item['instruction'] #+ '\n\n' + root_user_preference
instruction = instruction.strip()
response = item['response']
rewards = evaluate([instruction], [response])
print(rewards)
print('-------------------\n\n')
max_rewards.append(max(rewards))
mid_rewards.append(np.median(rewards))
item['scores'] = rewards
item['score'] = max(rewards)
data_to_save.append(item)
print(f'Ave max rewards: {np.mean(max_rewards)}')
print(f'Ave median rewards: {np.mean(mid_rewards)}')
save_data_to_json(data_to_save, args.output)
if __name__ == '__main__':
main()