Yingqian/dream_prm_math_7b

Usage

This checkpoint is used as a reward model for math chain-of-thought reasoning. It takes a chat-style conversation as input and outputs a scalar reward equal to the model’s preference for the "+" label over the "-" label.

Load the model

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "Yingqian/dream_prm_math_7b"  # <-- replace with this repo name

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# padding settings used in our code
tokenizer.padding_side = "right"
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id

model.eval()

Compute a reward score

The reward is implemented by looking at the model’s logits on the "+" and "-" tokens at a fixed position near the end of the sequence and taking the softmax probability of "+" as the score.

import torch

def get_reward(conversation):
    """
    conversation: list of dicts, e.g.
    [
        {"role": "user", "content": "<question and partial reasoning>"},
        {"role": "assistant", "content": "+"},
    ]
    Returns: float in (0, 1), probability of "+".
    """

    device = next(model.parameters()).device

    # Apply the model's chat template
    input_ids = tokenizer.apply_chat_template(
        conversation,
        return_tensors="pt"
    ).to(device)

    # IDs of "+" and "-" (last token of each encoding)
    plus_id = tokenizer.encode("+")[-1]
    minus_id = tokenizer.encode("-")[-1]
    candidate_ids = [plus_id, minus_id]

    with torch.no_grad():
        logits = model(input_ids).logits
        # In our experiments we read the logits from the 4th token from the end
        token_logits = logits[:, -4, candidate_ids]
        probs = token_logits.softmax(dim=-1)

    # Probability that the label is "+"
    return probs[0, 0].item()

Example for math reasoning

During our experiments, we call the reward model on a math question plus partial reasoning (either a planned sub-question or a candidate answer step). For example:

question = "Q: If 2x + 3 = 11, what is x?"
partial_step = "Step 1: Subtract 3 from both sides to get 2x = 8."
reward_context = question + "\n" + partial_step

conversation = [
    {"role": "user", "content": reward_context},
    {"role": "assistant", "content": "+"},  # target label
]

score = get_reward(conversation)
print("Reward score (probability of '+'):", score)

A higher score indicates that the model judges this step/answer as better (more preferred), and this scalar can be used as the reward signal to guide search or RL.