genesis-flow 1.0.0__py3-none-any.whl

mlflow/data/numpy_dataset.py

@@ -0,0 +1,223 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import compute_numpy_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class NumpyDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a NumPy dataset for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        features: Union[np.ndarray, dict[str, np.ndarray]],
+        source: DatasetSource,
+        targets: Union[np.ndarray, dict[str, np.ndarray]] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+    ):
+        """
+        Args:
+            features: A numpy array or dictionary of numpy arrays containing dataset features.
+            source: The source of the numpy dataset.
+            targets: A numpy array or dictionary of numpy arrays containing dataset targets.
+                Optional.
+            name: The name of the dataset. E.g. "wiki_train". If unspecified, a name is
+                automatically generated.
+            digest: The digest (hash, fingerprint) of the dataset. If unspecified, a digest
+                is automatically computed.
+        """
+        self._features = features
+        self._targets = targets
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_digest(self) -> str:
+        """
+        Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+        """
+        return compute_numpy_digest(self._features, self._targets)
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type, schema, and profile.
+        """
+        schema = json.dumps(self.schema.to_dict()) if self.schema else None
+        config = super().to_dict()
+        config.update(
+            {
+                "schema": schema,
+                "profile": json.dumps(self.profile),
+            }
+        )
+        return config
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def features(self) -> Union[np.ndarray, dict[str, np.ndarray]]:
+        """
+        The features of the dataset.
+        """
+        return self._features
+
+    @property
+    def targets(self) -> Optional[Union[np.ndarray, dict[str, np.ndarray]]]:
+        """
+        The targets of the dataset. May be ``None`` if no targets are available.
+        """
+        return self._targets
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be ``None`` if a profile cannot be computed.
+        """
+
+        def get_profile_attribute(numpy_data, attr_name):
+            if isinstance(numpy_data, dict):
+                return {key: getattr(array, attr_name) for key, array in numpy_data.items()}
+            else:
+                return getattr(numpy_data, attr_name)
+
+        profile = {
+            "features_shape": get_profile_attribute(self._features, "shape"),
+            "features_size": get_profile_attribute(self._features, "size"),
+            "features_nbytes": get_profile_attribute(self._features, "nbytes"),
+        }
+        if self._targets is not None:
+            profile.update(
+                {
+                    "targets_shape": get_profile_attribute(self._targets, "shape"),
+                    "targets_size": get_profile_attribute(self._targets, "size"),
+                    "targets_nbytes": get_profile_attribute(self._targets, "nbytes"),
+                }
+            )
+
+        return profile
+
+    @cached_property
+    def schema(self) -> Optional[TensorDatasetSchema]:
+        """
+        MLflow TensorSpec schema representing the dataset features and targets (optional).
+        """
+        try:
+            features_schema = _infer_schema(self._features)
+            targets_schema = None
+            if self._targets is not None:
+                targets_schema = _infer_schema(self._targets)
+            return TensorDatasetSchema(features=features_schema, targets=targets_schema)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for NumPy dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        return PyFuncInputsOutputs(self._features, self._targets)
+
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation. Required
+        for use with mlflow.sklearn.evaluate().
+        """
+        return EvaluationDataset(
+            data=self._features,
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+            name=self.name,
+            digest=self.digest,
+        )
+
+
+def from_numpy(
+    features: Union[np.ndarray, dict[str, np.ndarray]],
+    source: Union[str, DatasetSource] = None,
+    targets: Union[np.ndarray, dict[str, np.ndarray]] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+) -> NumpyDataset:
+    """
+    Constructs a :py:class:`NumpyDataset <mlflow.data.numpy_dataset.NumpyDataset>` object from
+    NumPy features, optional targets, and source. If the source is path like, then this will
+    construct a DatasetSource object from the source path. Otherwise, the source is assumed to
+    be a DatasetSource object.
+
+    Args:
+        features: NumPy features, represented as an np.ndarray or dictionary of named np.ndarrays.
+        source: The source from which the numpy data was derived, e.g. a filesystem path, an S3 URI,
+            an HTTPS URL, a delta table name with version, or spark table etc. ``source`` may be
+            specified as a URI, a path-like string, or an instance of
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`. If unspecified,
+            the source is assumed to be the code location (e.g. notebook cell, script, etc.) where
+            :py:func:`from_numpy <mlflow.data.from_numpy>` is being called.
+        targets: Optional NumPy targets, represented as an np.ndarray or dictionary of named
+            np.ndarrays.
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: The dataset digest (hash). If unspecified, a digest is computed automatically.
+
+    .. code-block:: python
+        :test:
+        :caption: Basic Example
+
+        import mlflow
+        import numpy as np
+
+        x = np.random.uniform(size=[2, 5, 4])
+        y = np.random.randint(2, size=[2])
+        dataset = mlflow.data.from_numpy(x, targets=y)
+
+    .. code-block:: python
+        :test:
+        :caption: Dict Example
+
+        import mlflow
+        import numpy as np
+
+        x = {
+            "feature_1": np.random.uniform(size=[2, 5, 4]),
+            "feature_2": np.random.uniform(size=[2, 5, 4]),
+        }
+        y = np.random.randint(2, size=[2])
+        dataset = mlflow.data.from_numpy(x, targets=y)
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return NumpyDataset(
+        features=features, source=resolved_source, targets=targets, name=name, digest=digest
+    )

mlflow/data/pandas_dataset.py

@@ -0,0 +1,231 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import pandas as pd
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import compute_pandas_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.types import Schema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class PandasDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a Pandas DataFrame for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        source: DatasetSource,
+        targets: Optional[str] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+        predictions: Optional[str] = None,
+    ):
+        """
+        Args:
+            df: A pandas DataFrame.
+            source: The source of the pandas DataFrame.
+            targets: The name of the target column. Optional.
+            name: The name of the dataset. E.g. "wiki_train". If unspecified, a name is
+                automatically generated.
+            digest: The digest (hash, fingerprint) of the dataset. If unspecified, a digest
+                is automatically computed.
+            predictions: Optional. The name of the column containing model predictions,
+                if the dataset contains model predictions. If specified, this column
+                must be present in the dataframe (``df``).
+        """
+        if targets is not None and targets not in df.columns:
+            raise MlflowException(
+                f"The specified pandas DataFrame does not contain the specified targets column"
+                f" '{targets}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+        if predictions is not None and predictions not in df.columns:
+            raise MlflowException(
+                f"The specified pandas DataFrame does not contain the specified predictions column"
+                f" '{predictions}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+        self._df = df
+        self._targets = targets
+        self._predictions = predictions
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_digest(self) -> str:
+        """
+        Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+        """
+        return compute_pandas_digest(self._df)
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type, schema, and profile.
+        """
+        schema = json.dumps({"mlflow_colspec": self.schema.to_dict()}) if self.schema else None
+        config = super().to_dict()
+        config.update(
+            {
+                "schema": schema,
+                "profile": json.dumps(self.profile),
+            }
+        )
+        return config
+
+    @property
+    def df(self) -> pd.DataFrame:
+        """
+        The underlying pandas DataFrame.
+        """
+        return self._df
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def targets(self) -> Optional[str]:
+        """
+        The name of the target column. May be ``None`` if no target column is available.
+        """
+        return self._targets
+
+    @property
+    def predictions(self) -> Optional[str]:
+        """
+        The name of the predictions column. May be ``None`` if no predictions column is available.
+        """
+        return self._predictions
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be ``None`` if a profile cannot be computed.
+        """
+        return {
+            "num_rows": len(self._df),
+            "num_elements": int(self._df.size),
+        }
+
+    @cached_property
+    def schema(self) -> Optional[Schema]:
+        """
+        An instance of :py:class:`mlflow.types.Schema` representing the tabular dataset. May be
+        ``None`` if the schema cannot be inferred from the dataset.
+        """
+        try:
+            return _infer_schema(self._df)
+        except Exception as e:
+            _logger.debug("Failed to infer schema for Pandas dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        if self._targets:
+            inputs = self._df.drop(columns=[self._targets])
+            outputs = self._df[self._targets]
+            return PyFuncInputsOutputs(inputs, outputs)
+        else:
+            return PyFuncInputsOutputs(self._df)
+
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation. Required
+        for use with mlflow.evaluate().
+        """
+        return EvaluationDataset(
+            data=self._df,
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+            predictions=self._predictions,
+            name=self.name,
+            digest=self.digest,
+        )
+
+
+def from_pandas(
+    df: pd.DataFrame,
+    source: Union[str, DatasetSource] = None,
+    targets: Optional[str] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+    predictions: Optional[str] = None,
+) -> PandasDataset:
+    """
+    Constructs a :py:class:`PandasDataset <mlflow.data.pandas_dataset.PandasDataset>` instance from
+    a Pandas DataFrame, optional targets, optional predictions, and source.
+
+    Args:
+        df: A Pandas DataFrame.
+        source: The source from which the DataFrame was derived, e.g. a filesystem
+            path, an S3 URI, an HTTPS URL, a delta table name with version, or
+            spark table etc. ``source`` may be specified as a URI, a path-like string,
+            or an instance of
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+            If unspecified, the source is assumed to be the code location
+            (e.g. notebook cell, script, etc.) where
+            :py:func:`from_pandas <mlflow.data.from_pandas>` is being called.
+        targets: An optional target column name for supervised training. This column
+            must be present in the dataframe (``df``).
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: The dataset digest (hash). If unspecified, a digest is computed
+            automatically.
+        predictions: An optional predictions column name for model evaluation. This column
+            must be present in the dataframe (``df``).
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+        import pandas as pd
+
+        x = pd.DataFrame(
+            [["tom", 10, 1, 1], ["nick", 15, 0, 1], ["july", 14, 1, 1]],
+            columns=["Name", "Age", "Label", "ModelOutput"],
+        )
+        dataset = mlflow.data.from_pandas(x, targets="Label", predictions="ModelOutput")
+    """
+
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return PandasDataset(
+        df=df,
+        source=resolved_source,
+        targets=targets,
+        name=name,
+        digest=digest,
+        predictions=predictions,
+    )