Fix tests after refactor

argilla-io · Jun 6, 2024 · ed319dc · ed319dc
1 parent 76be9a7
commit ed319dc
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 12 deletions.
diff --git a/tests/unit/llms/test_vertexai.py b/tests/unit/llms/test_vertexai.py
@@ -115,7 +115,6 @@ def test_serialization(self, _: MagicMock) -> None:
         _dump = {
             "model": "gemini-1.0-pro",
             "generation_kwargs": {},
-            "structured_output": None,
             "type_info": {
                 "module": "distilabel.llms.vertexai",
                 "name": "VertexAILLM",

diff --git a/tests/unit/steps/tasks/evol_instruct/test_base.py b/tests/unit/steps/tasks/evol_instruct/test_base.py
@@ -127,7 +127,6 @@ def test_serialization(self, dummy_llm: LLM) -> None:
             "input_batch_size": task.input_batch_size,
             "llm": {
                 "generation_kwargs": {},
-                "structured_output": None,
                 "type_info": {
                     "module": task.llm.__module__,
                     "name": task.llm.__class__.__name__,

diff --git a/tests/unit/steps/tasks/evol_instruct/test_generator.py b/tests/unit/steps/tasks/evol_instruct/test_generator.py
@@ -117,7 +117,6 @@ def test_serialization(self, dummy_llm: LLM) -> None:
             "name": "task",
             "llm": {
                 "generation_kwargs": {},
-                "structured_output": None,
                 "type_info": {
                     "module": task.llm.__class__.__module__,
                     "name": task.llm.__class__.__name__,

diff --git a/tests/unit/steps/tasks/evol_quality/test_base.py b/tests/unit/steps/tasks/evol_quality/test_base.py
@@ -98,7 +98,6 @@ def test_serialization(self, dummy_llm: LLM) -> None:
             "input_batch_size": task.input_batch_size,
             "llm": {
                 "generation_kwargs": {},
-                "structured_output": None,
                 "type_info": {
                     "module": task.llm.__module__,
                     "name": task.llm.__class__.__name__,

diff --git a/tests/unit/steps/tasks/structured_outputs/test_outlines.py b/tests/unit/steps/tasks/structured_outputs/test_outlines.py
@@ -17,9 +17,10 @@
 import pytest
 from distilabel.llms.huggingface.transformers import TransformersLLM
 from distilabel.steps.tasks.structured_outputs.outlines import (
-    StructuredOutputType,
+    # StructuredOutputType,
     model_to_schema,
 )
+from distilabel.steps.tasks.typing import OutlinesStructuredOutputType
 from pydantic import BaseModel
 
 
@@ -88,18 +89,14 @@ class DummyUserTest(BaseModel):
 
 
 class TestOutlinesIntegration:
-    # @pytest.mark.skipif(
-    #     not DISTILABEL_RUN_SLOW_TESTS,
-    #     reason="Slow tests, run locally when needed.",
-    # )
     @pytest.mark.parametrize(
         "format, schema, prompt",
         [
             (
                 "json",
                 DummyUserTest,
                 "Create a user profile with the fields name, last_name and id",
-            ),  #
+            ),
             (
                 "json",
                 model_to_schema(DummyUserTest),
@@ -117,7 +114,9 @@ def test_generation(
     ) -> None:
         llm = TransformersLLM(
             model="openaccess-ai-collective/tiny-mistral",
-            structured_output=StructuredOutputType(format=format, schema=schema),
+            structured_output=OutlinesStructuredOutputType(
+                format=format, schema=schema
+            ),
         )
         llm.load()
 
@@ -154,7 +153,9 @@ def test_serialization(
     ) -> None:
         llm = TransformersLLM(
             model="openaccess-ai-collective/tiny-mistral",
-            structured_output=StructuredOutputType(format=format, schema=schema),
+            structured_output=OutlinesStructuredOutputType(
+                format=format, schema=schema
+            ),
         )
         llm.load()
         assert llm.dump() == dump