Style edits

ludwig/models/ecd.py

@@ -180,10 +180,11 @@ def get_augmentation_pipelines(self) -> AugmentationPipelines:
 
 
 class RWD(ECD):
-    """
-    This class represents a Reward Model, a model that inputs some feature (i.e. text transcript) and predicts
-    a single scalar output representing the reward/preference of that input. This model type is used for applications
-    such as RLHF fine-tuning of LLMs. This model class is a subclass of ECD, and uses most of ECD's code and pathways.
+    """This class represents a Reward Model, a model type that takes as input some feature (i.e. text) and predicts
+    a single scalar output representing the reward/preference of that input.
+
+    This model type is used for applications such as RLHF fine-tuning of LLMs. This model class is a subclass of ECD,
+    and uses most of ECD's code and pathways.
     """
     @staticmethod
     def type() -> str:

ludwig/schema/decoders/base.py

@@ -3,7 +3,16 @@
 
 from ludwig.api_annotations import DeveloperAPI
 from ludwig.constants import (
-    BINARY, CATEGORY, MODEL_ECD, MODEL_GBM, MODEL_LLM, MODEL_RWD, NUMBER, SET, TIMESERIES, VECTOR)
+    BINARY,
+    CATEGORY,
+    MODEL_ECD,
+    MODEL_GBM,
+    MODEL_LLM,
+    MODEL_RWD,
+    NUMBER,
+    SET,
+    TIMESERIES,
+    VECTOR)
 from ludwig.schema import common_fields
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.decoders.utils import register_decoder_config

ludwig/schema/features/number_feature.py

@@ -16,12 +16,12 @@
     ecd_defaults_config_registry,
     ecd_input_config_registry,
     ecd_output_config_registry,
-    rwd_output_config_registry,
     gbm_defaults_config_registry,
     gbm_input_config_registry,
     gbm_output_config_registry,
     input_mixin_registry,
     output_mixin_registry,
+    rwd_output_config_registry,
 )
 from ludwig.schema.metadata import FEATURE_METADATA
 from ludwig.schema.metadata.parameter_metadata import INTERNAL_ONLY

ludwig/trainers/trainer_rlhf.py

@@ -16,7 +16,9 @@
 
 
 @register_trainer(MODEL_RWD)
-class RewardModelTrainer(Trainer):
+class RWDTrainer(Trainer):
+    """This class trains models of type Reward Model."""
+
     @staticmethod
     def get_schema_cls():
         return RWDTrainerConfig
@@ -77,9 +79,7 @@ def train_step(
 
         # Validate inputs and targets
         if not len(inputs) == 1:
-            raise ValueError(
-                f"Invalid reward model training data inputs, expect 1 input feature, got {len(inputs)}."
-            )
+            raise ValueError(f"Invalid reward model training data inputs, expect 1 input feature, got {len(inputs)}.")
         if not len(targets) == 1:
             raise ValueError(
                 f"Invalid reward model training data targets, expect 1 target feature, got {len(targets)}."