fix num_labels= 1 test fail (#3493) [skip ci]

* trl_num_lables=1

* casual num_lables=1,rwd model

* lint

src/axolotl/core/builders/causal.py

@@ -421,6 +421,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
             trainer_kwargs["dataset_tags"] = [
                 d["path"] for d in self.cfg.datasets if not Path(d["path"]).is_dir()
             ]
+        # TRL's RewardTrainer validates num_labels=1 on pre-loaded models; ensure the
+        # config reflects this regardless of how the model was instantiated.
+        if (
+            self.cfg.reward_model
+            and getattr(self.model.config, "num_labels", None) != 1
+        ):
+            self.model.config.num_labels = 1
         trainer = trainer_cls(
             model=self.model,
             train_dataset=self.train_dataset,

src/axolotl/utils/schemas/validation.py

@@ -253,6 +253,23 @@ class TrainingValidationMixin:
                 data["pad_to_sequence_len"] = True
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def set_reward_model_defaults(cls, data):
+        if data.get("reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 1
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForSequenceClassification"
+
+        if data.get("process_reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 2
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForTokenClassification"
+
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_gas_bsz(cls, data):