Qwen2.5-1.5B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_PATH = "Qwen2.5-1.5B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter"

tokenizer = None
model = None
device = None

def load_model():
    global tokenizer, model, device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading model from {MODEL_PATH}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_PATH,
        num_labels=1,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    )
    model.to(device)
    model.eval()
    print(f"Model loaded on {device}")

load_model()

def get_reward_score(prompt, response):
    messages = [
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": response}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    
    inputs = tokenizer(text, return_tensors="pt", truncation=False, max_length=None).to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        reward = outputs.logits[0, 0].item()
    
    return reward