Views
No views yet
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
rm_tokenizer = AutoTokenizer.from_pretrained('Ray2333/gpt2-large-harmless-reward_model')
reward_model = AutoModelForSequenceClassification.from_pretrained(
'Ray2333/gpt2-large-harmless-reward_model',
num_labels=1, torch_dtype=torch.bfloat16,
device_map=0,
)
q, a = "\n\nHuman: I just came out of from jail, any suggestion of my future? \n\nAssistant:", "Go back to jail you scum"
inputs = rm_tokenizer(q, a, return_tensors='pt', truncation=True)
with torch.no_grad():
reward = reward_model(**(inputs.to(0))).logits[0].cpu().detach().item()@article{yang2024rewards,
title={Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment},
author={Yang, Rui and Pan, Xiaoman and Luo, Feng and Qiu, Shuang and Zhong, Han and Yu, Dong and Chen, Jianshu},
journal={International Conference on Machine Learning},
year={2024}
}