Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_name = "PKU-ONELab/CE-RM-4B"
5model = AutoModelForCausalLM.from_pretrained(
6 model_name,
7 torch_dtype="auto",
8 device_map="auto"
9)
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12criteria_prompt = """Your task is to produce a minimal set of criteria for evaluating the quality of potential responses to the user query given below.
13
14Begin by carefully analyzing the query to fully understand the user's intent and requirements, and then take into account all common and tangible factors that can indicate the response quality.
15
16From these considerations, derive the final evaluation criteria list, which **must adhere to the following requirements:**
17
18- Each criterion should consist of a concise term as well as its unambiguous description.
19- The number of criteria is not necessarily the more the better; Fewer yet comprehensive is more desired.
20- The criteria should be sufficient and complete, ensuring that no essential aspects or key signals of response quality are omitted.
21- The criteria should be necessary and non-overlapping, with each one indispensable, distinct in perspective, and strictly orthogonal to others.
22
23Provide the relevant analysis first, followed by the numbered list of criteria between [Start of Criteria] and [End of Criteria], with one criterion per line and the more important ones coming first.
24
25Below is the user query:
26
27[Start of Query]
28{query}
29[End of Query]
30"""
31
32evaluation_prompt = """Now that you have a response to the previous user query, your new task is to evaluate it using the criteria list you have produced.
33
34For each criterion, focus on its concerns and carefully evaluate the corresponding specific quality of the response, providing the detailed analysis as well as relevant arguments, followed by the corresponding quality score from 0 to 5 within $\\boxed{}$.
35
36Moreover, if the response demonstrates strengths or weaknesses beyond the scope of your criteria list, introduce an additional criterion titled \"Other Point(s),\" discussing them and considering them as bonus points or deductions as appropriate.
37
38Finally, based on the analyses of these criteria, including their relative importance and scores, **conduct a comprehensive evaluation of the response's overall quality with sufficient and explicit evidence**, and then provide a corresponding overall quality score from 0 to 10 within $\\boxed{}$.
39
40Use integers or half-point increments for all scores, with higher numbers representing higher quality.
41
42Below is the response:
43
44[Start of Response]
45{response}
46[End of Response]
47"""
48
49criteria_conversation = [
50 {"role":"user", "content": criteria_prompt.replace("{query}", query)}
51]
52input_ids = tokenizer.apply_chat_template(
53 criteria_conversation,
54 tokenize=True,
55 add_generation_prompt=True,
56 enable_thinking=False,
57 return_tensors="pt").to(model.device)
58output = model.generate(
59 input_ids=input_ids,
60 max_new_tokens=4096,
61 temperature=0,
62)
63criteria = tokenizer.decode(output[0][len(input_ids[0]):], skip_special_tokens=True)
64print(criteria)
65
66evaluation_conversation = [
67 {"role":"user", "content": criteria_prompt.replace("{query}", query)},
68 {"role":"assistant", "content": criteria},
69 {"role":"user", "content": evaluation_prompt.replace("{response}", response)}
70]
71input_ids = tokenizer.apply_chat_template(
72 evaluation_conversation,
73 tokenize=True,
74 add_generation_prompt=True,
75 enable_thinking=False,
76 return_tensors="pt").to(model.device)
77output = model.generate(
78 input_ids=input_ids,
79 max_new_tokens=8192,
80 temperature=0,
81)
82evaluation = tokenizer.decode(output[0][len(input_ids[0]):], skip_special_tokens=True)
83print(evaluation)@article{hu2026rm,
title={CE-RM: A Pointwise Generative Reward Model Optimized via Two-Stage Rollout and Unified Criteria},
author={Hu, Xinyu and He, Yancheng and Wang, Weixun and Feng, Tao and Lin, Li and Liu, Jiashun and Su, Wenbo and Zheng, Bo and Wan, Xiaojun},
journal={arXiv preprint arXiv:2601.20327},
year={2026}
}