Implement "Improving Text Embeddings with LLMs"

src/distilabel/steps/tasks/__init__.py

@@ -23,6 +23,12 @@
 from distilabel.steps.tasks.evol_quality.base import EvolQuality
 from distilabel.steps.tasks.generate_embeddings import GenerateEmbeddings
 from distilabel.steps.tasks.genstruct import Genstruct
+from distilabel.steps.tasks.improving_text_embeddings import (
+    BrainstormEmbeddingTasks,
+    GenerateTextClassificationData,
+    GenerateTextMatchingData,
+    GenerateTextRetrievalData,
+)
 from distilabel.steps.tasks.instruction_backtranslation import (
     InstructionBacktranslation,
 )
@@ -38,6 +44,7 @@
 __all__ = [
     "Task",
     "GeneratorTask",
+    "BrainstormEmbeddingTasks",
     "ChatGeneration",
     "ChatItem",
     "ChatType",
@@ -48,6 +55,9 @@
     "EvolInstructGenerator",
     "EvolQuality",
     "GenerateEmbeddings",
+    "GenerateTextClassificationData",
+    "GenerateTextMatchingData",
+    "GenerateTextRetrievalData",
     "Genstruct",
     "InstructionBacktranslation",
     "PairRM",

src/distilabel/steps/tasks/base.py

@@ -64,7 +64,9 @@ def load(self) -> None:
 
     @abstractmethod
     def format_output(
-        self, output: Union[str, None], input: Dict[str, Any]
+        self,
+        output: Union[str, None],
+        input: Union[Dict[str, Any], None] = None,
     ) -> Dict[str, Any]:
         """Abstract method to format the outputs of the task. It needs to receive an output
         as a string, and generates a Python dictionary with the outputs of the task. In
@@ -74,7 +76,9 @@ def format_output(
         pass
 
     def _format_outputs(
-        self, outputs: "GenerateOutput", inputs: List[Dict[str, Any]]
+        self,
+        outputs: "GenerateOutput",
+        inputs: Union[List[Dict[str, Any]], None] = None,
     ) -> List[Dict[str, Any]]:
         """Formats the outputs of the task using the `format_output` method. If the output
         is `None` (i.e. the LLM failed to generate a response), then the outputs will be
@@ -87,8 +91,11 @@ def _format_outputs(
         Returns:
             A list containing a dictionary with the outputs of the task for each input.
         """
+        if inputs is None:
+            inputs = [None]  # type: ignore
+
         formatted_outputs = []
-        for output, input in zip(outputs, inputs * len(outputs)):
+        for output, input in zip(outputs, inputs * len(outputs)):  # type: ignore
             try:
                 formatted_output = self.format_output(output, input)
                 formatted_output = self._maybe_add_raw_output(
@@ -103,7 +110,7 @@ def _format_outputs(
         return formatted_outputs
 
     def _output_on_failure(
-        self, output: Union[str, None], input: Dict[str, Any]
+        self, output: Union[str, None], input: Union[Dict[str, Any], None] = None
     ) -> Dict[str, Any]:
         """In case of failure to format the output, this method will return a dictionary including
         a new field `distilabel_meta` with the raw output of the LLM.
@@ -183,14 +190,14 @@ def process(self, inputs: StepInput) -> "StepOutput":  # type: ignore
             if self.group_generations:
                 combined = combine_dicts(*formatted_outputs)
                 task_outputs.append(
-                    {**input, "model_name": self.llm.model_name, **combined}
+                    {**input, **combined, "model_name": self.llm.model_name}
                 )
                 continue
 
             # Create a row per generation
             for formatted_output in formatted_outputs:
                 task_outputs.append(
-                    {**input, "model_name": self.llm.model_name, **formatted_output}
+                    {**input, **formatted_output, "model_name": self.llm.model_name}
                 )
 
         yield task_outputs