RLAIF

rlaif.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained('NousResearch/Nous-Hermes-Llama2-13b', device_map = 'auto')
tokenizer = AutoTokenizer.from_pretrained('NousResearch/Nous-Hermes-Llama2-13b')
model.eval()

print(tokenizer('yes')) # [1, 4871]
print(tokenizer.decode(4874)) # yes

print(tokenizer('no')) # [1, 694]
print(tokenizer.decode(694)) # no


question = 'Do pineapples belong on pizza? \nAnswer:'


outputs = model.generate(tokenizer([question], return_tensors = "pt").input_ids.to('cuda'),
                         return_dict_in_generate = True, output_scores = True, max_new_tokens = 1)


p_yes = torch.exp(outputs.scores[0][:, 4874]).cpu().numpy()[0]
p_no = torch.exp(outputs.scores[0][:, 694]).cpu().numpy()[0]

print('Prob yes:', p_yes) # p_yes: 490.5 (the probs are not normalized)
print('Prob no:', p_no) # p_no: 424.5

reward = (p_yes / (p_yes + p_no))
print('Reward:', reward) # reward: 0.536