agadetskii/Qwen2.5-14B-Instruct-uPRM-T80-adapters

import torch
from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F

# load model and tokenizer

model_name = "agadetskii/Qwen2.5-14B-Instruct-uPRM-T80-adapters"

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True,
    force_download=True
)

model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype="bfloat16",
    device_map="auto",            # or "cuda" for single GPU
    attn_implementation="flash_attention_2",
)

device = torch.device("cuda")
model = model.to(device)


# example data

DEFAULT_SYSTEM_PROMPT = """You are a strict mathematical reasoning judge.

Your task is to evaluate one individual reasoning step of a math problem at a time.

- If the step is mathematically correct, respond with `+`.  
- If the step is mathematically incorrect or logically flawed, respond with `-`.  
- Do not provide any explanation, comment, or feedback — only respond with `+` or `-`, and nothing else.  
- Each input is either a single reasoning step or a new problem followed by its first reasoning step. In both cases, evaluate only the validity of the reasoning step.
- For each new problem, once you determine that a step is incorrect, you must consider all subsequent steps for that problem to also be incorrect, and respond with `-` for them as well.

Your response must only be one of these two symbols: `+` or `-`.
"""

data = {
    "system": DEFAULT_SYSTEM_PROMPT,
    "query": "Sue lives in a fun neighborhood. One weekend, the neighbors decided to play a prank on Sue.  On Friday morning, the neighbors placed 18 pink plastic flamingos out on Sue's front yard.  On Saturday morning, the neighbors took back one third of the flamingos, painted them white, and put these newly painted white flamingos back out on Sue's front yard.  Then, on Sunday morning, they added another 18 pink plastic flamingos to the collection. At noon on Sunday, how many more pink plastic flamingos were out than white plastic flamingos?",
    "response": [
      "To find out how many more pink plastic flamingos were out than white plastic flamingos at noon on Sunday, we can break down the problem into steps. First, on Friday, the neighbors start with 18 pink plastic flamingos.",
      "On Saturday, they take back one third of the flamingos. Since there were 18 flamingos, (1/3 \\times 18 = 6) flamingos are taken back. So, they have (18 - 6 = 12) flamingos left in their possession. Then, they paint these 6 flamingos white and put them back out on Sue's front yard. Now, Sue has the original 12 pink flamingos plus the 6 new white ones. Thus, by the end of Saturday, Sue has (12 + 6 = 18) pink flamingos and 6 white flamingos.",
      "On Sunday, the neighbors add another 18 pink plastic flamingos to Sue's front yard. By the end of Sunday morning, Sue has (18 + 18 = 36) pink flamingos and still 6 white flamingos.",
      "To find the difference, subtract the number of white flamingos from the number of pink flamingos: (36 - 6 = 30). Therefore, at noon on Sunday, there were 30 more pink plastic flamingos out than white plastic flamingos. The answer is (\\boxed{30})."
    ]
}

messages = [
    {"role": "system", "content": data['system']},
]
for i in range(len(data["response"])):
    if i == 0:
        usr_msg = data["query"] + " " + data["response"][i]
    else:
        usr_msg = data["response"][i]
    messages.append({"role": "user", "content": usr_msg})
    messages.append({"role": "assistant", "content": "<|*|>"})


conversation_str = tokenizer.apply_chat_template(
    messages, 
    tokenize=False, 
    add_generation_prompt=False
)

input_ids = tokenizer.encode(
    conversation_str, 
    return_tensors="pt", 
).to(model.device)

with torch.no_grad():
    outputs = model(input_ids=input_ids)



def make_step_rewards(logits, token_masks):
    probabilities = F.softmax(logits, dim=-1)
    probabilities = probabilities * token_masks.unsqueeze(-1) # bs, seq_len, num_labels
    
    all_scores_res = []
    for i in range(probabilities.size(0)):
        sample = probabilities[i] # seq_len, num_labels
        positive_probs = sample[sample != 0].view(-1, 2)[:, 0] # valid_tokens, num_labels
        non_zero_elements_list = positive_probs.cpu().tolist()
        all_scores_res.append(non_zero_elements_list)
    return all_scores_res


step_sep_id = tokenizer.encode("<|*|>")[0]
token_masks = (input_ids == step_sep_id)
step_reward = make_step_rewards(outputs[0], token_masks)
print(step_reward)
# returns [[0.9668280482292175, 0.2960759401321411, 0.9475431442260742, 0.9833130240440369]]