snowflake-ml-python 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -2,7 +2,6 @@ import os
 import pathlib
 import sys
 import tempfile
-import warnings
 import zipfile
 from contextlib import contextmanager
 from datetime import datetime
@@ -18,7 +17,6 @@ from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import (
-    _core_requirements,
     _packaging_requirements,
     model_blob_meta,
     model_meta_schema,
@@ -29,14 +27,10 @@ from snowflake.ml.model._packager.model_runtime import model_runtime
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
 
-_PACKAGING_CORE_DEPENDENCIES = [
-    str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
-    for r in _core_requirements.REQUIREMENTS
-]  # Legacy Model only
 _PACKAGING_REQUIREMENTS = [
     str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
     for r in _packaging_requirements.REQUIREMENTS
-]  # New Model only
+]
 _SNOWFLAKE_PKG_NAME = "snowflake"
 _SNOWFLAKE_ML_PKG_NAME = f"{_SNOWFLAKE_PKG_NAME}.ml"
 
@@ -55,6 +49,7 @@ def create_model_metadata(
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
     python_version: Optional[str] = None,
+    task: model_types.Task = model_types.Task.UNKNOWN,
     **kwargs: Any,
 ) -> Generator["ModelMetadata", None, None]:
     """Create a generator for model metadata object. Use generator to ensure correct register and unregister for
@@ -74,6 +69,9 @@ def create_model_metadata(
         pip_requirements: List of pip Python packages requirements for running the model. Defaults to None.
         python_version: A string of python version where model is run. Used for user override. If specified as None,
             current version would be captured. Defaults to None.
+        task: The task of the Model Version. It is an enum class Task with values TABULAR_REGRESSION,
+            TABULAR_BINARY_CLASSIFICATION, TABULAR_MULTI_CLASSIFICATION, TABULAR_RANKING, or UNKNOWN. By default,
+            it is set to Task.UNKNOWN and may be overridden by inferring from the Model Object.
         **kwargs: Dict of attributes and values of the metadata. Used when loading from file.
 
     Raises:
@@ -84,18 +82,6 @@ def create_model_metadata(
     """
     model_dir_path = os.path.normpath(model_dir_path)
     embed_local_ml_library = kwargs.pop("embed_local_ml_library", False)
-    legacy_save = kwargs.pop("_legacy_save", False)
-    if "relax_version" not in kwargs:
-        warnings.warn(
-            (
-                "`relax_version` is not set and therefore defaulted to True. Dependency version constraints relaxed "
-                "from ==x.y.z to >=x.y, <(x+1). To use specific dependency versions for compatibility, "
-                "reproducibility, etc., set `options={'relax_version': False}` when logging the model."
-            ),
-            category=UserWarning,
-            stacklevel=2,
-        )
-    relax_version = kwargs.pop("relax_version", True)
 
     if embed_local_ml_library:
         # Use the last one which is loaded first, that is mean, it is loaded from site-packages.
@@ -118,7 +104,6 @@ def create_model_metadata(
         pip_requirements=pip_requirements,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
-        legacy_save=legacy_save,
     )
 
     if embed_local_ml_library:
@@ -131,17 +116,13 @@ def create_model_metadata(
         model_type=model_type,
         signatures=signatures,
         function_properties=function_properties,
+        task=task,
     )
 
     code_dir_path = os.path.join(model_dir_path, MODEL_CODE_DIR)
-    if (embed_local_ml_library and legacy_save) or code_paths:
+    if code_paths:
         os.makedirs(code_dir_path, exist_ok=True)
 
-    if embed_local_ml_library and legacy_save:
-        snowml_path_in_code = os.path.join(code_dir_path, _SNOWFLAKE_PKG_NAME)
-        os.makedirs(snowml_path_in_code, exist_ok=True)
-        file_utils.copy_file_or_tree(path_to_copy, snowml_path_in_code)
-
     if code_paths:
         for code_path in code_paths:
             # This part is to prevent users from providing code following our naming and overwrite our code.
@@ -160,8 +141,6 @@ def create_model_metadata(
                     cloudpickle.register_pickle_by_value(mod)
                     imported_modules.append(mod)
         yield model_meta
-        if relax_version:
-            model_meta.env.relax_version()
         model_meta.save(model_dir_path)
     finally:
         for mod in imported_modules:
@@ -174,7 +153,6 @@ def _create_env_for_model_metadata(
     pip_requirements: Optional[List[str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
-    legacy_save: bool = False,
 ) -> model_env.ModelEnv:
     env = model_env.ModelEnv()
 
@@ -184,7 +162,7 @@ def _create_env_for_model_metadata(
     env.python_version = python_version  # type: ignore[assignment]
     env.snowpark_ml_version = snowml_env.VERSION
 
-    requirements_to_add = _PACKAGING_CORE_DEPENDENCIES if legacy_save else _PACKAGING_REQUIREMENTS
+    requirements_to_add = _PACKAGING_REQUIREMENTS
 
     if embed_local_ml_library:
         env.include_if_absent(
@@ -237,7 +215,7 @@ class ModelMetadata:
         function_properties: A dict mapping function names to dict mapping function property key to value.
         metadata: User provided key-value metadata of the model. Defaults to None.
         creation_timestamp: Unix timestamp when the model metadata is created.
-        model_objective: Model objective like regression, classification etc.
+        task: Model task like TABULAR_REGRESSION, tabular_classification, timeseries_forecasting etc.
     """
 
     def telemetry_metadata(self) -> ModelMetadataTelemetryDict:
@@ -261,7 +239,7 @@ class ModelMetadata:
         min_snowpark_ml_version: Optional[str] = None,
         models: Optional[Dict[str, model_blob_meta.ModelBlobMeta]] = None,
         original_metadata_version: Optional[str] = model_meta_schema.MODEL_METADATA_VERSION,
-        model_objective: Optional[model_meta_schema.ModelObjective] = model_meta_schema.ModelObjective.UNKNOWN,
+        task: model_types.Task = model_types.Task.UNKNOWN,
         explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = None,
     ) -> None:
         self.name = name
@@ -287,9 +265,7 @@ class ModelMetadata:
 
         self.original_metadata_version = original_metadata_version
 
-        self.model_objective: model_meta_schema.ModelObjective = (
-            model_objective or model_meta_schema.ModelObjective.UNKNOWN
-        )
+        self.task: model_types.Task = task
         self.explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = explain_algorithm
 
     @property
@@ -306,10 +282,10 @@ class ModelMetadata:
         if self._runtimes and "cpu" in self._runtimes:
             return self._runtimes
         runtimes = {
-            "cpu": model_runtime.ModelRuntime("cpu", self.env),
+            "cpu": model_runtime.ModelRuntime("cpu", self.env, is_warehouse=False),
         }
         if self.env.cuda_version:
-            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_gpu=True)})
+            runtimes.update({"gpu": model_runtime.ModelRuntime("gpu", self.env, is_warehouse=False, is_gpu=True)})
         return runtimes
 
     def save(self, model_dir_path: str) -> None:
@@ -343,7 +319,7 @@ class ModelMetadata:
                 "signatures": {func_name: sig.to_dict() for func_name, sig in self.signatures.items()},
                 "version": model_meta_schema.MODEL_METADATA_VERSION,
                 "min_snowpark_ml_version": self.min_snowpark_ml_version,
-                "model_objective": self.model_objective.value,
+                "task": self.task.value,
                 "explainability": (
                     model_meta_schema.ExplainabilityMetadataDict(algorithm=self.explain_algorithm.value)
                     if self.explain_algorithm
@@ -387,7 +363,7 @@ class ModelMetadata:
             signatures=loaded_meta["signatures"],
             version=original_loaded_meta_version,
             min_snowpark_ml_version=loaded_meta_min_snowpark_ml_version,
-            model_objective=loaded_meta.get("model_objective", model_meta_schema.ModelObjective.UNKNOWN.value),
+            task=loaded_meta.get("task", model_types.Task.UNKNOWN.value),
             explainability=loaded_meta.get("explainability", None),
             function_properties=loaded_meta.get("function_properties", {}),
         )
@@ -442,9 +418,7 @@ class ModelMetadata:
             min_snowpark_ml_version=model_dict["min_snowpark_ml_version"],
             models=models,
             original_metadata_version=model_dict["version"],
-            model_objective=model_meta_schema.ModelObjective(
-                model_dict.get("model_objective", model_meta_schema.ModelObjective.UNKNOWN.value)
-            ),
+            task=model_types.Task(model_dict.get("task", model_types.Task.UNKNOWN.value)),
             explain_algorithm=explanation_algorithm,
             function_properties=model_dict.get("function_properties", {}),
         )

snowflake/ml/model/_packager/model_meta/model_meta_schema.py

@@ -50,10 +50,6 @@ class LightGBMModelBlobOptions(BaseModelBlobOptions):
     lightgbm_estimator_type: Required[str]
 
 
-class LLMModelBlobOptions(BaseModelBlobOptions):
-    batch_size: Required[int]
-
-
 class MLFlowModelBlobOptions(BaseModelBlobOptions):
     artifact_path: Required[str]
 
@@ -65,7 +61,6 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
 ModelBlobOptions = Union[
     BaseModelBlobOptions,
     HuggingFacePipelineModelBlobOptions,
-    LLMModelBlobOptions,
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
 ]
@@ -96,18 +91,10 @@ class ModelMetadataDict(TypedDict):
     signatures: Required[Dict[str, Dict[str, Any]]]
     version: Required[str]
     min_snowpark_ml_version: Required[str]
-    model_objective: Required[str]
+    task: Required[str]
     explainability: NotRequired[Optional[ExplainabilityMetadataDict]]
     function_properties: NotRequired[Dict[str, Dict[str, Any]]]
 
 
-class ModelObjective(Enum):
-    UNKNOWN = "unknown"
-    BINARY_CLASSIFICATION = "binary_classification"
-    MULTI_CLASSIFICATION = "multi_classification"
-    REGRESSION = "regression"
-    RANKING = "ranking"
-
-
 class ModelExplainAlgorithm(Enum):
     SHAP = "shap"

snowflake/ml/model/_packager/model_packager.py

@@ -47,7 +47,9 @@ class ModelPackager:
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
         options: Optional[model_types.ModelSaveOption] = None,
+        task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:
+
         if (signatures is None) and (sample_input_data is None) and not model_handler.is_auto_signature_model(model):
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
@@ -56,17 +58,20 @@ class ModelPackager:
                 ),
             )
 
-        if (signatures is not None) and (sample_input_data is not None):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    "Signatures and sample_input_data both cannot be specified at the same time."
-                ),
-            )
-
         if not options:
             options = model_types.BaseModelSaveOption()
 
+        # here handling the case of enable_explainability is False/None
+        enable_explainability = options.get("enable_explainability", None)
+        if enable_explainability is False or enable_explainability is None:
+            if (signatures is not None) and (sample_input_data is not None):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INVALID_ARGUMENT,
+                    original_exception=ValueError(
+                        "Signatures and sample_input_data both cannot be specified at the same time."
+                    ),
+                )
+
         handler = model_handler.find_handler(model)
         if handler is None:
             raise snowml_exceptions.SnowflakeMLException(
@@ -84,6 +89,7 @@ class ModelPackager:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             python_version=python_version,
+            task=task,
             **options,
         ) as meta:
             model_blobs_path = os.path.join(self.local_dir_path, ModelPackager.MODEL_BLOBS_DIR)

snowflake/ml/model/_packager/model_runtime/model_runtime.py

@@ -36,6 +36,7 @@ class ModelRuntime:
         name: str,
         env: model_env.ModelEnv,
         imports: Optional[List[str]] = None,
+        is_warehouse: bool = False,
         is_gpu: bool = False,
         loading_from_file: bool = False,
     ) -> None:
@@ -60,6 +61,16 @@ class ModelRuntime:
             ],
         )
 
+        if not is_warehouse and self.embed_local_ml_library:
+            self.runtime_env.include_if_absent(
+                [
+                    model_env.ModelDependency(
+                        requirement="pyarrow",
+                        pip_name="pyarrow",
+                    )
+                ],
+            )
+
         if is_gpu:
             self.runtime_env.generate_env_for_cuda()
 

snowflake/ml/model/_signatures/pytorch_handler.py

@@ -30,7 +30,7 @@ class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Ten
 
     @staticmethod
     def count(data: Sequence["torch.Tensor"]) -> int:
-        return min(data_col.shape[0] for data_col in data)
+        return min(data_col.shape[0] for data_col in data)  # type: ignore[no-any-return]
 
     @staticmethod
     def truncate(data: Sequence["torch.Tensor"]) -> Sequence["torch.Tensor"]:

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -14,9 +14,10 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import type_hints as model_types
-from snowflake.ml.model._deploy_client.warehouse import infer_template
 from snowflake.ml.model._signatures import base_handler, core, pandas_handler
 
+_KEEP_ORDER_COL_NAME = "_ID"
+
 
 class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.DataFrame]):
     @staticmethod
@@ -109,7 +110,7 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         # Role will be no effect on the column index. That is to say, the feature name is the actual column name.
         if keep_order:
             df = df.reset_index(drop=True)
-            df[infer_template._KEEP_ORDER_COL_NAME] = df.index
+            df[_KEEP_ORDER_COL_NAME] = df.index
         sp_df = session.create_dataframe(df)
         column_names = []
         columns = []

snowflake/ml/model/_signatures/utils.py

@@ -110,6 +110,15 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ConversationalPipeline
     # Needs to convert to conversation object.
     if task == "conversational":
+        warnings.warn(
+            (
+                "Conversational pipeline is removed from transformers since 4.42.0. "
+                "Support will be removed from snowflake-ml-python soon."
+            ),
+            category=DeprecationWarning,
+            stacklevel=1,
+        )
+
         return core.ModelSignature(
             inputs=[
                 core.FeatureSpec(name="user_inputs", dtype=core.DataType.STRING, shape=(-1,)),

snowflake/ml/model/type_hints.py

@@ -1,22 +1,9 @@
 # mypy: disable-error-code="import"
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Sequence,
-    TypedDict,
-    TypeVar,
-    Union,
-)
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, Literal, Sequence, TypedDict, TypeVar, Union
 
 import numpy.typing as npt
-from typing_extensions import NotRequired, Required
-
-from snowflake.ml.model import deploy_platforms
-from snowflake.ml.model._signatures import core
+from typing_extensions import NotRequired
 
 if TYPE_CHECKING:
     import catboost
@@ -34,7 +21,6 @@ if TYPE_CHECKING:
 
     import snowflake.ml.model.custom_model
     import snowflake.ml.model.models.huggingface_pipeline
-    import snowflake.ml.model.models.llm
     import snowflake.snowpark
     from snowflake.ml.modeling.framework import base  # noqa: F401
 
@@ -90,7 +76,6 @@ SupportedNoSignatureRequirementsModelType = Union[
     "transformers.Pipeline",
     "sentence_transformers.SentenceTransformer",
     "snowflake.ml.model.models.huggingface_pipeline.HuggingFacePipelineModel",
-    "snowflake.ml.model.models.llm.LLM",
 ]
 
 SupportedModelType = Union[
@@ -133,86 +118,11 @@ SupportedModelHandlerType = Literal[
     "tensorflow",
     "torchscript",
     "xgboost",
-    "llm",
 ]
 
 _ModelType = TypeVar("_ModelType", bound=SupportedModelType)
 
 
-class DeployOptions(TypedDict):
-    """Common Options for deploying to Snowflake."""
-
-    ...
-
-
-class WarehouseDeployOptions(DeployOptions):
-    """Options for deploying to the Snowflake Warehouse.
-
-
-    permanent_udf_stage_location: A Snowflake stage option where the UDF should be persisted. If specified, the model
-        will be deployed as a permanent UDF, otherwise temporary.
-    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. It detects any
-        ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
-    replace_udf: Flag to indicate when deploying model as permanent UDF, whether overwriting existed UDF is allowed.
-        Default to False.
-    """
-
-    permanent_udf_stage_location: NotRequired[str]
-    relax_version: NotRequired[bool]
-    replace_udf: NotRequired[bool]
-
-
-class SnowparkContainerServiceDeployOptions(DeployOptions):
-    """Deployment options for deploying to SnowService.
-    When type hint is updated, please ensure the concrete class is updated accordingly at:
-    //snowflake/ml/model/_deploy_client/snowservice/_deploy_options
-
-    compute_pool[REQUIRED]: SnowService compute pool name. Please refer to official doc for how to create a
-        compute pool: https://docs.snowflake.com/LIMITEDACCESS/snowpark-containers/reference/compute-pool
-    image_repo: SnowService image repo path. e.g. "<image_registry>/<db>/<schema>/<repo>". Default to auto
-        inferred based on session information.
-    min_instances: Minimum number of service replicas. Default to 1.
-    max_instances: Maximum number of service replicas. Default to 1.
-    prebuilt_snowflake_image: When provided, the image-building step is skipped, and the pre-built image from
-        Snowflake is used as is. This option is for users who consistently use the same image for multiple use
-        cases, allowing faster deployment. The snowflake image used for deployment is logged to the console for
-        future use. Default to None.
-    num_gpus: Number of GPUs to be used for the service. Default to 0.
-    num_workers: Number of workers used for model inference. Please ensure that the number of workers is set lower than
-        the total available memory divided by the size of model to prevent memory-related issues. Default is number of
-        CPU cores * 2 + 1.
-    enable_remote_image_build: When set to True, will enable image build on a remote SnowService job. Default is True.
-    force_image_build: When set to True, an image rebuild will occur. The default is False, which means the system
-        will automatically check whether a previously built image can be reused
-    model_in_image: When set to True, image would container full model weights. The default if False, which
-                means image without model weights and we do stage mount to access weights.
-    debug_mode: When set to True, deployment artifacts will be persisted in a local temp directory.
-    enable_ingress: When set to True, will expose HTTP endpoint for access to the predict method of the created
-        service.
-    external_access_integrations: External Access Integrations name used to build image and deploy the model.
-        Please refer to the doc for how to create an External Access Integrations: https://docs.snowflake.com/
-        developer-guide/snowpark-container-services/additional-considerations-services-jobs
-        #configuring-network-capabilities .
-        To make sure your image could be built, access to the following endpoint must be allowed.
-        docker.com:80, docker.com:443, anaconda.com:80, anaconda.com:443, anaconda.org:80, anaconda.org:443,
-        pypi.org:80, pypi.org:443
-    """
-
-    compute_pool: str
-    image_repo: NotRequired[str]
-    min_instances: NotRequired[int]
-    max_instances: NotRequired[int]
-    prebuilt_snowflake_image: NotRequired[str]
-    num_gpus: NotRequired[int]
-    num_workers: NotRequired[int]
-    enable_remote_image_build: NotRequired[bool]
-    force_image_build: NotRequired[bool]
-    model_in_image: NotRequired[bool]
-    debug_mode: NotRequired[bool]
-    enable_ingress: NotRequired[bool]
-    external_access_integrations: List[str]
-
-
 class ModelMethodSaveOptions(TypedDict):
     case_sensitive: NotRequired[bool]
     max_batch_size: NotRequired[int]
@@ -223,16 +133,14 @@ class BaseModelSaveOption(TypedDict):
     """Options for saving the model.
 
     embed_local_ml_library: Embedding local SnowML into the code directory of the folder.
-    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. It detects any
-        ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
+    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable in Warehouse.
+        It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
     """
 
     embed_local_ml_library: NotRequired[bool]
     relax_version: NotRequired[bool]
-    _legacy_save: NotRequired[bool]
     function_type: NotRequired[Literal["FUNCTION", "TABLE_FUNCTION"]]
     method_options: NotRequired[Dict[str, ModelMethodSaveOptions]]
-    include_pip_dependencies: NotRequired[bool]
     enable_explainability: NotRequired[bool]
 
 
@@ -293,10 +201,6 @@ class SentenceTransformersSaveOptions(BaseModelSaveOption):
     cuda_version: NotRequired[str]
 
 
-class LLMSaveOptions(BaseModelSaveOption):
-    cuda_version: NotRequired[str]
-
-
 ModelSaveOption = Union[
     BaseModelSaveOption,
     CatBoostModelSaveOptions,
@@ -311,7 +215,6 @@ ModelSaveOption = Union[
     MLFlowSaveOptions,
     HuggingFaceSaveOptions,
     SentenceTransformersSaveOptions,
-    LLMSaveOptions,
 ]
 
 
@@ -369,10 +272,7 @@ class HuggingFaceLoadOptions(BaseModelLoadOption):
 
 class SentenceTransformersLoadOptions(BaseModelLoadOption):
     use_gpu: NotRequired[bool]
-
-
-class LLMLoadOptions(BaseModelLoadOption):
-    ...
+    device: NotRequired[str]
 
 
 ModelLoadOption = Union[
@@ -389,45 +289,12 @@ ModelLoadOption = Union[
     MLFlowLoadOptions,
     HuggingFaceLoadOptions,
     SentenceTransformersLoadOptions,
-    LLMLoadOptions,
 ]
 
 
-class SnowparkContainerServiceDeployDetails(TypedDict):
-    """
-    Attributes:
-        service_info: A snowpark row containing the result of "describe service"
-        service_function_sql: SQL for service function creation.
-    """
-
-    service_info: Optional[Dict[str, Any]]
-    service_function_sql: str
-
-
-class WarehouseDeployDetails(TypedDict):
-    ...
-
-
-DeployDetails = Union[
-    SnowparkContainerServiceDeployDetails,
-    WarehouseDeployDetails,
-]
-
-
-class Deployment(TypedDict):
-    """Deployment information.
-
-    Attributes:
-        name: Name of the deployment.
-        platform: Target platform to deploy the model.
-        target_method: Target method name.
-        signature: The signature of the model method.
-        options: Additional options when deploying the model.
-    """
-
-    name: Required[str]
-    platform: Required[deploy_platforms.TargetPlatform]
-    target_method: Required[str]
-    signature: core.ModelSignature
-    options: Required[DeployOptions]
-    details: NotRequired[DeployDetails]
+class Task(Enum):
+    UNKNOWN = "UNKNOWN"
+    TABULAR_BINARY_CLASSIFICATION = "TABULAR_BINARY_CLASSIFICATION"
+    TABULAR_MULTI_CLASSIFICATION = "TABULAR_MULTI_CLASSIFICATION"
+    TABULAR_REGRESSION = "TABULAR_REGRESSION"
+    TABULAR_RANKING = "TABULAR_RANKING"

snowflake/ml/modeling/_internal/constants.py

@@ -1 +1,2 @@
 IN_ML_RUNTIME_ENV_VAR = "IN_SPCS_ML_RUNTIME"
+USE_OPTIMIZED_DATA_INGESTOR = "USE_OPTIMIZED_DATA_INGESTOR"

snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py

@@ -166,10 +166,10 @@ class PandasTransformHandlers:
             SnowflakeMLException: The input column list does not have one of `X` and `X_test`.
         """
         assert hasattr(self.estimator, "score")  # make type checker happy
-        argspec = inspect.getfullargspec(self.estimator.score)
-        if "X" in argspec.args:
+        params = inspect.signature(self.estimator.score).parameters
+        if "X" in params:
             score_args = {"X": self.dataset[input_cols]}
-        elif "X_test" in argspec.args:
+        elif "X_test" in params:
             score_args = {"X_test": self.dataset[input_cols]}
         else:
             raise exceptions.SnowflakeMLException(
@@ -178,10 +178,10 @@ class PandasTransformHandlers:
             )
 
         if len(label_cols) > 0:
-            label_arg_name = "Y" if "Y" in argspec.args else "y"
+            label_arg_name = "Y" if "Y" in params else "y"
             score_args[label_arg_name] = self.dataset[label_cols].squeeze()
 
-        if sample_weight_col is not None and "sample_weight" in argspec.args:
+        if sample_weight_col is not None and "sample_weight" in params:
             score_args["sample_weight"] = self.dataset[sample_weight_col].squeeze()
 
         score = self.estimator.score(**score_args)

snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py

@@ -43,14 +43,14 @@ class PandasModelTrainer:
             Trained model
         """
         assert hasattr(self.estimator, "fit")  # Keep mypy happy
-        argspec = inspect.getfullargspec(self.estimator.fit)
+        params = inspect.signature(self.estimator.fit).parameters
         args = {"X": self.dataset[self.input_cols]}
 
         if self.label_cols:
-            label_arg_name = "Y" if "Y" in argspec.args else "y"
+            label_arg_name = "Y" if "Y" in params 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:
+        if self.sample_weight_col is not None and "sample_weight" in params:
             args["sample_weight"] = self.dataset[self.sample_weight_col].squeeze()
 
         return self.estimator.fit(**args)
@@ -59,6 +59,7 @@ class PandasModelTrainer:
         self,
         expected_output_cols_list: List[str],
         drop_input_cols: Optional[bool] = False,
+        example_output_pd_df: Optional[pd.DataFrame] = None,
     ) -> Tuple[pd.DataFrame, object]:
         """Trains the model using specified features and target columns from the dataset.
         This API is different from fit itself because it would also provide the predict
@@ -69,6 +70,8 @@ class PandasModelTrainer:
                 name as a list. Defaults to None.
             drop_input_cols (Optional[bool]): Boolean to determine whether to
                 drop the input columns from the output dataset.
+            example_output_pd_df (Optional[pd.DataFrame]): Example output dataframe
+                This is not used in PandasModelTrainer. It is used in SnowparkModelTrainer.
 
         Returns:
             Tuple[pd.DataFrame, object]: [predicted dataset, estimator]
@@ -108,13 +111,13 @@ class PandasModelTrainer:
         assert hasattr(self.estimator, "fit")  # make type checker happy
         assert hasattr(self.estimator, "fit_transform")  # make type checker happy
 
-        argspec = inspect.getfullargspec(self.estimator.fit)
+        params = inspect.signature(self.estimator.fit).parameters
         args = {"X": self.dataset[self.input_cols]}
         if self.label_cols:
-            label_arg_name = "Y" if "Y" in argspec.args else "y"
+            label_arg_name = "Y" if "Y" in params 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:
+        if self.sample_weight_col is not None and "sample_weight" in params:
             args["sample_weight"] = self.dataset[self.sample_weight_col].squeeze()
 
         inference_res = self.estimator.fit_transform(**args)

snowflake/ml/modeling/_internal/model_specifications.py

@@ -53,11 +53,13 @@ class SKLearnModelSpecifications(ModelSpecifications):
 
 class XGBoostModelSpecifications(ModelSpecifications):
     def __init__(self) -> None:
+        import sklearn
         import xgboost
 
         imports: List[str] = ["xgboost"]
         pkgDependencies: List[str] = [
             f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
             f"xgboost=={xgboost.__version__}",
             f"cloudpickle=={cp.__version__}",
         ]