Fix unit tests

Author: gabrielmbmb

Diff for: src/distilabel/llms/huggingface/inference_endpoints.py

@@ -159,9 +159,19 @@ def only_one_of_model_id_endpoint_name_or_base_url_provided(
 
         if self.base_url and (self.model_id or self.endpoint_name):
             self._logger.warning(  # type: ignore
-                f"Since the `base_url={self.base_url}` is available and either one of `model_id` or `endpoint_name`"
-                " is also provided, the `base_url` will either be ignored or overwritten with the one generated"
-                " from either of those args, for serverless or dedicated inference endpoints, respectively."
+                f"Since the `base_url={self.base_url}` is available and either one of `model_id`"
+                " or `endpoint_name` is also provided, the `base_url` will either be ignored"
+                " or overwritten with the one generated from either of those args, for serverless"
+                " or dedicated inference endpoints, respectively."
+            )
+
+        if self.model_id and self.tokenizer_id is None:
+            self.tokenizer_id = self.model_id
+
+        if self.use_magpie_template and self.tokenizer_id is None:
+            raise ValueError(
+                "`use_magpie_template` cannot be `True` if `tokenizer_id` is `None`. Please,"
+                " set a `tokenizer_id` and try again."
             )
 
         if self.base_url and not (self.model_id or self.endpoint_name):
@@ -174,9 +184,9 @@ def only_one_of_model_id_endpoint_name_or_base_url_provided(
             return self
 
         raise ValidationError(
-            "Only one of `model_id` or `endpoint_name` must be provided. If `base_url` is provided too,"
-            " it will be overwritten instead. Found `model_id`={self.model_id}, `endpoint_name`={self.endpoint_name},"
-            f" and `base_url`={self.base_url}."
+            f"Only one of `model_id` or `endpoint_name` must be provided. If `base_url` is"
+            f" provided too, it will be overwritten instead. Found `model_id`={self.model_id},"
+            f" `endpoint_name`={self.endpoint_name}, and `base_url`={self.base_url}."
         )
 
     def load(self) -> None:  # noqa: C901

Diff for: src/distilabel/llms/mixins/magpie.py

@@ -48,7 +48,7 @@ class MagpieChatTemplateMixin(BaseModel, validate_assignment=True):
         - [Magpie: Alignment Data Synthesis from Scratch by Prompting Aligned LLMs with Nothing](https://arxiv.org/abs/2406.08464)
     """
 
-    use_magpie_template: bool = True
+    use_magpie_template: bool = False
     magpie_pre_query_template: Union[MagpieAvailablePreQueryTemplates, str, None] = None
 
     @field_validator("magpie_pre_query_template")

Diff for: src/distilabel/pipeline/step_wrapper.py

@@ -16,7 +16,7 @@
 from queue import Queue
 from typing import Any, Dict, List, Optional, Union, cast
 
-from distilabel.llms.mixins import CudaDevicePlacementMixin
+from distilabel.llms.mixins.cuda_device_placement import CudaDevicePlacementMixin
 from distilabel.pipeline.batch import _Batch
 from distilabel.pipeline.constants import LAST_BATCH_SENT_FLAG
 from distilabel.pipeline.typing import StepLoadStatus

Diff for: tests/unit/llms/huggingface/test_inference_endpoints.py

@@ -14,19 +14,38 @@
 
 import random
 from unittest import mock
-from unittest.mock import AsyncMock, MagicMock, Mock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import nest_asyncio
 import pytest
 from distilabel.llms.huggingface.inference_endpoints import InferenceEndpointsLLM
 
 
 @patch("huggingface_hub.AsyncInferenceClient")
-@patch("openai.AsyncOpenAI")
 class TestInferenceEndpointsLLM:
-    def test_load_no_api_key(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+    def test_no_tokenizer_magpie_raise_value_error(
+        self, mock_inference_client: MagicMock
     ) -> None:
+        with pytest.raises(
+            ValueError,
+            match="`use_magpie_template` cannot be `True` if `tokenizer_id` is `None`",
+        ):
+            InferenceEndpointsLLM(
+                base_url="http://localhost:8000",
+                use_magpie_template=True,
+                magpie_pre_query_template="llama3",
+            )
+
+    def test_tokenizer_id_set_if_model_id(
+        self, mock_inference_client: MagicMock
+    ) -> None:
+        llm = InferenceEndpointsLLM(
+            model_id="distilabel-internal-testing/tiny-random-mistral"
+        )
+
+        assert llm.tokenizer_id == llm.model_id
+
+    def test_load_no_api_key(self, mock_inference_client: MagicMock) -> None:
         llm = InferenceEndpointsLLM(
             model_id="distilabel-internal-testing/tiny-random-mistral"
         )
@@ -40,12 +59,8 @@ def test_load_no_api_key(
             ):
                 llm.load()
 
-    def test_load_with_cached_token(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
-    ) -> None:
-        llm = InferenceEndpointsLLM(
-            model_id="distilabel-internal-testing/tiny-random-mistral"
-        )
+    def test_load_with_cached_token(self, mock_inference_client: MagicMock) -> None:
+        llm = InferenceEndpointsLLM(base_url="http://localhost:8000")
 
         # Mock `huggingface_hub.constants.HF_TOKEN_PATH` to exist
         with (
@@ -58,7 +73,7 @@ def test_load_with_cached_token(
             llm.load()
 
     def test_serverless_inference_endpoints_llm(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
             model_id="distilabel-internal-testing/tiny-random-mistral"
@@ -68,7 +83,7 @@ def test_serverless_inference_endpoints_llm(
         assert llm.model_name == "distilabel-internal-testing/tiny-random-mistral"
 
     def test_dedicated_inference_endpoints_llm(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
             endpoint_name="tiny-random-mistral",
@@ -79,11 +94,12 @@ def test_dedicated_inference_endpoints_llm(
         assert llm.model_name == "tiny-random-mistral"
 
     def test_dedicated_inference_endpoints_llm_via_url(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
             base_url="https://api-inference.huggingface.co/models/distilabel-internal-testing/tiny-random-mistral"
         )
+        llm.load()
 
         assert isinstance(llm, InferenceEndpointsLLM)
         assert (
@@ -93,12 +109,12 @@ def test_dedicated_inference_endpoints_llm_via_url(
 
     @pytest.mark.asyncio
     async def test_agenerate_via_inference_client(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
             model_id="distilabel-internal-testing/tiny-random-mistral"
         )
-        llm._aclient = mock_inference_client
+        llm.load()
 
         llm._aclient.text_generation = AsyncMock(
             return_value=" Aenean hendrerit aliquam velit. ..."
@@ -113,39 +129,14 @@ async def test_agenerate_via_inference_client(
             ]
         ) == [" Aenean hendrerit aliquam velit. ..."]
 
-    @pytest.mark.asyncio
-    async def test_agenerate_via_openai_client(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
-    ) -> None:
-        llm = InferenceEndpointsLLM(
-            model_id="distilabel-internal-testing/tiny-random-mistral",
-            use_openai_client=True,
-        )
-        llm._aclient = mock_openai_client
-
-        mocked_completion = Mock(
-            choices=[Mock(message=Mock(content=" Aenean hendrerit aliquam velit. ..."))]
-        )
-        llm._aclient.chat.completions.create = AsyncMock(return_value=mocked_completion)
-
-        assert await llm.agenerate(
-            input=[
-                {"role": "system", "content": ""},
-                {
-                    "role": "user",
-                    "content": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
-                },
-            ]
-        ) == [" Aenean hendrerit aliquam velit. ..."]
-
     @pytest.mark.asyncio
     async def test_generate_via_inference_client(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
-            model_id="distilabel-internal-testing/tiny-random-mistral"
+            model_id="distilabel-internal-testing/tiny-random-mistral",
         )
-        llm._aclient = mock_inference_client
+        llm.load()
 
         llm._aclient.text_generation = AsyncMock(
             return_value=" Aenean hendrerit aliquam velit. ..."
@@ -165,45 +156,15 @@ async def test_generate_via_inference_client(
             ]
         ) == [(" Aenean hendrerit aliquam velit. ...",)]
 
-    @pytest.mark.asyncio
-    async def test_generate_via_openai_client(
-        self, mock_inference_client: MagicMock, mock_openai_client: MagicMock
-    ) -> None:
-        llm = InferenceEndpointsLLM(
-            model_id="distilabel-internal-testing/tiny-random-mistral",
-            use_openai_client=True,
-        )
-        llm._aclient = mock_openai_client
-
-        mocked_completion = Mock(
-            choices=[Mock(message=Mock(content=" Aenean hendrerit aliquam velit. ..."))]
-        )
-        llm._aclient.chat.completions.create = AsyncMock(return_value=mocked_completion)
-
-        ...
-        nest_asyncio.apply()
-
-        assert llm.generate(
-            inputs=[
-                [
-                    {"role": "system", "content": ""},
-                    {
-                        "role": "user",
-                        "content": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
-                    },
-                ]
-            ]
-        ) == [(" Aenean hendrerit aliquam velit. ...",)]
-
     @pytest.mark.asyncio
     async def test_agenerate_with_structured_output(
-        self, mock_inference_client: MagicMock, _: MagicMock
+        self, mock_inference_client: MagicMock
     ) -> None:
         llm = InferenceEndpointsLLM(
             model_id="distilabel-internal-testing/tiny-random-mistral",
             structured_output={"format": "regex", "schema": r"\b[A-Z][a-z]*\b"},
         )
-        llm._aclient = mock_inference_client
+        llm.load()
 
         llm._aclient.text_generation = AsyncMock(
             return_value=" Aenean hendrerit aliquam velit. ..."
@@ -223,7 +184,7 @@ async def test_agenerate_with_structured_output(
         ) == [" Aenean hendrerit aliquam velit. ..."]
 
         kwargs = {
-            "prompt": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+            "prompt": "<s>[INST] Lorem ipsum dolor sit amet, consectetur adipiscing elit. [/INST]",
             "max_new_tokens": 128,
             "do_sample": False,
             "typical_p": None,
@@ -235,15 +196,11 @@ async def test_agenerate_with_structured_output(
             "return_full_text": False,
             "watermark": False,
             "grammar": {"type": "regex", "value": "\\b[A-Z][a-z]*\\b"},
-            "seed": 478163327,  # pre-computed random value with `random.seed(42)`
+            "seed": 2053695854357871005,  # pre-computed random value with `random.seed(42)`
         }
-        mock_inference_client.text_generation.assert_called_with(**kwargs)
+        llm._aclient.text_generation.assert_called_with(**kwargs)
 
-    def test_serialization(
-        self,
-        mock_inference_client: MagicMock,
-        mock_openai_client: MagicMock,
-    ) -> None:
+    def test_serialization(self, mock_inference_client: MagicMock) -> None:
         llm = InferenceEndpointsLLM(
             model_id="distilabel-internal-testing/tiny-random-mistral",
         )
@@ -253,11 +210,12 @@ def test_serialization(
             "endpoint_name": None,
             "endpoint_namespace": None,
             "base_url": None,
-            "tokenizer_id": None,
+            "tokenizer_id": "distilabel-internal-testing/tiny-random-mistral",
             "generation_kwargs": {},
+            "magpie_pre_query_template": None,
             "structured_output": None,
             "model_display_name": None,
-            "use_openai_client": False,
+            "use_magpie_template": False,
             "type_info": {
                 "module": "distilabel.llms.huggingface.inference_endpoints",
                 "name": "InferenceEndpointsLLM",