Add ratings_values to RatingToArgillaMixin.to_argilla_record (#227)

argilla-io · Jan 9, 2024 · 9d27b61 · 9d27b61
1 parent d9ad890
commit 9d27b61
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 2 deletions.
diff --git a/src/distilabel/tasks/critique/base.py b/src/distilabel/tasks/critique/base.py
@@ -67,12 +67,14 @@ def to_argilla_record(
         generations_column: str = "generations",
         score_column: str = "score",
         critique_column: str = "critique",
+        score_values: Optional[List[int]] = None,
     ) -> Union["FeedbackRecord", List["FeedbackRecord"]]:
         return super().to_argilla_record(
             dataset_row=dataset_row,
             generations_column=generations_column,
             ratings_column=score_column,
             rationale_column=critique_column,
+            ratings_values=score_values or [1, 2, 3, 4, 5],
         )
 
 

diff --git a/src/distilabel/tasks/mixins.py b/src/distilabel/tasks/mixins.py
@@ -136,6 +136,7 @@ def to_argilla_record(  # noqa: C901
         generations_column: str = "generations",
         ratings_column: str = "rating",
         rationale_column: str = "rationale",
+        ratings_values: Optional[List[int]] = None,
     ) -> "FeedbackRecord":
         """Converts a dataset row to an Argilla `FeedbackRecord`."""
         # We start off with the fields, which are the inputs of the LLM, but also
@@ -189,7 +190,12 @@ def to_argilla_record(  # noqa: C901
                             "value": 1
                             if value < 1
                             else int(value)
-                            if value < 10
+                            if value
+                            <= (
+                                max(ratings_values)
+                                if ratings_values is not None
+                                else 10
+                            )
                             else None,
                         }
                     )

diff --git a/src/distilabel/tasks/preference/ultrafeedback.py b/src/distilabel/tasks/preference/ultrafeedback.py
@@ -29,7 +29,7 @@
 from distilabel.tasks.prompt import Prompt
 
 if TYPE_CHECKING:
-    from argilla import FeedbackDataset
+    from argilla import FeedbackDataset, FeedbackRecord
 
 _ULTRAFEEDBACK_TEMPLATE = get_template("ultrafeedback.jinja2")
 
@@ -141,6 +141,25 @@ def to_argilla_dataset(
             ratings_values=ratings_values or [1, 2, 3, 4, 5],
         )
 
+    # Override the default `to_argilla_record` method to provide the `ratings_values` of
+    # UltraFeedback, as the default goes from 1-10 while UltraFeedback's default is 1-5
+    # (0-4 actually, but Argilla doesn't support 0s).
+    def to_argilla_record(
+        self,
+        dataset_row: Dict[str, Any],
+        generations_column: str = "generations",
+        ratings_column: str = "rating",
+        rationale_column: str = "rationale",
+        ratings_values: Optional[List[int]] = None,
+    ) -> "FeedbackRecord":
+        return super().to_argilla_record(
+            dataset_row=dataset_row,
+            generations_column=generations_column,
+            ratings_column=ratings_column,
+            rationale_column=rationale_column,
+            ratings_values=ratings_values or [1, 2, 3, 4, 5],
+        )
+
     @classmethod
     def for_overall_quality(
         cls,