snowflake-ml-python 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

snowflake/ml/modeling/_internal/model_specifications.py

@@ -0,0 +1,146 @@
+import inspect
+from typing import List
+
+import cloudpickle as cp
+import numpy as np
+
+from snowflake.ml._internal.exceptions import error_codes, exceptions
+
+
+class ModelSpecifications:
+    """
+    A dataclass to define model based specifications like required imports, and package dependencies for Sproc/Udfs.
+    """
+
+    def __init__(self, imports: List[str], pkgDependencies: List[str]) -> None:
+        self.imports = imports
+        self.pkgDependencies = pkgDependencies
+
+
+class SKLearnModelSpecifications(ModelSpecifications):
+    def __init__(self) -> None:
+        import sklearn
+
+        imports: List[str] = ["sklearn"]
+        # TODO(snandamuri): Replace cloudpickle with joblib after latest version of joblib is added to snowflake conda.
+        pkgDependencies = [
+            f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
+            f"cloudpickle=={cp.__version__}",
+        ]
+
+        # A change from previous implementation.
+        # When reusing the Sprocs for all the fit() call in the session, the static dpendencies list should include
+        # all the possible dependencies required during the lifetime.
+
+        # Include XGBoost in the dependencies if it is installed.
+        try:
+            import xgboost
+        except ModuleNotFoundError:
+            pass
+        else:
+            pkgDependencies.append(f"xgboost=={xgboost.__version__}")
+
+        # Include lightgbm in the dependencies if it is installed.
+        try:
+            import lightgbm
+        except ModuleNotFoundError:
+            pass
+        else:
+            pkgDependencies.append(f"lightgbm=={lightgbm.__version__}")
+
+        super().__init__(imports=imports, pkgDependencies=pkgDependencies)
+
+
+class XGBoostModelSpecifications(ModelSpecifications):
+    def __init__(self) -> None:
+        import xgboost
+
+        imports: List[str] = ["xgboost"]
+        pkgDependencies: List[str] = [
+            f"numpy=={np.__version__}",
+            f"xgboost=={xgboost.__version__}",
+            f"cloudpickle=={cp.__version__}",
+        ]
+        super().__init__(imports=imports, pkgDependencies=pkgDependencies)
+
+
+class LightGBMModelSpecifications(ModelSpecifications):
+    def __init__(self) -> None:
+        import lightgbm
+
+        imports: List[str] = ["lightgbm"]
+        pkgDependencies: List[str] = [
+            f"numpy=={np.__version__}",
+            f"lightgbm=={lightgbm.__version__}",
+            f"cloudpickle=={cp.__version__}",
+        ]
+        super().__init__(imports=imports, pkgDependencies=pkgDependencies)
+
+
+class SklearnModelSelectionModelSpecifications(ModelSpecifications):
+    def __init__(self) -> None:
+        import sklearn
+        import xgboost
+
+        imports: List[str] = ["sklearn", "xgboost"]
+        pkgDependencies: List[str] = [
+            f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
+            f"cloudpickle=={cp.__version__}",
+            f"xgboost=={xgboost.__version__}",
+        ]
+
+        # Only include lightgbm in the dependencies if it is installed.
+        try:
+            import lightgbm
+        except ModuleNotFoundError:
+            pass
+        else:
+            imports.append("lightgbm")
+            pkgDependencies.append(f"lightgbm=={lightgbm.__version__}")
+
+        super().__init__(imports=imports, pkgDependencies=pkgDependencies)
+
+
+class ModelSpecificationsBuilder:
+    """
+    A factory class to build ModelSpecifications object for different types of models.
+    """
+
+    @classmethod
+    def build(cls, model: object) -> ModelSpecifications:
+        """
+        A static factory method that builds ModelSpecifications object based on the module name of native model object.
+
+        Args:
+            model: Native model object to be trained.
+
+        Returns:
+            Appropriate ModelSpecification object
+
+        Raises:
+            SnowflakeMLException: Raises an exception the module of given model can't be determined.
+            TypeError: Raises the exception for unsupported modules.
+        """
+        module = inspect.getmodule(model)
+        if module is None:
+            raise exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_TYPE,
+                original_exception=ValueError("Unable to infer model type of the given native model object."),
+            )
+        root_module_name = module.__name__.split(".")[0]
+        if root_module_name == "sklearn":
+            from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
+
+            if isinstance(model, GridSearchCV) or isinstance(model, RandomizedSearchCV):
+                return SklearnModelSelectionModelSpecifications()
+            return SKLearnModelSpecifications()
+        elif root_module_name == "xgboost":
+            return XGBoostModelSpecifications()
+        elif root_module_name == "lightgbm":
+            return LightGBMModelSpecifications()
+        else:
+            raise TypeError(
+                f"Unexpected module type: {root_module_name}." "Supported module types: sklearn, xgboost, lightgbm."
+            )

snowflake/ml/modeling/_internal/model_trainer.py

@@ -0,0 +1,13 @@
+from typing import Protocol
+
+
+class ModelTrainer(Protocol):
+    """
+    Interface for model trainer implementations.
+
+    There are multiple flavors of training like training with pandas datasets, training with
+    Snowpark datasets using sprocs, and out of core training with Snowpark datasets etc.
+    """
+
+    def train(self) -> object:
+        raise NotImplementedError

snowflake/ml/modeling/_internal/model_trainer_builder.py

@@ -0,0 +1,78 @@
+from typing import List, Optional, Union
+
+import pandas as pd
+from sklearn import model_selection
+
+from snowflake.ml.modeling._internal.distributed_hpo_trainer import (
+    DistributedHPOTrainer,
+)
+from snowflake.ml.modeling._internal.estimator_utils import is_single_node
+from snowflake.ml.modeling._internal.model_trainer import ModelTrainer
+from snowflake.ml.modeling._internal.pandas_trainer import PandasModelTrainer
+from snowflake.ml.modeling._internal.snowpark_trainer import SnowparkModelTrainer
+from snowflake.snowpark import DataFrame, Session
+
+_PROJECT = "ModelDevelopment"
+
+
+class ModelTrainerBuilder:
+    """
+    A builder class to create instances of ModelTrainer for different models and training conditions.
+
+    This class provides methods to build instances of ModelTrainer tailored to specific machine learning
+    models and training configurations like dataset's location etc. It abstracts the creation process,
+    allowing the user to obtain a configured ModelTrainer for a particular model architecture or configuration.
+    """
+
+    _ENABLE_DISTRIBUTED = True
+
+    @classmethod
+    def _check_if_distributed_hpo_enabled(cls, session: Session) -> bool:
+        return not is_single_node(session) and ModelTrainerBuilder._ENABLE_DISTRIBUTED is True
+
+    @classmethod
+    def build(
+        cls,
+        estimator: object,
+        dataset: Union[DataFrame, pd.DataFrame],
+        input_cols: Optional[List[str]] = None,
+        label_cols: Optional[List[str]] = None,
+        sample_weight_col: Optional[str] = None,
+        autogenerated: bool = False,
+        subproject: str = "",
+    ) -> ModelTrainer:
+        """
+        Builder method that creates an approproiate ModelTrainer instance based on the given params.
+        """
+        assert input_cols is not None  # Make MyPy happpy
+        if isinstance(dataset, pd.DataFrame):
+            return PandasModelTrainer(
+                estimator=estimator,
+                dataset=dataset,
+                input_cols=input_cols,
+                label_cols=label_cols,
+                sample_weight_col=sample_weight_col,
+            )
+        elif isinstance(dataset, DataFrame):
+            trainer_klass = SnowparkModelTrainer
+            assert dataset._session is not None  # Make MyPy happpy
+            if isinstance(estimator, model_selection.GridSearchCV) or isinstance(
+                estimator, model_selection.RandomizedSearchCV
+            ):
+                if ModelTrainerBuilder._check_if_distributed_hpo_enabled(session=dataset._session):
+                    trainer_klass = DistributedHPOTrainer
+            return trainer_klass(
+                estimator=estimator,
+                dataset=dataset,
+                session=dataset._session,
+                input_cols=input_cols,
+                label_cols=label_cols,
+                sample_weight_col=sample_weight_col,
+                autogenerated=autogenerated,
+                subproject=subproject,
+            )
+        else:
+            raise TypeError(
+                f"Unexpected dataset type: {type(dataset)}."
+                "Supported dataset types: snowpark.DataFrame, pandas.DataFrame."
+            )

snowflake/ml/modeling/_internal/pandas_trainer.py

@@ -0,0 +1,54 @@
+import inspect
+from typing import List, Optional
+
+import pandas as pd
+
+
+class PandasModelTrainer:
+    """
+    A class for training machine learning models using Pandas datasets.
+    """
+
+    def __init__(
+        self,
+        estimator: object,
+        dataset: pd.DataFrame,
+        input_cols: List[str],
+        label_cols: Optional[List[str]],
+        sample_weight_col: Optional[str],
+    ) -> None:
+        """
+        Initializes the PandasModelTrainer with a model, a Pandas DataFrame, feature, and label column names.
+
+        Args:
+            estimator: SKLearn compatible estimator or transformer object.
+            dataset: The dataset used for training the model.
+            input_cols: The name(s) of one or more columns in a DataFrame containing a feature to be used for training.
+            label_cols: The name(s) of one or more columns in a DataFrame representing the target variable(s) to learn.
+            sample_weight_col: The column name representing the weight of training examples.
+        """
+        self.estimator = estimator
+        self.dataset = dataset
+        self.input_cols = input_cols
+        self.label_cols = label_cols
+        self.sample_weight_col = sample_weight_col
+
+    def train(self) -> object:
+        """
+        Trains the model using specified features and target columns from the dataset.
+
+        Returns:
+            Trained model
+        """
+        assert hasattr(self.estimator, "fit")  # Keep mypy happy
+        argspec = inspect.getfullargspec(self.estimator.fit)
+        args = {"X": self.dataset[self.input_cols]}
+
+        if self.label_cols:
+            label_arg_name = "Y" if "Y" in argspec.args else "y"
+            args[label_arg_name] = self.dataset[self.label_cols].squeeze()
+
+        if self.sample_weight_col is not None and "sample_weight" in argspec.args:
+            args["sample_weight"] = self.dataset[self.sample_weight_col].squeeze()
+
+        return self.estimator.fit(**args)