Views
No views yet

1import vllm
2
3
4def apply_qwen_math_template(question: str):
5 return (
6 "<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n<|im_start|>user\n"
7 + question
8 + "<|im_end|>\n<|im_start|>assistant\n"
9 )
10
11def apply_r1_template(question: str):
12 return (
13 "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. "
14 "The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.\nUser: "
15 + question
16 + "\nAssistant: <think>"
17 )
18
19model_name = "sail/Qwen2.5-Math-1.5B-Oat-Zero"
20
21sampling_params = vllm.SamplingParams(
22 n=1,
23 temperature=0,
24 top_p=1,
25 max_tokens=3000,
26)
27
28model = vllm.LLM(
29 model_name,
30 max_model_len=4096,
31 dtype="bfloat16",
32 enable_prefix_caching=True,
33)
34
35if "Llama-3.2-3B-Oat-Zero" in model_name:
36 apply_template = apply_r1_template
37else:
38 apply_template = apply_qwen_math_template
39
40prompts = [
41 "How many positive whole-number divisors does 196 have?"
42]
43prompts = list(map(apply_template, prompts))
44outputs = model.generate(prompts, sampling_params)
45
46print(outputs)1@article{liu2025understanding,
2 title={Understanding r1-zero-like training: A critical perspective},
3 author={Liu, Zichen and Chen, Changyu and Li, Wenjun and Qi, Penghui and Pang, Tianyu and Du, Chao and Lee, Wee Sun and Lin, Min},
4 journal={arXiv preprint arXiv:2503.20783},
5 year={2025}
6}