Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4model_name = "THU-KEG/WildReward-8B"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForSequenceClassification.from_pretrained(model_name)
7
8def build_text(query, response, history_str=""):
9 """Format input text for reward model scoring."""
10 text = f"""
11# Task Description
12You are an expert conversation evaluator. Your task is to judge the **User's Satisfaction** with the Assistant's response based on the conversation context.
13Please rate the response on a scale of 1 to 5 integers.
14
15# Scoring Criteria
16[1] CLEARLY NEGATIVE / REJECTION
17[2] CORRECTION / ERROR POINTER (Negative)
18[3] NEUTRAL
19[4] POSITIVE ENGAGEMENT
20[5] CLEAR SATISFACTION
21
22# Input Data
23## Context (History)
24{history_str}
25
26## User Query
27{query}
28
29## Assistant Response
30{response}
31
32# Output
33Based on the criteria above, please output ONLY the integer score (1, 2, 3, 4, or 5).
34"""
35 return text.strip()
36
37# Prepare query and response
38query = "Explain quantum computing in simple terms."
39response = "Quantum computing uses quantum bits or 'qubits' that can exist in multiple states simultaneously, unlike classical bits..."
40
41# Build formatted text
42text = build_text(query, response)
43
44# Tokenize
45inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=4096).to(model.device)
46
47# Get reward score
48with torch.no_grad():
49 outputs = model(**inputs)
50 logits = outputs.logits
51
52 # CORAL / Ordinal Regression (output shape: 1, K-1)
53 probs = torch.sigmoid(logits)
54 reward = 1 + torch.sum(probs).item()
55
56print(f"Reward score: {reward:.2f} (scale: 1-5)")1@misc{peng2026wildrewardlearningrewardmodels,
2 title={WildReward: Learning Reward Models from In-the-Wild Human Interactions},
3 author={Hao Peng and Yunjia Qi and Xiaozhi Wang and Zijun Yao and Lei Hou and Juanzi Li},
4 year={2026},
5 eprint={2602.08829},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2602.08829},
9}