File size: 470 Bytes
c42ab87 |
1 2 3 4 5 6 7 8 9 10 11 12 |
from transformers import AutoTokenizer
from safe_rlhf.models import AutoModelForScore
model = AutoModelForScore.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', device_map='auto')
tokenizer = AutoTokenizer.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', use_fast=False)
input = 'BEGINNING OF CONVERSATION: USER: hello ASSISTANT:Hello! How can I help you today?'
input_ids = tokenizer(input, return_tensors='pt')
output = model(**input_ids)
print(output)
|