README.md

OPRMs (Outcome \& Process Reward Models) are trained to predict the correctness of each step on the position of "\n\n", as well as the correctness of the whole solution on the position of "\<eos\>".

Usage:

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "ScalableMath/llemma-7b-oprm-prm800k-level-1to3-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/llemma_7b")

qa_example = """# Question

Convert the point $(0,3)$ in rectangular coordinates to polar coordinates.  Enter your answer in the form $(r,\theta),$ where $r > 0$ and $0 \le \theta < 2 \pi.$

# Solution

To convert from rectangular to polar coordinates, I need to use the formulas $r = \sqrt{x^2 + y^2}$ and $\theta = \tan^{-1}(y/x).$

In this case, $x = 0$ and $y = 3,$ so I can plug them into the formulas.

For $r,$ I get $r = \sqrt{0^2 + 3^2} = \sqrt{9} = 3.$

For $\theta,$ I get $\theta = \tan^{-1}(3/0).$

This is undefined, since the tangent function is not defined at $0.$

However, I can use the fact that the point $(0,3)$ lies on the positive $y$-axis, which has an angle of $\pi/2$ radians or $90^\circ.$

Therefore, I can choose any angle in the range $(0,\pi/2)$ as the value of $\theta.$

I will choose $\theta = \pi/2,$ since it is the simplest and most natural choice.

Therefore, the polar coordinates of the point $(0,3)$ are $(3,\pi/2).$

# Answer

(3,\pi/2)"""

begin_solution_tokens = tokenizer.encode("\n\n# Solution", add_special_tokens=False)[1:]
scoring_tokens = tokenizer.encode("\n\n", add_special_tokens=False)[1:]
eos_token = tokenizer.eos_token_id

input_ids = tokenizer.encode(qa_example)

begin_solution_flag = False

candidate_positions = []

for start_idx in range(len(input_ids)):
    if tuple(input_ids[start_idx:start_idx+len(begin_solution_tokens)]) == tuple(begin_solution_tokens):
        begin_solution_flag = True

    if begin_solution_flag and tuple(input_ids[start_idx:start_idx+len(scoring_tokens)]) == tuple(scoring_tokens):
        candidate_positions.append(start_idx)

    if input_ids[start_idx] == eos_token:
        candidate_positions.append(start_idx)
        break

# maybe delete the first and the second to last candidate_positions
# because they are "\n\n" after "# Solution" and after "# Answer"
del candidate_positions[0]
del candidate_positions[-2]

input_tensor = torch.tensor([input_ids])
candidate_positions = torch.tensor(candidate_positions)

with torch.no_grad():
    logits = model(input_tensor).logits
    scores =logits.mean(dim=-1)
    step_scores = scores[0][candidate_positions]
    step_probs = torch.sigmoid(step_scores)

print(step_probs)

# tensor([0.8093, 0.9566, 0.9872, 0.9890, 0.9797, 0.3090, 0.8044, 0.7677, 0.8105, 0.5247])
```