feat: 230 feature add sentence transformers support for the to argilla method

pyproject.toml

@@ -40,7 +40,7 @@ openai = ["openai >= 1.0.0"]
 vllm = ["vllm >= 0.2.1"]
 vertexai = ["google-cloud-aiplatform >= 1.38.0"]
 together = ["together"]
-argilla = ["argilla >= 1.18.0"]
+argilla = ["argilla >= 1.22.0", "sentence-transformers >= 2.0.0"]
 tests = ["pytest >= 7.4.0"]
 docs = [
     "mkdocs-material >= 9.5.0",

src/distilabel/tasks/base.py

@@ -18,16 +18,25 @@
     import importlib_resources
 else:
     import importlib.resources as importlib_resources
-
+import warnings
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Union
 
 from jinja2 import Template
 
 from distilabel.tasks.prompt import Prompt
+from distilabel.utils.imports import _ARGILLA_AVAILABLE
+
+if _ARGILLA_AVAILABLE:
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 
 def get_template(template_name: str) -> str:
@@ -119,6 +128,44 @@ def to_argilla_record(
             " `FeedbackDataset` you will need to implement this method first."
         )
 
+    def add_vectors_to_argilla_dataset(
+        self,
+        dataset: Union["FeedbackRecord", List["FeedbackRecord"], "FeedbackDataset"],
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"],
+    ) -> Union["FeedbackRecord", List["FeedbackRecord"], "FeedbackDataset"]:
+        if _ARGILLA_AVAILABLE and vector_strategy:
+            try:
+                if isinstance(vector_strategy, SentenceTransformersExtractor):
+                    ste = vector_strategy
+
+                elif vector_strategy is True:
+                    ste = SentenceTransformersExtractor(
+                        model="en",
+                        show_progress=True,
+                    )
+                else:
+                    raise ValueError(
+                        "The `vector_strategy` must be either `True` or a `SentenceTransformersExtractor` instance."
+                    )
+
+                dataset = ste.update_dataset(dataset=dataset)
+            except Exception as e:
+                warnings.warn(
+                    f"An error occurred while adding vectors to the dataset: {e}",
+                    stacklevel=2,
+                )
+
+        elif not _ARGILLA_AVAILABLE and vector_strategy:
+            warnings.warn(
+                "An error occurred while adding vectors to the dataset: "
+                "The `argilla`/`sentence-transformers` packages are not installed or the installed version is not compatible with the"
+                " required version. If you want to add vectors to your dataset, please run `pip install 'distilabel[vectors]'`.",
+                stacklevel=2,
+            )
+        else:
+            pass
+        return dataset
+
     # Renamed to _to_argilla_record instead of renaming `to_argilla_record` to protected, as that would
     # imply more breaking changes.
     def _to_argilla_record(  # noqa: C901

src/distilabel/tasks/critique/base.py

@@ -22,6 +22,9 @@
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 
 @dataclass
@@ -52,13 +55,15 @@ def to_argilla_dataset(
         score_column: str = "score",
         critique_column: str = "critique",
         score_values: Optional[List[int]] = None,
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
     ) -> "FeedbackDataset":
         return super().to_argilla_dataset(
             dataset_row=dataset_row,
             generations_column=generations_column,
             ratings_column=score_column,
             rationale_column=critique_column,
             ratings_values=score_values or [1, 2, 3, 4, 5],
+            vector_strategy=vector_strategy,
         )
 
     def to_argilla_record(

src/distilabel/tasks/critique/ultracm.py

@@ -14,14 +14,17 @@
 
 import re
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
 
 from distilabel.tasks.base import get_template
 from distilabel.tasks.critique.base import CritiqueTask, CritiqueTaskOutput
 from distilabel.tasks.prompt import Prompt
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 _ULTRACM_TEMPLATE = get_template("ultracm.jinja2")
 
@@ -102,11 +105,13 @@ def to_argilla_dataset(
         score_column: str = "score",
         critique_column: str = "critique",
         score_values: Optional[List[int]] = None,
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
     ) -> "FeedbackDataset":
         return super().to_argilla_dataset(
             dataset_row=dataset_row,
             generations_column=generations_column,
             score_column=score_column,
             critique_column=critique_column,
             score_values=score_values or [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            vector_strategy=vector_strategy,
         )

src/distilabel/tasks/mixins.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Protocol
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Protocol, Union
 
+from distilabel.tasks.base import Task
 from distilabel.utils.argilla import (
     infer_fields_from_dataset_row,
     model_metadata_from_dataset_row,
@@ -26,6 +27,9 @@
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 
 class TaskProtocol(Protocol):
@@ -46,6 +50,7 @@ def to_argilla_dataset(
         ratings_column: str = "rating",
         rationale_column: str = "rationale",
         ratings_values: Optional[List[int]] = None,
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
     ) -> "FeedbackDataset":
         # First we infer the fields from the input_args_names, but we could also
         # create those manually instead using `rg.TextField(...)`
@@ -116,11 +121,15 @@ def to_argilla_dataset(
         )
         # Then we just return the `FeedbackDataset` with the fields, questions, and metadata properties
         # defined above.
-        return rg.FeedbackDataset(
+        dataset = rg.FeedbackDataset(
             fields=fields,
             questions=questions,
             metadata_properties=metadata_properties,  # Note that these are always optional
         )
+        dataset = Task.add_vectors_to_argilla_dataset(
+            dataset=dataset, vector_strategy=vector_strategy
+        )
+        return dataset
 
     def _merge_rationales(
         self, rationales: List[str], generations_column: str = "generations"

src/distilabel/tasks/preference/ultrafeedback.py

@@ -22,6 +22,7 @@
     List,
     Optional,
     TypedDict,
+    Union,
 )
 
 from distilabel.tasks.base import get_template
@@ -30,6 +31,9 @@
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentence_transformers import (
+        SentenceTransformersExtractor,
+    )
 
 _ULTRAFEEDBACK_TEMPLATE = get_template("ultrafeedback.jinja2")
 
@@ -132,13 +136,15 @@ def to_argilla_dataset(
         ratings_column: str = "rating",
         rationale_column: str = "rationale",
         ratings_values: Optional[List[int]] = None,
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
     ) -> "FeedbackDataset":
         return super().to_argilla_dataset(
             dataset_row=dataset_row,
             generations_column=generations_column,
             ratings_column=ratings_column,
             rationale_column=rationale_column,
             ratings_values=ratings_values or [1, 2, 3, 4, 5],
+            vector_strategy=vector_strategy,
         )
 
     # Override the default `to_argilla_record` method to provide the `ratings_values` of

src/distilabel/tasks/text_generation/base.py

@@ -29,8 +29,12 @@
 if _ARGILLA_AVAILABLE:
     import argilla as rg
 
+
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 
 @dataclass
@@ -154,6 +158,7 @@ def to_argilla_dataset(
         self,
         dataset_row: Dict[str, Any],
         generations_column: Optional[str] = "generations",
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
     ) -> "FeedbackDataset":
         # First we infer the fields from the input_args_names, but we could also
         # create those manually instead using `rg.TextField(...)`
@@ -201,11 +206,15 @@ def to_argilla_dataset(
                 )
         # Then we just return the `FeedbackDataset` with the fields, questions, and metadata properties
         # defined above.
-        return rg.FeedbackDataset(
+        dataset = rg.FeedbackDataset(
             fields=fields,
             questions=questions,
             metadata_properties=metadata_properties,  # Note that these are always optional
         )
+        dataset = Task.add_vectors_to_argilla_dataset(
+            dataset=dataset, vector_strategy=vector_strategy
+        )
+        return dataset
 
     def to_argilla_record(self, dataset_row: Dict[str, Any]) -> "FeedbackRecord":
         """Converts a dataset row to an Argilla `FeedbackRecord`."""

src/distilabel/tasks/text_generation/self_instruct.py

@@ -15,9 +15,9 @@
 import re
 import warnings
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from distilabel.tasks.base import get_template
+from distilabel.tasks.base import Task, get_template
 from distilabel.tasks.prompt import Prompt
 from distilabel.tasks.text_generation.base import TextGenerationTask
 from distilabel.utils.argilla import (
@@ -31,6 +31,9 @@
 
 if TYPE_CHECKING:
     from argilla import FeedbackDataset, FeedbackRecord
+    from argilla.client.feedback.integrations.sentencetransformers import (
+        SentenceTransformersExtractor,
+    )
 
 _SELF_INSTRUCT_TEMPLATE = get_template("self-instruct.jinja2")
 
@@ -103,7 +106,11 @@ def parse_output(self, output: str) -> Dict[str, List[str]]:
         pattern = re.compile(r"\d+\.\s*(.*?)\n")
         return {"instructions": pattern.findall(output)}
 
-    def to_argilla_dataset(self, dataset_row: Dict[str, Any]) -> "FeedbackDataset":
+    def to_argilla_dataset(
+        self,
+        dataset_row: Dict[str, Any],
+        vector_strategy: Union[bool, "SentenceTransformersExtractor"] = True,
+    ) -> "FeedbackDataset":
         # First we infer the fields from the input_args_names, but we could also
         # create those manually instead using `rg.TextField(...)`
         fields = infer_fields_from_dataset_row(
@@ -149,11 +156,17 @@ def to_argilla_dataset(self, dataset_row: Dict[str, Any]) -> "FeedbackDataset":
         )  # type: ignore
         # Then we just return the `FeedbackDataset` with the fields, questions, and metadata properties
         # defined above.
-        return rg.FeedbackDataset(
+        dataset = rg.FeedbackDataset(
             fields=fields,
             questions=questions,  # type: ignore
             metadata_properties=metadata_properties,  # Note that these are always optional
         )
+        dataset: (
+            FeedbackRecord | List[FeedbackRecord] | FeedbackDataset
+        ) = Task.add_vectors_to_argilla_dataset(
+            dataset=dataset, vector_strategy=vector_strategy
+        )
+        return dataset
 
     def to_argilla_record(
         self,

src/distilabel/utils/imports.py

@@ -91,7 +91,9 @@ def _check_package_is_available(
 
 
 _ARGILLA_AVAILABLE = _check_package_is_available(
-    "argilla", min_version="1.16.0", greater_or_equal=True
+    "argilla", min_version="1.22.0", greater_or_equal=True
+) and _check_package_is_available(
+    "sentence-transformers", min_version="2.0.0", greater_or_equal=True
 )
 _OPENAI_AVAILABLE = _check_package_is_available(
     "openai", min_version="1.0.0", greater_or_equal=True