Add magpie unit tests

Author: gabrielmbmb

Diff for: src/distilabel/steps/tasks/magpie/base.py

@@ -61,13 +61,6 @@ class MagpieBase(RuntimeParametersMixin):
         " content of certain topic, guide the style, etc.",
     )
 
-    @property
-    def outputs(self) -> List[str]:
-        """Either a multi-turn conversation or the instruction generated."""
-        if self.only_instruction:
-            return ["instruction"]
-        return ["conversation"]
-
     def _prepare_inputs_for_instruction_generation(
         self, inputs: List[Dict[str, Any]]
     ) -> List["FormattedInput"]:
@@ -121,7 +114,7 @@ def _generate_multi_turn_conversation(
     ) -> List[Dict[str, Any]]:
         conversations = self._prepare_inputs_for_instruction_generation(inputs)
 
-        for _ in range(self.n_turns - 1):  # type: ignore
+        for _ in range(self.n_turns):  # type: ignore
             # Generate instruction or user message
             outputs = self.llm.generate(
                 inputs=conversations,
@@ -355,6 +348,13 @@ def format_input(self, input: Dict[str, Any]) -> "ChatType":
         """Does nothing."""
         return []
 
+    @property
+    def outputs(self) -> List[str]:
+        """Either a multi-turn conversation or the instruction generated."""
+        if self.only_instruction:
+            return ["instruction"]
+        return ["conversation"]
+
     def format_output(
         self,
         output: Union[str, None],

Diff for: src/distilabel/steps/tasks/magpie/generator.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 from pydantic import Field
 
@@ -211,6 +211,13 @@ def format_output(
         """Does nothing."""
         return {}
 
+    @property
+    def outputs(self) -> List[str]:
+        """Either a multi-turn conversation or the instruction generated."""
+        if self.only_instruction:
+            return ["instruction"]
+        return ["conversation"]
+
     def process(self, offset: int = 0) -> "GeneratorStepOutput":
         """Generates the desired number of instructions or conversations using Magpie.
 

Diff for: tests/unit/conftest.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, List
 
 import pytest
-from distilabel.llms.base import AsyncLLM
+from distilabel.llms.base import LLM, AsyncLLM
+from distilabel.llms.mixins.magpie import MagpieChatTemplateMixin
 
 if TYPE_CHECKING:
     from distilabel.llms.typing import GenerateOutput
@@ -37,6 +38,22 @@ async def agenerate(
         return ["output" for _ in range(num_generations)]
 
 
+class DummyMagpieLLM(LLM, MagpieChatTemplateMixin):
+    def load(self) -> None:
+        pass
+
+    @property
+    def model_name(self) -> str:
+        return "test"
+
+    def generate(
+        self, inputs: List["FormattedInput"], num_generations: int = 1, **kwargs: Any
+    ) -> List["GenerateOutput"]:
+        return [
+            ["Hello Magpie" for _ in range(num_generations)] for _ in range(len(inputs))
+        ]
+
+
 @pytest.fixture
 def dummy_llm() -> AsyncLLM:
     return DummyLLM()

Diff for: tests/unit/llms/mixins/test_magpie.py

@@ -0,0 +1,60 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from distilabel.llms.mixins.magpie import MAGPIE_PRE_QUERY_TEMPLATES
+
+from tests.unit.conftest import DummyMagpieLLM
+
+
+class TestMagpieChatTemplateMixin:
+    def test_magpie_pre_query_template_set(self) -> None:
+        with pytest.raises(
+            ValueError,
+            match="Cannot set `use_magpie_template=True` if `magpie_pre_query_template` is `None`",
+        ):
+            DummyMagpieLLM(use_magpie_template=True)
+
+    def test_magpie_pre_query_template_alias_resolved(self) -> None:
+        llm = DummyMagpieLLM(magpie_pre_query_template="llama3")
+        assert llm.magpie_pre_query_template == MAGPIE_PRE_QUERY_TEMPLATES["llama3"]
+
+    def test_apply_magpie_pre_query_template(self) -> None:
+        llm = DummyMagpieLLM(magpie_pre_query_template="<user>")
+
+        assert (
+            llm.apply_magpie_pre_query_template(
+                prompt="<system>Hello hello</system>", input=[]
+            )
+            == "<system>Hello hello</system>"
+        )
+
+        llm = DummyMagpieLLM(
+            use_magpie_template=True, magpie_pre_query_template="<user>"
+        )
+
+        assert (
+            llm.apply_magpie_pre_query_template(
+                prompt="<system>Hello hello</system>", input=[]
+            )
+            == "<system>Hello hello</system><user>"
+        )
+
+        assert (
+            llm.apply_magpie_pre_query_template(
+                prompt="<system>Hello hello</system><user>Hey</user>",
+                input=[{"role": "user", "content": "Hey"}],
+            )
+            == "<system>Hello hello</system><user>Hey</user>"
+        )

Diff for: tests/unit/steps/tasks/magpie/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2023-present, Argilla, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+