Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4# tokenizer
5def get_tokenizer(model_id):
6 tokenizer = AutoTokenizer.from_pretrained(model_id)
7 tokenizer.pad_token = tokenizer.eos_token
8 tokenizer.padding_side = 'left'
9 tokenizer.truncation_side = 'left'
10 return tokenizer
11
12tokenizer = get_tokenizer('dongboklee/dORM-14B')
13candidate_tokens = [
14 self.tokenizer.encode("-", add_special_tokens=False)[-1],
15 self.tokenizer.encode("+", add_special_tokens=False)[-1]
16]
17tag_id = self.tokenizer.encode(" \n\n\n\n", add_special_tokens=False)[-1]
18
19# model
20device = 'cuda' if torch.cuda.is_available() else 'cpu'
21model = AutoModelForCausalLM.from_pretrained('dongboklee/dORM-14B')
22model.eval()
23model.to(device)
24
25question = '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])'
26solution = ["To convert a string to an integer in Python 3, we use the built-in function int().",
27 "The int() function takes two arguments: the string to be converted and an optional base (default is 10, which is for decimal).",
28 "For example: int(\"123\", 10) converts the string \"123\" to the integer 123.",
29 "Looking at the options, we can see that the correct function is option E: int(x [,base]).",
30 "The answer is (E)."]
31input_text = question + ' \n\n' + ' \n\n\n\n'.join(solution) + ' \n\n\n\n' # solution steps are separated by ' \n\n\n\n'
32input_id = torch.tensor([tokenizer.encode(input_text)]).to(device)
33
34with torch.no_grad():
35 logits = model(input_id).logits[:,:,candidate_tokens]
36 scores = logits.softmax(dim=-1)[:,:,1]
37 step_scores = scores[input_id == tag_id]
38 step_probs = step_scores.tolist()[:,:,-1]@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}
}