Views
No views yet
1from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
2
3def safe_execute(code_string: str, keys=None):
4 def execute(x):
5 try:
6 exec(x)
7 locals_ = locals()
8 if keys is None:
9 return locals_.get('ans', None)
10 else:
11 return [locals_.get(k, None) for k in keys]
12 except Exception:
13 return None
14 try:
15 ans = func_timeout.func_timeout(5, execute, args=(code_string,))
16 except func_timeout.FunctionTimedOut:
17 ans = None
18 return ans
19
20checkpoint = "zhuxunyu/math-codet5p-770m-py"
21device = "cuda" # for GPU usage or "cpu" for CPU usage
22
23tokenizer = AutoTokenizer.from_pretrained(checkpoint)
24model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint).to(device)
25
26question = "Question: Janet\u2019s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?\nProgram: Let’s design executable python program (return ans) to solve the question.".
27input = tokenizer(question, max_length=256, padding="max_length", truncation=True, return_tensors="pt").to(model.device)
28
29with torch.no_grad():
30 output = model.generate(**input, max_length=256)
31
32generation = tokenizer.decode(output, skip_special_tokens=True)
33ans = safe_execute(generation)
34print(float(ans))@misc{zhu2023mathcodet5plus,
title={math-codet5p-770m-py},
author={Xunyu Zhu, Jian Li, Yong Liu, Can Ma, Weiping Wang},
year={2023}
}