Process reward models

[SalmanMohammadi]

Fixing num_labels for reward models

1. Reward models should be initialized with num_labels=1. I've added a num_labels field which is set in cfg_kwargs, rather than model_kwargs, as this field is parsed from config.num_labels in transformers. This also allows us to correctly initialize process reward models with num_labels=2.

Adding support for process reward model training

I've added support for the PRMTrainer from trl, and also for the appropriate dataset format. Please see a screenshot of a successful training run below.

Resulting in the following trained model: https://huggingface.co/smohammadi/Qwen2.5-3B-MathShepherd

[winglian]

A couple of minor things, but should be good to go after.

[winglian]

Looks great to me. Thanks!