Views
No views yet
| Property | Value |
|---|---|
| Base Architecture | Qwen3-8B (Qwen3ForProcessRewardModel) |
| Parameters | ~8B |
| Precision | bfloat16 |
| Max Sequence Length | 40960 tokens |
| Output Labels | 2 (negative / positive per step) |
| Step Separator Token | <extra_0> |
pip install torch transformersThe model uses customtrust_remote_codeclasses (v1_fin_prm.Qwen3ForProcessRewardModelandv1_fin_config.Qwen3PRMConfig) that are loaded automatically via theauto_mapinconfig.json.
1import torch
2from transformers import AutoModel, AutoTokenizer
3
4MODEL_PATH = "path/to/Dianjin-PRM"
5
6model = AutoModel.from_pretrained(
7 MODEL_PATH,
8 trust_remote_code=True,
9 device_map=None,
10).eval()
11
12# Multi-GPU via DataParallel (optional)
13model = torch.nn.DataParallel(model).cuda()
14
15tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)<extra_0>:##Question
<your question here>
##Thinking Trajectory
<step 1><extra_0><step 2><extra_0>...<step N><extra_0>1question = "What is the present value of $1000 received in 5 years at a 10% discount rate?"
2
3steps = [
4 "We need to calculate the present value using the formula PV = FV / (1 + r)^n.",
5 "Substituting the values: PV = 1000 / (1 + 0.10)^5.",
6 "PV = 1000 / 1.61051 ≈ 620.92.",
7]
8
9trajectory = "<extra_0>".join(steps) + "<extra_0>"
10completion = f"##Question\n{question}\n\n##Thinking Trajectory\n{trajectory}"1def make_step_rewards(logits, token_masks):
2 """Extract per-step reward scores from model logits."""
3 probabilities = torch.nn.functional.softmax(logits, dim=-1)
4 probabilities = probabilities * token_masks.unsqueeze(-1)
5 all_scores_res = []
6 for i in range(probabilities.size(0)):
7 sample = probabilities[i]
8 positive_probs = sample[sample != 0].view(-1, 2)[:, 1]
9 all_scores_res.append(positive_probs.cpu().tolist())
10 return all_scores_res
11
12
13# Tokenize
14input_ids = tokenizer(
15 [completion],
16 return_tensors="pt",
17 padding=True,
18 truncation=True,
19)["input_ids"].to("cuda")
20
21# Forward pass
22with torch.inference_mode():
23 outputs = model(input_ids=input_ids)
24
25# Build step-separator mask and extract rewards
26step_sep_id = tokenizer.encode("<extra_0>")[0]
27token_masks = (input_ids == step_sep_id)
28step_rewards = make_step_rewards(outputs.logits, token_masks)
29
30print(step_rewards)
31# e.g. [[0.92, 0.87, 0.95]] — one score per step, per sample<extra_0> step boundary. Higher values indicate higher-quality reasoning steps.1import numpy as np
2
3candidates = [...] # list of (trajectory_string, final_answer) tuples
4all_rewards = []
5
6for trajectory, answer in candidates:
7 completion = f"##Question\n{question}\n\n##Thinking Trajectory\n{trajectory}"
8 input_ids = tokenizer(
9 [completion], return_tensors="pt", padding=True, truncation=True
10 )["input_ids"].to("cuda")
11
12 with torch.inference_mode():
13 outputs = model(input_ids=input_ids)
14
15 step_sep_id = tokenizer.encode("<extra_0>")[0]
16 token_masks = (input_ids == step_sep_id)
17 rewards = make_step_rewards(outputs.logits, token_masks)
18 # Use the minimum step score as the overall trajectory score
19 all_rewards.append(min(rewards[0]))
20
21best_idx = int(np.argmax(all_rewards))
22best_answer = candidates[best_idx][1]| Component | Description |
|---|---|
##Question | The original question/problem |
##Thinking Trajectory | Reasoning steps separated by <extra_0> |
<extra_0> | Special token used as step separator (token id: 151669) |
<extra_0> position.padding=True.torch.nn.DataParallel.trust_remote_code=True when loading the model, as it relies on custom architecture classes.