Update ultrafeedback from tests

argilla-io · Jan 8, 2024 · b8cc8cc · b8cc8cc
1 parent 95dac2c
commit b8cc8cc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -23,7 +23,7 @@
 @pytest.fixture
 def custom_dataset():
     ds = CustomDataset.from_dict({"input": ["a", "b"], "generations": ["c", "d"]})
-    ds.task = UltraFeedbackTask.for_text_quality()
+    ds.task = UltraFeedbackTask.for_overall_quality()
     return ds
 
 
@@ -60,7 +60,7 @@ def test_do_checkpoint(
     ds = CustomDataset.from_dict(
         {"input": ["a"] * dataset_len, "generations": ["a"] * dataset_len}
     )
-    ds.task = UltraFeedbackTask.for_text_quality()
+    ds.task = UltraFeedbackTask.for_overall_quality()
     chk = DatasetCheckpoint(save_frequency=save_frequency)
     ctr = 0
     with tempfile.TemporaryDirectory() as tmpdir: