Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4def get_tokenizer(model_id):
5 tokenizer = AutoTokenizer.from_pretrained(model_id)
6 tokenizer.pad_token = tokenizer.eos_token
7 tokenizer.padding_side = 'left'
8 tokenizer.truncation_side = 'left'
9 return tokenizer
10
11device = 'cuda' if torch.cuda.is_available() else 'cpu'
12tokenizer = get_tokenizer('UW-Madison-Lee-Lab/VersaPRM-Aug')
13model = AutoModelForCausalLM.from_pretrained('UW-Madison-Lee-Lab/VersaPRM-Aug')
14candidate_tokens = [12, 10]
15model.to(device)
16
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)."]
23input_text = question + ' \n\n' + ' \n\n\n\n'.join(solution) + ' \n\n\n\n' # solution steps are separated by ' \n\n\n\n'
24input_id = torch.tensor([tokenizer.encode(input_text)]).to(device)
25
26with torch.no_grad():
27 logits = model(input_id).logits[:,:,candidate_tokens]
28 scores = logits.softmax(dim=-1)[:,:,1]
29 step_scores = scores[input_id == 23535]
30 step_probs = step_scores.tolist()