Qwen3 Nemotron 32B GenRM Principle

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "nvidia/Qwen3-Nemotron-32B-GenRM-Principle"

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "What is 1+1?"
good_response = "1+1=2"
bad_response = "1+1=3"
principle = "correctness"

for response in [good_response, bad_response]:
    messages = [{'role': "user", "content": prompt}, {'role': "assistant", "content": response}, {'role': "principle", "content": principle}]
    tokenized_message = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
    response_token_ids = model.generate(tokenized_message['input_ids'].cuda(),attention_mask=tokenized_message['attention_mask'].cuda(),  max_new_tokens=16000, return_dict_in_generate=True, output_scores=True)
    response = tokenizer.decode(response_token_ids.sequences[0].tolist())
    score_max = torch.max(response_token_ids.scores[-2][0]).item() # normalize max score to zero to match vLLM logprobs and clip others to -50 if they return -inf for too small of a value. The same should be done if the required tokens is not returned by vLLM's top k logprobs.
    score_no = max(response_token_ids.scores[-2][0][2308].item() - score_max, -50)  # token for " No"
    score_yes = max(response_token_ids.scores[-2][0][7414].item() - score_max, -50) # token for " Yes"
    reward = score_yes-score_no
    print(response)
    print(reward)