Small edits

ludwig-ai · Jun 15, 2023 · 4860265 · 4860265
1 parent 1662827
commit 4860265
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/ludwig/schema/decoders/base.py b/ludwig/schema/decoders/base.py
@@ -12,7 +12,7 @@
     NUMBER,
     SET,
     TIMESERIES,
-    VECTOR
+    VECTOR,
 )
 from ludwig.schema import common_fields
 from ludwig.schema import utils as schema_utils

diff --git a/ludwig/schema/encoders/text_encoders.py b/ludwig/schema/encoders/text_encoders.py
@@ -1,7 +1,7 @@
 from typing import Callable, Dict, List, Optional, TYPE_CHECKING, Union
 
 from ludwig.api_annotations import DeveloperAPI
-from ludwig.constants import MODEL_ECD, MODEL_GBM, TEXT
+from ludwig.constants import MODEL_ECD, MODEL_GBM, MODEL_RWD, TEXT
 from ludwig.error import ConfigValidationError
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.encoders.sequence_encoders import SequenceEncoderConfig
@@ -599,7 +599,7 @@ def module_name():
 
 
 @DeveloperAPI
-@register_encoder_config("bert", TEXT)
+@register_encoder_config("bert", TEXT, model_types=[MODEL_RWD])
 @ludwig_dataclass
 class BERTConfig(HFEncoderConfig):
     """This dataclass configures the schema used for an BERT encoder."""

diff --git a/ludwig/schema/features/text_feature.py b/ludwig/schema/features/text_feature.py
@@ -100,7 +100,7 @@ class RWDTextInputFeatureConfig(TextInputFeatureConfig):
     encoder: BaseEncoderConfig = EncoderDataclassField(
         MODEL_RWD,
         feature_type=TEXT,
-        default="parallel_cnn",
+        default="bert",
     )
 
 

diff --git a/tests/integration_tests/test_trainer.py b/tests/integration_tests/test_trainer.py
@@ -252,7 +252,7 @@ def test_rlhf_reward_model_trainer(tmpdir):
             "transcript_column": transcript_column,
         }
     }
-    config["model_type"] = "reward_model"
+    config["model_type"] = "rwd"
 
     # Train Ludwig model with the dataset
     ludwig_model = LudwigModel(config, backend=backend)