From 06a3610376b1cecc9681c1d76cabbe15030181bd Mon Sep 17 00:00:00 2001 From: Francisco Aranda Date: Wed, 17 Jul 2024 09:49:06 +0200 Subject: [PATCH] fix: wrong argilla module import --- src/distilabel/steps/argilla/base.py | 2 +- tests/unit/steps/argilla/test_preference.py | 2 +- tests/unit/steps/argilla/test_text_generation.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/distilabel/steps/argilla/base.py b/src/distilabel/steps/argilla/base.py index d7a8af154e..90ae6d4a11 100644 --- a/src/distilabel/steps/argilla/base.py +++ b/src/distilabel/steps/argilla/base.py @@ -92,7 +92,7 @@ def model_post_init(self, __context: Any) -> None: super().model_post_init(__context) try: - import argilla_sdk as rg # noqa + import argilla as rg # noqa except ImportError as ie: raise ImportError( "Argilla is not installed. Please install it using `pip install argilla --upgrade`." diff --git a/tests/unit/steps/argilla/test_preference.py b/tests/unit/steps/argilla/test_preference.py index 5ff95782b0..edeb566849 100644 --- a/tests/unit/steps/argilla/test_preference.py +++ b/tests/unit/steps/argilla/test_preference.py @@ -15,7 +15,7 @@ import os from unittest.mock import patch -import argilla_sdk as rg +import argilla as rg from distilabel.pipeline.local import Pipeline from distilabel.steps.argilla.preference import PreferenceToArgilla diff --git a/tests/unit/steps/argilla/test_text_generation.py b/tests/unit/steps/argilla/test_text_generation.py index 5679c54517..884ae51cff 100644 --- a/tests/unit/steps/argilla/test_text_generation.py +++ b/tests/unit/steps/argilla/test_text_generation.py @@ -15,7 +15,7 @@ import os from unittest.mock import patch -import argilla_sdk as rg +import argilla as rg from distilabel.pipeline.local import Pipeline from distilabel.steps.argilla.text_generation import TextGenerationToArgilla