More test edits

ludwig-ai · Jun 15, 2023 · 9b46959 · 9b46959
1 parent a0808cf
commit 9b46959
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/tests/integration_tests/test_preprocessing.py b/tests/integration_tests/test_preprocessing.py
@@ -164,17 +164,17 @@ def test_strip_whitespace_category(csv_filename, tmpdir):
 
 
 def test_rlhf_reward_model_data_preprocessor():
-    id_column = "Reward_Session_ID"
-    outcome_column = "Human_Feedback_Outcome"
+    id_column = "reward_session_id"
+    outcome_column = "outcome"
     chosen_value = "some_value_1"
     rejected_value = "some_value_2"
-    transcript_column = "Transcript"
+    transcript_column = "transcript"
 
     # Define the features
     input_features = [
         text_feature(
             name=transcript_column,
-            encoder={"type": "auto_transformer", "pretrained_model_name_or_path": "bert-base-uncased"}
+            encoder={"type": "auto_transformer", "pretrained_model_name_or_path": "bert-base-uncased"},
         )
     ]
     output_features = [number_feature(name=id_column)]

diff --git a/tests/integration_tests/test_trainer.py b/tests/integration_tests/test_trainer.py
@@ -218,17 +218,17 @@ def test_changing_parameters_on_plateau(tmpdir):
 
 
 def test_rlhf_reward_model_trainer(tmpdir):
-    id_column = "Reward_Session_ID"
-    outcome_column = "Human_Feedback_Outcome"
+    id_column = "reward_session_id"
+    outcome_column = "outcome"
     chosen_value = "some_value_1"
     rejected_value = "some_value_2"
-    transcript_column = "Transcript"
+    transcript_column = "transcript"
 
     # Define the features
     input_features = [
         text_feature(
             name=transcript_column,
-            encoder={"type": "auto_transformer", "pretrained_model_name_or_path": "bert-base-uncased"}
+            encoder={"type": "auto_transformer", "pretrained_model_name_or_path": "bert-base-uncased"},
         )
     ]
     output_features = [number_feature(name=id_column)]