import torch
import torch.nn.functional as F
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
prm_path = "Diginyx/Qwen3.5-27B-prm-ep1"
peft_cfg = PeftConfig.from_pretrained(prm_path)
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
base = AutoModelForSequenceClassification.from_pretrained(
peft_cfg.base_model_name_or_path,
num_labels=1,
quantization_config=bnb_config,
device_map="auto",
)
prm = PeftModel.from_pretrained(base, prm_path)
tokenizer = AutoTokenizer.from_pretrained(prm_path)
prm.eval()
messages = [
{"role": "user", "content": "Guess a 5-letter word."},
{"role": "assistant", "content": "Guess: crane"},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024,
truncation_side="left").to("cuda")
with torch.no_grad():
score = torch.sigmoid(prm(**inputs).logits[0, 0]).item()
print(f"P(success) = {score:.4f}")