import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "cyberagent/ca-reward-3b-ja",
    device_map="auto",
    num_labels=1,
)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/ca-reward-3b-ja")

prompt = """手軽に栄養を補給できる食事を教えてください。"""
response1 = """栄養補給が手軽にできるものとしては、野菜たっぷりのスムージー、ゆで卵とサラダ、納豆ご飯などがおすすめです。特に納豆は手間なく良質なタンパク質が摂れますよ。お身体を大切にしてくださいね。"""
response2 = """手軽な栄養補給なら冷凍食品でいいと思います。レンジで温めるだけだし、時間がない時はコンビニ弁当でも悪くないですよ。"""

chat1 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response1}]
chat2 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response2}]

chat1_formatted = tokenizer.apply_chat_template(chat1, tokenize=False)
chat2_formatted = tokenizer.apply_chat_template(chat2, tokenize=False)
chat1_tokenized = tokenizer(chat1_formatted, return_tensors="pt", max_length=4096, truncation=True, padding="max_length",).to(model.device)
chat2_tokenized = tokenizer(chat2_formatted, return_tensors="pt", max_length=4096, truncation=True, padding="max_length",).to(model.device)

with torch.no_grad():
    chat1_score = model(**chat1_tokenized).logits.item()
    chat2_score = model(**chat2_tokenized).logits.item()

print(f"Score for response 1: {chat1_score}")
print(f"Score for response 2: {chat2_score}")

# Score for response 1: 1.2595189809799194
# Score for response 2: -2.917454242706299

@misc{cyberagent-ca-reward-3b-ja,
      title={cyberagent/ca-reward-3b-ja},
      url={https://huggingface.co/cyberagent/ca-reward-3b-ja},
      author={Ryota Mitsuhashi},
      year={2025},
}