Views
No views yet
1import math
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5# tokenizer
6tokenizer = AutoTokenizer.from_pretrained('dongboklee/gORM-14B')
7yes_id = tokenizer.encode(" Yes", add_special_tokens=False)[-1]
8no_id = tokenizer.encode(" No", add_special_tokens=False)[-1]
9
10# model
11device = 'cuda' if torch.cuda.is_available() else 'cpu'
12model = AutoModelForCausalLM.from_pretrained('dongboklee/gORM-14B')
13model.eval()
14model.to(device)
15
16# prompt formatting
17question = 'Question: In Python 3, which of the following function convert a string to an int in python?\nA. short(x)\nB. float(x)\nC. integer(x [,base])\nD. double(x)\nE. int(x [,base])\nF. long(x [,base] )\nG. num(x)\nH. str(x)\nI. char(x)\nJ. digit(x [,base])'
18solution = ["To convert a string to an integer in Python 3, we use the built-in function int().",
19 "The int() function takes two arguments: the string to be converted and an optional base (default is 10, which is for decimal).",
20 "For example: int(\"123\", 10) converts the string \"123\" to the integer 123.",
21 "Looking at the options, we can see that the correct function is option E: int(x [,base]).",
22 "The answer is (E)."]
23category_name = "computer science"
24prefix = "\n\n".join(solution)
25
26# Create the prompt
27prompt_text = (
28 f"You are a {category_name} teacher. Grade the solution, verifying correctness step by step.\n"
29 "At the end of Solution verification, when you give your final grade, write it in the form \"Verification: Is the answer correct (Yes/No)? X\", where X is either Yes or No.\n\n"
30 f"[{category_name.capitalize()} Problem]\n{question.strip()}\n\n"
31 f"[Solution]\n{prefix.strip()}\n"
32)
33
34prompt = tokenizer.apply_chat_template(
35 [{'role': "user", "content": prompt_text}],
36 tokenize=False, add_generation_prompt=True, add_special_tokens=False
37) + "Let's verify step by step:"
38
39# Tokenize the prompt
40inputs = tokenizer(prompt, return_tensors="pt").to(device)
41
42# generate
43with torch.no_grad():
44 outputs = model.generate(
45 **inputs,
46 max_new_tokens=8192,
47 return_dict_in_generate=True,
48 output_scores=True,
49 pad_token_id=tokenizer.eos_token_id
50 )
51
52# compute reward
53logits = outputs.logits[0, -2, :]
54yes_logit, no_logit = logits[yes_id].item(), logits[no_id].item()
55reward = math.exp(yes_logit) / (math.exp(yes_logit) + math.exp(no_logit))@article{multi-rm,
title = {Rethinking Reward Models for Multi-Domain Test-Time Scaling},
author = {Lee, Dong Bok and Lee, Seanie and Park, Sangwoo and Kang, Minki and Baek, Jinheon and Kim, Dongki and Wagner, Dominik and Jin, Jiongdao and Lee, Heejun and Bocklet, Tobias and Wang, Jinyu and Fu, Jingjing and Hwang, Sung Ju and Bian, Jiang and Song, Lei},
journal = {arXiv preprint arXiv:2510.00492},
year = {2025}
}