Untitled

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

# Model and tokenizer loading
prm_model_path = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
tokenizer = AutoTokenizer.from_pretrained(prm_model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(prm_model_path, trust_remote_code=True).eval().to("cuda", dtype=torch.float16)

# Data preparation
datas = [
    {
        "problem": "Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?",
        "response": "To determine how much money Janet makes every day at the farmers' market, we need to follow these steps:\n1. **Calculate the total number of eggs laid by the ducks per day.**\n   Janet's ducks lay 16 eggs per day.\n2. **Determine the number of eggs Janet uses each day.**\n   - She eats 3 eggs for breakfast every morning.\n   - She bakes muffins for her friends every day with 4 eggs.\n   So, the total number of eggs used per day is:\n   \\[\n   3 + 4 = 7 \\text{ eggs}\n   \\]\n3. **Calculate the number of eggs Janet sells at the farmers' market each day.**\n   Subtract the number of eggs used from the total number of eggs laid:\n   \\[\n   16 - 7 = 9 \\text{ eggs}\n   \\]\n4. **Determine how much money Janet makes from selling the eggs.**\n   She sells each egg for $2, so the total amount of money she makes is:\n   \\[\n   9 \\times 2 = 18 \\text{ dollars}\n   \\]\nTherefore, the amount of money Janet makes every day at the farmers' market is \\(\\boxed{18}\\)."
    }
]
problems = [data["problem"] for data in datas]
responses = [data["response"] for data in datas]

# Tokenize input efficiently
encoded_input = tokenizer.batch_encode_plus(
    [(problems[i], responses[i]) for i in range(len(problems))],
    return_tensors="pt",
    padding=True,
    truncation=True,
    max_length=128
).to("cuda")

# Inference with mixed precision
with torch.inference_mode():
    outputs = model(**encoded_input)
    logits = outputs.logits.float()  # Convert back to float for stability
    rewards = torch.sigmoid(logits).squeeze()

print("Normalized Step Rewards:", rewards)