Fix StepOutput type (#1072)

argilla-io · Nov 28, 2024 · 8dd6405 · 8dd6405
1 parent a8d02c2
commit 8dd6405
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 18 deletions.
diff --git a/src/distilabel/steps/typing.py b/src/distilabel/steps/typing.py
@@ -14,24 +14,8 @@
 
 from typing import Any, Dict, Iterator, List, Tuple, Union
 
-StepData = List[Dict[str, Any]]
-StepStatistics = Dict[str, Any]
-StepOutput = Iterator[Dict[str, Union[StepData, StepStatistics]]]
-r"""`StepOutput` is an alias of the typing.
-A step output is a dict of the form:
-{
-		"outputs": [
-            {"col1": "val1", "col2": "val2"},
-            {"col1": "val1", "col2": "val2"},
-            {"col1": "val1", "col2": "val2"},
-		],
-		"statistics": {
-            "llm": {},
-            "time": 12341234,
-            ...
-		}
-}
-"""
+StepOutput = Iterator[List[Dict[str, Any]]]
+
 
 GeneratorStepOutput = Iterator[Tuple[List[Dict[str, Any]], bool]]
 """`GeneratorStepOutput` is an alias of the typing `Iterator[Tuple[List[Dict[str, Any]], bool]]`"""

diff --git a/tests/unit/steps/argilla/test_preference.py b/tests/unit/steps/argilla/test_preference.py
@@ -85,6 +85,8 @@ def test_process(self, mock_dataset) -> None:
             step.load()
         step._instruction = "instruction"
         step._generations = "generations"
+        step._ratings = "ratings"
+        step._rationales = "rationales"
         step._dataset = mock_dataset  # type: ignore
 
         step._dataset.records.log = lambda x: x  # type: ignore