snowflake-ml-python 1.9.0__py3-none-any.whl → 1.9.2__py3-none-any.whl

snowflake/ml/model/_client/sql/service.py

@@ -2,7 +2,7 @@ import dataclasses
 import enum
 import logging
 import textwrap
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import (
@@ -69,43 +69,6 @@ class ServiceSQLClient(_base._BaseSQLClient):
     CONTAINER_STATUS = "status"
     MESSAGE = "message"
 
-    def build_model_container(
-        self,
-        *,
-        database_name: Optional[sql_identifier.SqlIdentifier],
-        schema_name: Optional[sql_identifier.SqlIdentifier],
-        model_name: sql_identifier.SqlIdentifier,
-        version_name: sql_identifier.SqlIdentifier,
-        compute_pool_name: sql_identifier.SqlIdentifier,
-        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_name: sql_identifier.SqlIdentifier,
-        gpu: Optional[Union[str, int]],
-        force_rebuild: bool,
-        external_access_integration: sql_identifier.SqlIdentifier,
-        statement_params: Optional[dict[str, Any]] = None,
-    ) -> None:
-        actual_image_repo_database = image_repo_database_name or self._database_name
-        actual_image_repo_schema = image_repo_schema_name or self._schema_name
-        actual_model_database = database_name or self._database_name
-        actual_model_schema = schema_name or self._schema_name
-        fq_model_name = self.fully_qualified_object_name(actual_model_database, actual_model_schema, model_name)
-        fq_image_repo_name = identifier.get_schema_level_object_identifier(
-            actual_image_repo_database.identifier(),
-            actual_image_repo_schema.identifier(),
-            image_repo_name.identifier(),
-        )
-        is_gpu_str = "TRUE" if gpu else "FALSE"
-        force_rebuild_str = "TRUE" if force_rebuild else "FALSE"
-        query_result_checker.SqlResultValidator(
-            self._session,
-            (
-                f"CALL SYSTEM$BUILD_MODEL_CONTAINER('{fq_model_name}', '{version_name}', '{compute_pool_name}',"
-                f" '{fq_image_repo_name}', '{is_gpu_str}', '{force_rebuild_str}', '', '{external_access_integration}')"
-            ),
-            statement_params=statement_params,
-        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
-
     def deploy_model(
         self,
         *,

snowflake/ml/model/_model_composer/model_composer.py

@@ -3,7 +3,7 @@ import tempfile
 import uuid
 import warnings
 from types import ModuleType
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 from urllib import parse
 
 from absl import logging
@@ -21,6 +21,9 @@ from snowflake.ml.model._packager.model_meta import model_meta
 from snowflake.snowpark import Session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
+if TYPE_CHECKING:
+    from snowflake.ml.experiment._experiment_info import ExperimentInfo
+
 
 class ModelComposer:
     """Top-level class to construct contents in a MODEL object in SQL.
@@ -136,6 +139,7 @@ class ModelComposer:
         ext_modules: Optional[list[ModuleType]] = None,
         code_paths: Optional[list[str]] = None,
         task: model_types.Task = model_types.Task.UNKNOWN,
+        experiment_info: Optional["ExperimentInfo"] = None,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
         # set enable_explainability=False if the model is not runnable in WH or the target platforms include SPCS
@@ -230,6 +234,7 @@ class ModelComposer:
             options=options,
             user_files=user_files,
             data_sources=self._get_data_sources(model, sample_input_data),
+            experiment_info=experiment_info,
             target_platforms=target_platforms,
         )
 

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import pathlib
 import warnings
-from typing import Optional, cast
+from typing import TYPE_CHECKING, Optional, cast
 
 import yaml
 
@@ -23,6 +23,9 @@ from snowflake.ml.model._packager.model_meta import (
 )
 from snowflake.ml.model._packager.model_runtime import model_runtime
 
+if TYPE_CHECKING:
+    from snowflake.ml.experiment._experiment_info import ExperimentInfo
+
 logger = logging.getLogger(__name__)
 
 
@@ -49,6 +52,7 @@ class ModelManifest:
         user_files: Optional[dict[str, list[str]]] = None,
         options: Optional[type_hints.ModelSaveOption] = None,
         data_sources: Optional[list[data_source.DataSource]] = None,
+        experiment_info: Optional["ExperimentInfo"] = None,
         target_platforms: Optional[list[type_hints.TargetPlatform]] = None,
     ) -> None:
         if options is None:
@@ -183,7 +187,7 @@ class ModelManifest:
         if self.user_files:
             manifest_dict["user_files"] = [user_file.save(self.workspace_path) for user_file in self.user_files]
 
-        lineage_sources = self._extract_lineage_info(data_sources)
+        lineage_sources = self._extract_lineage_info(data_sources, experiment_info)
         if lineage_sources:
             manifest_dict["lineage_sources"] = lineage_sources
 
@@ -210,7 +214,9 @@ class ModelManifest:
         return res
 
     def _extract_lineage_info(
-        self, data_sources: Optional[list[data_source.DataSource]]
+        self,
+        data_sources: Optional[list[data_source.DataSource]],
+        experiment_info: Optional["ExperimentInfo"],
     ) -> list[model_manifest_schema.LineageSourceDict]:
         result = []
         if data_sources:
@@ -229,4 +235,12 @@ class ModelManifest:
                             type=model_manifest_schema.LineageSourceTypes.QUERY.value, entity=source.sql
                         )
                     )
+        if experiment_info:
+            result.append(
+                model_manifest_schema.LineageSourceDict(
+                    type=model_manifest_schema.LineageSourceTypes.EXPERIMENT.value,
+                    entity=experiment_info.fully_qualified_name,
+                    version=experiment_info.run_name,
+                )
+            )
         return result

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -83,6 +83,7 @@ class SnowparkMLDataDict(TypedDict):
 class LineageSourceTypes(enum.Enum):
     DATASET = "DATASET"
     QUERY = "QUERY"
+    EXPERIMENT = "EXPERIMENT"
 
 
 class LineageSourceDict(TypedDict):

snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 import warnings
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
@@ -23,9 +24,13 @@ from snowflake.ml.model._signatures import utils as model_signature_utils
 from snowflake.ml.model.models import huggingface_pipeline
 from snowflake.snowpark._internal import utils as snowpark_utils
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     import transformers
 
+DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"  # noqa: E501
+
 
 def get_requirements_from_task(task: str, spcs_only: bool = False) -> list[model_env.ModelDependency]:
     # Text
@@ -326,6 +331,23 @@ class HuggingFacePipelineHandler(
                 **device_config,
             )
 
+            # If the task is text-generation, and the tokenizer does not have a chat_template,
+            # set the default chat template.
+            if (
+                hasattr(m, "task")
+                and m.task == "text-generation"
+                and hasattr(m.tokenizer, "chat_template")
+                and not m.tokenizer.chat_template
+            ):
+                warnings.warn(
+                    "The tokenizer does not have default chat_template. "
+                    "Setting the chat_template to default ChatML template.",
+                    UserWarning,
+                    stacklevel=1,
+                )
+                logger.info(DEFAULT_CHAT_TEMPLATE)
+                m.tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE
+
             m.__dict__.update(pipeline_params)
 
         else:
@@ -481,8 +503,25 @@ class HuggingFacePipelineHandler(
 
         # To enable batch_size > 1 for LLM
         # Pipe might not have tokenizer, but should always have a model, and model should always have a config.
-        if getattr(pipe, "tokenizer", None) is not None and pipe.tokenizer.pad_token_id is None:
-            pipe.tokenizer.pad_token_id = pipe.model.config.eos_token_id
+        if (
+            getattr(pipe, "tokenizer", None) is not None
+            and pipe.tokenizer.pad_token_id is None
+            and hasattr(pipe.model.config, "eos_token_id")
+        ):
+            if isinstance(pipe.model.config.eos_token_id, int):
+                pipe.tokenizer.pad_token_id = pipe.model.config.eos_token_id
+            elif (
+                isinstance(pipe.model.config.eos_token_id, list)
+                and len(pipe.model.config.eos_token_id) > 0
+                and isinstance(pipe.model.config.eos_token_id[0], int)
+            ):
+                pipe.tokenizer.pad_token_id = pipe.model.config.eos_token_id[0]
+            else:
+                warnings.warn(
+                    f"Unexpected type of eos_token_id: {type(pipe.model.config.eos_token_id)}. "
+                    "Not setting pad_token_id to eos_token_id.",
+                    stacklevel=2,
+                )
 
         _HFPipelineModel = _create_custom_model(pipe, model_meta)
         hg_pipe_model = _HFPipelineModel(custom_model.ModelContext())

snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import warnings
 from typing import TYPE_CHECKING, Callable, Optional, Sequence, Union, cast, final
@@ -24,6 +25,8 @@ if TYPE_CHECKING:
     import sklearn.base
     import sklearn.pipeline
 
+logger = logging.getLogger(__name__)
+
 
 def _unpack_container_runtime_pipeline(model: "sklearn.pipeline.Pipeline") -> "sklearn.pipeline.Pipeline":
     new_steps = []
@@ -201,13 +204,13 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 explain_target_method = str(explain_target_method)  # mypy complains if we don't cast to str here
 
                 input_signature = handlers_utils.get_input_signature(model_meta, explain_target_method)
-                transformed_background_data = _apply_transforms_up_to_last_step(
-                    model=model,
-                    data=background_data,
-                    input_feature_names=[spec.name for spec in input_signature],
-                )
 
                 try:
+                    transformed_background_data = _apply_transforms_up_to_last_step(
+                        model=model,
+                        data=background_data,
+                        input_feature_names=[spec.name for spec in input_signature],
+                    )
                     model_meta = handlers_utils.add_inferred_explain_method_signature(
                         model_meta=model_meta,
                         explain_method="explain",
@@ -217,6 +220,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                         output_feature_names=transformed_background_data.columns,
                     )
                 except Exception:
+                    logger.debug("Explainability is disabled due to an exception.", exc_info=True)
                     if kwargs.get("enable_explainability", None):
                         # user explicitly enabled explainability, so we should raise the error
                         raise ValueError(

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

@@ -10,9 +10,10 @@ REQUIREMENTS = [
     "cryptography",
     "fsspec>=2024.6.1,<2026",
     "importlib_resources>=6.1.1, <7",
-    "numpy>=1.23,<2",
+    "numpy>=1.23,<3",
     "packaging>=20.9,<25",
     "pandas>=2.1.4,<3",
+    "platformdirs<5",
     "pyarrow",
     "pydantic>=2.8.2, <3",
     "pyjwt>=2.0.0, <3",
@@ -28,6 +29,7 @@ REQUIREMENTS = [
     "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0",
     "snowflake.core>=1.0.2,<2",
     "sqlparse>=0.4,<1",
+    "tqdm<5",
     "typing-extensions>=4.1.0,<5",
     "xgboost>=1.7.3,<3",
 ]

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

@@ -98,9 +98,9 @@ class ModelRuntime:
             dependencies=model_meta_schema.ModelRuntimeDependenciesDict(
                 conda=env_dict["conda"],
                 pip=env_dict["pip"],
-                artifact_repository_map=env_dict["artifact_repository_map"]
-                if env_dict.get("artifact_repository_map") is not None
-                else {},
+                artifact_repository_map=(
+                    env_dict["artifact_repository_map"] if env_dict.get("artifact_repository_map") is not None else {}
+                ),
             ),
             resource_constraint=env_dict["resource_constraint"],
         )

snowflake/ml/model/_signatures/pandas_handler.py

@@ -86,6 +86,9 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
                 df_col_data = utils.series_dropna(df_col_data)
                 df_col_dtype = df_col_data.dtype
 
+            if utils.check_if_series_is_empty(df_col_data):
+                continue
+
             if df_col_dtype == np.dtype("O"):
                 # Check if all objects have the same type
                 if not all(isinstance(data_row, type(df_col_data.iloc[0])) for data_row in df_col_data):

snowflake/ml/model/_signatures/utils.py

@@ -412,3 +412,7 @@ def infer_dict(name: str, data: dict[str, Any]) -> core.FeatureGroupSpec:
             specs.append(core.FeatureSpec(name=key, dtype=core.DataType.from_numpy_type(np.array(value).dtype)))
 
     return core.FeatureGroupSpec(name=name, specs=specs)
+
+
+def check_if_series_is_empty(series: Optional[pd.Series]) -> bool:
+    return series is None or series.empty

snowflake/ml/model/event_handler.py

@@ -0,0 +1,117 @@
+import os
+import sys
+from typing import Any, Optional
+
+
+class _TqdmStatusContext:
+    """A tqdm-based context manager for status updates."""
+
+    def __init__(self, label: str, tqdm_module: Any, total: Optional[int] = None) -> None:
+        self._label = label
+        self._tqdm = tqdm_module
+        self._total = total or 1
+
+    def __enter__(self) -> "_TqdmStatusContext":
+        self._progress_bar = self._tqdm.tqdm(desc=self._label, file=sys.stdout, total=self._total, leave=True)
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self._progress_bar.close()
+
+    def update(self, label: str, *, state: str = "running", expanded: bool = True) -> None:
+        """Update the status by updating the tqdm description."""
+        if state == "complete":
+            self._progress_bar.update(self._progress_bar.total - self._progress_bar.n)
+            self._progress_bar.set_description(label)
+        else:
+            self._progress_bar.set_description(f"{self._label}: {label}")
+
+    def increment(self, n: int = 1) -> None:
+        """Increment the progress bar."""
+        self._progress_bar.update(n)
+
+
+class _StreamlitStatusContext:
+    """A streamlit-based context manager for status updates with progress bar support."""
+
+    def __init__(self, label: str, streamlit_module: Any, total: Optional[int] = None) -> None:
+        self._label = label
+        self._streamlit = streamlit_module
+        self._total = total
+        self._current = 0
+        self._progress_bar = None
+
+    def __enter__(self) -> "_StreamlitStatusContext":
+        self._status_container = self._streamlit.status(self._label, state="running", expanded=True)
+        if self._total is not None:
+            with self._status_container:
+                self._progress_bar = self._streamlit.progress(0, text=f"0/{self._total}")
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self._status_container.update(state="complete")
+
+    def update(self, label: str, *, state: str = "running", expanded: bool = True) -> None:
+        """Update the status label."""
+        if state != "complete":
+            label = f"{self._label}: {label}"
+        self._status_container.update(label=label, state=state, expanded=expanded)
+        if self._progress_bar is not None:
+            self._progress_bar.progress(
+                self._current / self._total if self._total > 0 else 0,
+                text=f"{label} - {self._current}/{self._total}",
+            )
+
+    def increment(self, n: int = 1) -> None:
+        """Increment the progress."""
+        if self._total is not None:
+            self._current = min(self._current + n, self._total)
+            if self._progress_bar is not None:
+                progress_value = self._current / self._total if self._total > 0 else 0
+                self._progress_bar.progress(progress_value, text=f"{self._current}/{self._total}")
+
+
+class ModelEventHandler:
+    """Event handler for model operations with streamlit-aware status updates."""
+
+    def __init__(self) -> None:
+        self._streamlit = None
+
+        # Try streamlit first
+        try:
+            import streamlit as st
+
+            if st.runtime.exists():
+                USE_STREAMLIT_WIDGETS = os.getenv("USE_STREAMLIT_WIDGETS", "1") == "1"
+                if USE_STREAMLIT_WIDGETS:
+                    self._streamlit = st
+        except ImportError:
+            pass
+
+        import tqdm
+
+        self._tqdm = tqdm
+
+    def update(self, message: str) -> None:
+        """Write a message using streamlit if available, otherwise use tqdm."""
+        if self._streamlit is not None:
+            self._streamlit.write(message)
+        else:
+            self._tqdm.tqdm.write(message)
+
+    def status(self, label: str, *, state: str = "running", expanded: bool = True, total: Optional[int] = None) -> Any:
+        """Context manager that provides status updates with optional enhanced display capabilities.
+
+        Args:
+            label: The status label
+            state: The initial state ("running", "complete", "error")
+            expanded: Whether to show expanded view (streamlit only)
+            total: Total number of steps for progress tracking (optional)
+
+        Returns:
+            Status context (Streamlit or Tqdm)
+        """
+        if self._streamlit is not None:
+            return _StreamlitStatusContext(label, self._streamlit, total)
+        else:
+            return _TqdmStatusContext(label, self._tqdm, total)

snowflake/ml/model/model_signature.py

@@ -16,7 +16,7 @@ from snowflake.ml._internal.exceptions import (
     exceptions as snowml_exceptions,
 )
 from snowflake.ml._internal.utils import formatting, identifier, sql_identifier
-from snowflake.ml.model import type_hints as model_types
+from snowflake.ml.model import type_hints
 from snowflake.ml.model._signatures import (
     base_handler,
     builtins_handler,
@@ -55,9 +55,9 @@ _MODEL_TELEMETRY_SUBPROJECT = "ModelSignature"
 
 
 def _truncate_data(
-    data: model_types.SupportedDataType,
+    data: type_hints.SupportedDataType,
     length: Optional[int] = 100,
-) -> model_types.SupportedDataType:
+) -> type_hints.SupportedDataType:
     for handler in _ALL_DATA_HANDLERS:
         if handler.can_handle(data):
             # If length is None, return the original data
@@ -89,7 +89,7 @@ def _truncate_data(
 
 
 def _infer_signature(
-    data: model_types.SupportedLocalDataType, role: Literal["input", "output"], use_snowflake_identifiers: bool = False
+    data: type_hints.SupportedLocalDataType, role: Literal["input", "output"], use_snowflake_identifiers: bool = False
 ) -> Sequence[core.BaseFeatureSpec]:
     """Infer the inputs/outputs signature given a data that could be dataframe, numpy array or list.
         Dispatching is used to separate logic for different types.
@@ -142,7 +142,7 @@ def _rename_signature_with_snowflake_identifiers(
 
 
 def _validate_array_or_series_type(
-    arr: Union[model_types._SupportedNumpyArray, pd.Series], feature_type: core.DataType, strict: bool = False
+    arr: Union[type_hints._SupportedNumpyArray, pd.Series], feature_type: core.DataType, strict: bool = False
 ) -> bool:
     original_dtype = arr.dtype
     dtype = arr.dtype
@@ -272,6 +272,8 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
                     ),
                 )
         else:
+            if utils.check_if_series_is_empty(data_col):
+                continue
             if isinstance(data_col.iloc[0], list):
                 if not ft_shape:
                     raise snowml_exceptions.SnowflakeMLException(
@@ -649,7 +651,7 @@ def _validate_snowpark_type_feature(
 
 
 def _convert_local_data_to_df(
-    data: model_types.SupportedLocalDataType, ensure_serializable: bool = False
+    data: type_hints.SupportedLocalDataType, ensure_serializable: bool = False
 ) -> pd.DataFrame:
     """Convert local data to pandas DataFrame or Snowpark DataFrame
 
@@ -679,7 +681,7 @@ def _convert_local_data_to_df(
 
 
 def _convert_and_validate_local_data(
-    data: model_types.SupportedLocalDataType, features: Sequence[core.BaseFeatureSpec], strict: bool = False
+    data: type_hints.SupportedLocalDataType, features: Sequence[core.BaseFeatureSpec], strict: bool = False
 ) -> pd.DataFrame:
     """Validate the data with features in model signature and convert to DataFrame
 
@@ -703,8 +705,8 @@ def _convert_and_validate_local_data(
     subproject=_MODEL_TELEMETRY_SUBPROJECT,
 )
 def infer_signature(
-    input_data: model_types.SupportedLocalDataType,
-    output_data: model_types.SupportedLocalDataType,
+    input_data: type_hints.SupportedLocalDataType,
+    output_data: type_hints.SupportedLocalDataType,
     input_feature_names: Optional[list[str]] = None,
     output_feature_names: Optional[list[str]] = None,
     input_data_limit: Optional[int] = 100,