Views
No views yet
1from transformers import AutoModelForSequenceClassification,AutoTokenizer
2import torch
3device = 'cuda'
4path = "RLHFlow/RewardModel-Mistral-7B-for-DPA-v1"
5rm = AutoModelForSequenceClassification.from_pretrained(path, trust_remote_code=True).to(device)
6tokenizer = AutoTokenizer.from_pretrained(path)
7
8input_template = "[INST] You must read the following conversation carefully and rate the assistant's response from score 0-100 in these aspects: helpfulness, correctness, coherence, honesty, complexity, verbosity\n\nUser: {prompt}\n\nAssistant: {response} [/INST]"
9
10# Use a sample from HelpSteer validation set
11prompt = 'What are some synonyms for the word "beautiful"?'
12response = "Nicely, Beautifully, Handsome, Stunning, Wonderful, Gorgeous, Pretty, Stunning, Elegant"
13
14model_inputs = tokenizer(input_template.format(prompt=prompt, response=response), return_tensors="pt").to(device)
15with torch.no_grad():
16 score = rm(**model_inputs).logits.squeeze().cpu().float().numpy()
17
18print(score)
19# [68.99269 69.62718 76.23071 33.48785 35.853596 63.833366 55.58917 68.7175 59.552124 46.465595]
20
21# Convert from our scale (0-100) to HelpSteer scale (0-4)
22helpsteer_rewards_pred = (score[:5]-10)/20
23print(helpsteer_rewards_pred)
24# [2.9496346 2.981359 3.3115356 1.1743925 1.2926798]
25# The actual rewards from the HelpSteer dataset for this sample are [3,3,4,2,2]

@inproceedings{wang2024arithmetic,
title={Arithmetic Control of LLMs for Diverse User Preferences: Directional Preference Alignment with Multi-Objective Rewards},
author={Haoxiang Wang and Yong Lin and Wei Xiong and Rui Yang and Shizhe Diao and Shuang Qiu and Han Zhao and Tong Zhang},
year={2024},
booktitle={ACL},
}