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

import os
from trl import RewardTrainer, RewardConfig
from datasets import load_dataset


os.environ["WANDB_PROJECT"] = "hh"

training_args = RewardConfig(
    output_dir="output/",
    report_to="wandb",
    run_name="Qwen2.5-1.5B-Instruct-ultrafeedback_binarized-reward",
    num_train_epochs=3,
    per_device_train_batch_size=16,  
    gradient_accumulation_steps=4,  
    learning_rate=1e-5,                      
    warmup_ratio=0.1,
    center_rewards_coefficient=1e-2, 
    bf16=True, 
)

trainer = RewardTrainer(
    model="model/Qwen/Qwen2.5-1.5B-Instruct",
    args=training_args,
    train_dataset=load_dataset("trl-lib/ultrafeedback_binarized", split="train"),
)
trainer.train()

from transformers import pipeline

text = "The capital of France is Paris."
rewarder = pipeline(model="None", device="cuda")
output = rewarder(text)[0]
print(output["score"])