zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/pigeon/flavors/pigeon_annotator_flavor.py

@@ -0,0 +1,104 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""Pigeon annotator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.annotators.base_annotator import (
+    BaseAnnotatorConfig,
+    BaseAnnotatorFlavor,
+)
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.pigeon import PIGEON_ANNOTATOR_FLAVOR
+from zenml.stack.authentication_mixin import AuthenticationConfigMixin
+
+if TYPE_CHECKING:
+    from zenml.integrations.pigeon.annotators import PigeonAnnotator
+
+
+class PigeonAnnotatorSettings(BaseSettings):
+    """Settings for the Pigeon annotator."""
+
+
+class PigeonAnnotatorConfig(
+    BaseAnnotatorConfig, PigeonAnnotatorSettings, AuthenticationConfigMixin
+):
+    """Config for the Pigeon annotator.
+
+    Attributes:
+        output_dir: The directory to store the annotations.
+        notebook_only: Whether the annotator only works within a notebook.
+    """
+
+    output_dir: str = "annotations"
+
+
+class PigeonAnnotatorFlavor(BaseAnnotatorFlavor):
+    """Pigeon annotator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return PIGEON_ANNOTATOR_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/annotator/pigeon.png"
+
+    @property
+    def config_class(self) -> Type[PigeonAnnotatorConfig]:
+        """Returns `PigeonAnnotatorConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return PigeonAnnotatorConfig
+
+    @property
+    def implementation_class(self) -> Type["PigeonAnnotator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.pigeon.annotators import PigeonAnnotator
+
+        return PigeonAnnotator

zenml/integrations/pillow/materializers/pillow_image_materializer.py

@@ -14,7 +14,6 @@
 """Materializer for Pillow Image objects."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 from PIL import Image
@@ -57,16 +56,15 @@ class PillowImageMaterializer(BaseMaterializer):
         files = io_utils.find_files(self.uri, f"{DEFAULT_IMAGE_FILENAME}.*")
         filepath = [file for file in files if not fileio.isdir(file)][0]
 
-        # create a temporary folder
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        temp_file = os.path.join(
-            temp_dir.name,
-            f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
-        )
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            temp_file = os.path.join(
+                temp_dir,
+                f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
+            )
 
-        # copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        return Image.open(temp_file)
+            # copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            return Image.open(temp_file)
 
     def save(self, image: Image.Image) -> None:
         """Write to artifact store.
@@ -74,18 +72,17 @@ class PillowImageMaterializer(BaseMaterializer):
         Args:
             image: An Image.Image object.
         """
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        file_extension = image.format or DEFAULT_IMAGE_EXTENSION
-        full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
-        temp_image_path = os.path.join(temp_dir.name, full_filename)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            file_extension = image.format or DEFAULT_IMAGE_EXTENSION
+            full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
+            temp_image_path = os.path.join(temp_dir, full_filename)
 
-        # save the image in a temporary directory
-        image.save(temp_image_path)
+            # save the image in a temporary directory
+            image.save(temp_image_path)
 
-        # copy the saved image to the artifact store
-        artifact_store_path = os.path.join(self.uri, full_filename)
-        io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
-        temp_dir.cleanup()
+            # copy the saved image to the artifact store
+            artifact_store_path = os.path.join(self.uri, full_filename)
+            io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
 
     def save_visualizations(
         self, image: Image.Image

zenml/integrations/polars/__init__.py

@@ -25,6 +25,7 @@ class PolarsIntegration(Integration):
         "polars>=0.19.5",
         "pyarrow>=12.0.0"
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyarrow"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/polars/materializers/dataframe_materializer.py

@@ -14,7 +14,6 @@
 """Polars materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type, Union
 
 import polars as pl
@@ -22,7 +21,6 @@ import pyarrow as pa  # type: ignore
 import pyarrow.parquet as pq  # type: ignore
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -45,35 +43,29 @@ class PolarsMaterializer(BaseMaterializer):
         Returns:
             A Polars data frame or series.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            io_utils.copy_dir(self.uri, temp_dir)
 
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the data from the temporary directory
-        table = pq.read_table(
-            os.path.join(temp_dir.name, "dataframe.parquet").replace("\\", "/")
-        )
-
-        # If the data is of type pl.Series, convert it back to a pyarrow array
-        # instead of a table.
-        if (
-            table.schema.metadata
-            and b"zenml_is_pl_series" in table.schema.metadata
-        ):
-            isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
-            isinstance_series = bool.from_bytes(isinstance_bytes, "big")
-            if isinstance_series:
-                table = table.column(0)
+            # Load the data from the temporary directory
+            table = pq.read_table(
+                os.path.join(temp_dir, "dataframe.parquet").replace("\\", "/")
+            )
 
-        # Convert the table to a Polars data frame or series
-        data = pl.from_arrow(table)
+            # If the data is of type pl.Series, convert it back to a pyarrow array
+            # instead of a table.
+            if (
+                table.schema.metadata
+                and b"zenml_is_pl_series" in table.schema.metadata
+            ):
+                isinstance_bytes = table.schema.metadata[b"zenml_is_pl_series"]
+                isinstance_series = bool.from_bytes(isinstance_bytes, "big")
+                if isinstance_series:
+                    table = table.column(0)
 
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
+            # Convert the table to a Polars data frame or series
+            data = pl.from_arrow(table)
 
-        return data
+            return data
 
     def save(self, data: Union[pl.DataFrame, pl.Series]) -> None:
         """Writes Polars data to the artifact store.
@@ -107,15 +99,10 @@ class PolarsMaterializer(BaseMaterializer):
             {b"zenml_is_pl_series": isinstance_bytes}
         )
 
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Write the table to a Parquet file
-        path = os.path.join(temp_dir.name, "dataframe.parquet").replace(
-            "\\", "/"
-        )
-        pq.write_table(table, path)  # Uses lz4 compression by default
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            # Write the table to a Parquet file
+            path = os.path.join(temp_dir, "dataframe.parquet").replace(
+                "\\", "/"
+            )
+            pq.write_table(table, path)  # Uses lz4 compression by default
+            io_utils.copy_dir(temp_dir, self.uri)

zenml/integrations/prodigy/__init__.py

@@ -0,0 +1,48 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""Initialization of the Prodigy integration."""
+from typing import List, Type
+
+from zenml.integrations.constants import PRODIGY
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+PRODIGY_ANNOTATOR_FLAVOR = "prodigy"
+
+
+class ProdigyIntegration(Integration):
+    """Definition of Prodigy integration for ZenML."""
+
+    NAME = PRODIGY
+    REQUIREMENTS = [
+        "prodigy",
+        "urllib3<2",
+    ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["urllib3"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Prodigy integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.prodigy.flavors import (
+            ProdigyAnnotatorFlavor,
+        )
+
+        return [ProdigyAnnotatorFlavor]
+
+
+ProdigyIntegration.check_installation()

zenml/integrations/prodigy/annotators/__init__.py

@@ -0,0 +1,20 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""Initialization of the Prodigy annotators submodule."""
+
+from zenml.integrations.prodigy.annotators.prodigy_annotator import (
+    ProdigyAnnotator,
+)
+
+__all__ = ["ProdigyAnnotator"]

zenml/integrations/prodigy/annotators/prodigy_annotator.py

@@ -0,0 +1,275 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""Implementation of the Prodigy annotation integration."""
+
+import json
+from typing import Any, List, Optional, Tuple, cast
+
+import prodigy
+from peewee import Database as PeeweeDatabase
+from prodigy.components.db import (
+    Database as ProdigyDatabase,
+)
+from prodigy.components.db import connect
+from prodigy.errors import ProdigyError
+
+from zenml.annotators.base_annotator import BaseAnnotator
+from zenml.integrations.prodigy.flavors.prodigy_annotator_flavor import (
+    ProdigyAnnotatorConfig,
+)
+from zenml.logger import get_logger
+from zenml.stack.authentication_mixin import AuthenticationMixin
+
+logger = get_logger(__name__)
+
+DEFAULT_LOCAL_INSTANCE_HOST = "localhost"
+DEFAULT_LOCAL_PRODIGY_PORT = 8080
+
+
+class ProdigyAnnotator(BaseAnnotator, AuthenticationMixin):
+    """Class to interact with the Prodigy annotation interface."""
+
+    @property
+    def config(self) -> ProdigyAnnotatorConfig:
+        """Returns the `ProdigyAnnotatorConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return cast(ProdigyAnnotatorConfig, self._config)
+
+    def get_url(self) -> str:
+        """Gets the top-level URL of the annotation interface.
+
+        Returns:
+            The URL of the annotation interface.
+        """
+        instance_url = DEFAULT_LOCAL_INSTANCE_HOST
+        port = DEFAULT_LOCAL_PRODIGY_PORT
+        if self.config.custom_config_path:
+            with open(self.config.custom_config_path, "r") as f:
+                config = json.load(f)
+            instance_url = config.get("instance_url", instance_url)
+            port = config.get("port", port)
+        return f"http://{instance_url}:{port}"
+
+    def get_url_for_dataset(self, dataset_name: str) -> str:
+        """Gets the URL of the annotation interface for the given dataset.
+
+        Prodigy does not support dataset-specific URLs, so this method returns
+        the top-level URL since that's what will be served for the user.
+
+        Args:
+            dataset_name: The name of the dataset. (Unuse)
+
+        Returns:
+            The URL of the annotation interface.
+        """
+        return self.get_url()
+
+    def get_datasets(self) -> List[Any]:
+        """Gets the datasets currently available for annotation.
+
+        Returns:
+            A list of datasets (str).
+        """
+        datasets = self._get_db().datasets
+        return cast(List[Any], datasets)
+
+    def get_dataset_names(self) -> List[str]:
+        """Gets the names of the datasets.
+
+        Returns:
+            A list of dataset names.
+        """
+        return self.get_datasets()
+
+    def get_dataset_stats(self, dataset_name: str) -> Tuple[int, int]:
+        """Gets the statistics of the given dataset.
+
+        Args:
+            dataset_name: The name of the dataset.
+
+        Returns:
+            A tuple containing (labeled_task_count, unlabeled_task_count) for
+                the dataset.
+
+        Raises:
+            IndexError: If the dataset does not exist.
+        """
+        db = self._get_db()
+        try:
+            labeled_data_count = db.count_dataset(name=dataset_name)
+        except ValueError as e:
+            raise IndexError(
+                f"Dataset {dataset_name} does not exist. Please use `zenml "
+                f"annotator dataset list` to list the available datasets."
+            ) from e
+        return (labeled_data_count, 0)
+
+    def launch(self, **kwargs: Any) -> None:
+        """Launches the annotation interface.
+
+        This method extracts the 'command' and additional config
+            parameters from kwargs.
+
+        Args:
+            **kwargs: Should include:
+                - command: The full recipe command without "prodigy".
+                - Any additional config parameters to overwrite the
+                    project-specific, global, and recipe config.
+
+        Raises:
+            ValueError: If the 'command' keyword argument is not provided.
+        """
+        command = kwargs.get("command")
+        if not command:
+            raise ValueError(
+                "The 'command' keyword argument is required for launching Prodigy."
+            )
+
+        # Remove 'command' from kwargs to pass the rest as config parameters
+        config = {
+            key: value for key, value in kwargs.items() if key != "command"
+        }
+        prodigy.serve(command=command, **config)
+
+    def _get_db(
+        self,
+        custom_database: PeeweeDatabase = None,
+        display_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+    ) -> ProdigyDatabase:
+        """Gets Prodigy database / client.
+
+        Args:
+            custom_database: Custom database to use.
+            display_id: The display id of the database.
+            display_name: The display name of the database.
+
+        Returns:
+            Prodigy database client.
+        """
+        db_kwargs = {}
+        if custom_database:
+            db_kwargs["db"] = custom_database
+        if display_id:
+            db_kwargs["display_id"] = display_id
+        if display_name:
+            db_kwargs["display_name"] = display_name
+
+        # database is passed in without the keyword argument
+        if custom_database:
+            return connect(custom_database, **db_kwargs)
+        return connect(**db_kwargs)
+
+    def add_dataset(self, **kwargs: Any) -> Any:
+        """Registers a dataset for annotation.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Prodigy client.
+
+        Returns:
+            A Prodigy list representing the dataset.
+
+        Raises:
+            ValueError: if 'dataset_name' and 'label_config' aren't provided.
+        """
+        db = self._get_db()
+        dataset_kwargs = {"dataset_name": kwargs.get("dataset_name")}
+        if not dataset_kwargs["dataset_name"]:
+            raise ValueError("`dataset_name` keyword argument is required.")
+
+        if kwargs.get("dataset_meta"):
+            dataset_kwargs["dataset_meta"] = kwargs.get("dataset_meta")
+        return db.add_dataset(**dataset_kwargs)
+
+    def delete_dataset(self, **kwargs: Any) -> None:
+        """Deletes a dataset from the annotation interface.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Prodigy
+                client.
+
+        Raises:
+            ValueError: If the dataset name is not provided or if the dataset
+                does not exist.
+        """
+        db = self._get_db()
+        if not (dataset_name := kwargs.get("dataset_name")):
+            raise ValueError("`dataset_name` keyword argument is required.")
+        try:
+            db.drop_dataset(name=dataset_name)
+        except ProdigyError as e:
+            # see https://support.prodi.gy/t/how-to-import-datasetdoesnotexist-error/7205
+            if type(e).__name__ == "DatasetNotFound":
+                raise ValueError(
+                    f"Dataset name '{dataset_name}' does not exist."
+                ) from e
+
+    def get_dataset(self, **kwargs: Any) -> Any:
+        """Gets the dataset metadata for the given name.
+
+        If you would like the labeled data, use `get_labeled_data` instead.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Prodigy client.
+
+        Returns:
+            The metadata associated with a Prodigy dataset
+
+        Raises:
+            ValueError: If the dataset name is not provided or if the dataset
+                does not exist.
+        """
+        db = self._get_db()
+        if dataset_name := kwargs.get("dataset_name"):
+            try:
+                return db.get_meta(name=dataset_name)
+            except Exception as e:
+                raise ValueError(
+                    f"Dataset name '{dataset_name}' does not exist."
+                ) from e
+
+    def get_labeled_data(self, **kwargs: Any) -> Any:
+        """Gets the labeled data for the given dataset.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Prodigy client.
+
+        Returns:
+            A list of all examples in the dataset serialized to the
+                Prodigy Task format.
+
+        Raises:
+            ValueError: If the dataset name is not provided or if the dataset
+                does not exist.
+        """
+        if dataset_name := kwargs.get("dataset_name"):
+            return self._get_db().get_dataset_examples(dataset_name)
+        else:
+            raise ValueError("`dataset_name` keyword argument is required.")
+
+    def get_unlabeled_data(self, **kwargs: str) -> Any:
+        """Gets the unlabeled data for the given dataset.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Prodigy client.
+
+        Raises:
+            NotImplementedError: Prodigy doesn't allow fetching unlabeled data.
+        """
+        raise NotImplementedError(
+            "Prodigy doesn't allow fetching unlabeled data."
+        )

zenml/integrations/prodigy/flavors/__init__.py

@@ -0,0 +1,24 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""Prodigy integration flavors."""
+
+from zenml.integrations.prodigy.flavors.prodigy_annotator_flavor import (
+    ProdigyAnnotatorConfig,
+    ProdigyAnnotatorFlavor,
+)
+
+__all__ = [
+    "ProdigyAnnotatorFlavor",
+    "ProdigyAnnotatorConfig",
+]