Fix some issues related to Magpie task (#783)

* Fix input columns not included in output

* Include `model_name` column as output

* Include `model_name` column

* Fix unit tests magpie

* Fix typo in docstring

Author: gabrielmbmb

src/distilabel/steps/tasks/magpie/base.py

@@ -109,6 +109,17 @@ def _append_messages_to_conversations(
             conversation.append({"role": role, "content": instruction})
         return conversations
 
+    def _generate_instruction(
+        self, inputs: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        prepared_inputs = self._prepare_inputs_for_instruction_generation(inputs)
+        outputs = self.llm.generate(
+            inputs=prepared_inputs,
+            num_generations=1,
+            **self.llm.generation_kwargs,  # type: ignore
+        )
+        return [{"instruction": output[0]} for output in outputs]
+
     def _generate_multi_turn_conversation(
         self, inputs: List[Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
@@ -156,17 +167,16 @@ def _generate_with_pre_query_template(
         Returns:
             The list of generated conversations.
         """
+        outputs = (
+            self._generate_instruction(inputs)
+            if self.only_instruction
+            else self._generate_multi_turn_conversation(inputs)
+        )
 
-        if self.only_instruction:
-            prepared_inputs = self._prepare_inputs_for_instruction_generation(inputs)
-            outputs = self.llm.generate(
-                inputs=prepared_inputs,
-                num_generations=1,
-                **self.llm.generation_kwargs,  # type: ignore
-            )
-            return [{"instruction": output[0]} for output in outputs]
-
-        return self._generate_multi_turn_conversation(inputs)
+        return [
+            {**input, **output, "model_name": self.llm.model_name}
+            for input, output in zip(inputs, outputs)
+        ]
 
 
 class Magpie(Task, MagpieBase):
@@ -209,8 +219,9 @@ class Magpie(Task, MagpieBase):
 
     Output columns:
         - conversation (`ChatType`): the generated conversation which is a list of chat
-            items with a role and a message. Only if `only_instructions=False`.
+            items with a role and a message. Only if `only_instruction=False`.
         - instruction (`str`): the generated instructions if `only_instruction=True`.
+        - model_name (`str`): The model name used to generate the `conversation` or `instruction`.
 
     Categories:
         - text-generation
@@ -352,8 +363,8 @@ def format_input(self, input: Dict[str, Any]) -> "ChatType":
     def outputs(self) -> List[str]:
         """Either a multi-turn conversation or the instruction generated."""
         if self.only_instruction:
-            return ["instruction"]
-        return ["conversation"]
+            return ["instruction", "model_name"]
+        return ["conversation", "model_name"]
 
     def format_output(
         self,

src/distilabel/steps/tasks/magpie/generator.py

@@ -62,6 +62,7 @@ class MagpieGenerator(GeneratorTask, MagpieBase):
         - conversation (`ChatType`): the generated conversation which is a list of chat
             items with a role and a message.
         - instruction (`str`): the generated instructions if `only_instruction=True`.
+        - model_name (`str`): The model name used to generate the `conversation` or `instruction`.
 
     Categories:
         - text-generation
@@ -215,8 +216,8 @@ def format_output(
     def outputs(self) -> List[str]:
         """Either a multi-turn conversation or the instruction generated."""
         if self.only_instruction:
-            return ["instruction"]
-        return ["conversation"]
+            return ["instruction", "model_name"]
+        return ["conversation", "model_name"]
 
     def process(self, offset: int = 0) -> "GeneratorStepOutput":
         """Generates the desired number of instructions or conversations using Magpie.

tests/unit/steps/tasks/magpie/test_base.py

@@ -30,14 +30,14 @@ def test_raise_value_error_llm_no_magpie_mixin(self) -> None:
     def test_outputs(self) -> None:
         task = Magpie(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))
 
-        assert task.outputs == ["conversation"]
+        assert task.outputs == ["conversation", "model_name"]
 
         task = Magpie(
             llm=DummyMagpieLLM(magpie_pre_query_template="llama3"),
             only_instruction=True,
         )
 
-        assert task.outputs == ["instruction"]
+        assert task.outputs == ["instruction", "model_name"]
 
     def test_process(self) -> None:
         task = Magpie(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"), n_turns=1)
@@ -50,18 +50,21 @@ def test_process(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -79,6 +82,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
@@ -88,6 +92,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
@@ -97,6 +102,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -115,31 +121,37 @@ def test_process_with_system_prompt_per_row(self) -> None:
             )
         ) == [
             {
+                "system_prompt": "You're a math expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a math expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
+                "system_prompt": "You're a florist expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a florist expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
+                "system_prompt": "You're a plumber expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a plumber expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -152,9 +164,18 @@ def test_process_only_instruction(self) -> None:
         task.load()
 
         assert next(task.process(inputs=[{}, {}, {}])) == [
-            {"instruction": "Hello Magpie"},
-            {"instruction": "Hello Magpie"},
-            {"instruction": "Hello Magpie"},
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
         ]
 
     def test_serialization(self) -> None:

tests/unit/steps/tasks/magpie/test_generator.py

@@ -30,14 +30,14 @@ def test_raise_value_error_llm_no_magpie_mixin(self) -> None:
     def test_outputs(self) -> None:
         task = MagpieGenerator(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))
 
-        assert task.outputs == ["conversation"]
+        assert task.outputs == ["conversation", "model_name"]
 
         task = MagpieGenerator(
             llm=DummyMagpieLLM(magpie_pre_query_template="llama3"),
             only_instruction=True,
         )
 
-        assert task.outputs == ["instruction"]
+        assert task.outputs == ["instruction", "model_name"]
 
     def test_serialization(self) -> None:
         task = MagpieGenerator(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))