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

snowflake/ml/model/models/huggingface_pipeline.py

@@ -1,8 +1,22 @@
+import logging
 import warnings
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from packaging import version
 
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.human_readable_id import hrid_generator
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._client.ops import service_ops
+from snowflake.snowpark import async_job, session
+
+logger = logging.getLogger(__name__)
+
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
 
 class HuggingFacePipelineModel:
     def __init__(
@@ -214,4 +228,159 @@ class HuggingFacePipelineModel:
         self.token = token
         self.trust_remote_code = trust_remote_code
         self.model_kwargs = model_kwargs
+        self.tokenizer = tokenizer
         self.__dict__.update(kwargs)
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=[
+            "service_name",
+            "image_build_compute_pool",
+            "service_compute_pool",
+            "image_repo",
+            "gpu_requests",
+            "num_workers",
+            "max_batch_rows",
+        ],
+    )
+    @snowpark._internal.utils.private_preview(version="1.9.1")
+    def create_service(
+        self,
+        *,
+        session: session.Session,
+        # registry.log_model parameters
+        model_name: str,
+        version_name: Optional[str] = None,
+        pip_requirements: Optional[list[str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        comment: Optional[str] = None,
+        # model_version_impl.create_service parameters
+        service_name: str,
+        service_compute_pool: str,
+        image_repo: str,
+        image_build_compute_pool: Optional[str] = None,
+        ingress_enabled: bool = False,
+        max_instances: int = 1,
+        cpu_requests: Optional[str] = None,
+        memory_requests: Optional[str] = None,
+        gpu_requests: Optional[Union[str, int]] = None,
+        num_workers: Optional[int] = None,
+        max_batch_rows: Optional[int] = None,
+        force_rebuild: bool = False,
+        build_external_access_integrations: Optional[list[str]] = None,
+        block: bool = True,
+    ) -> Union[str, async_job.AsyncJob]:
+        """Logs a Hugging Face model and creates a service in Snowflake.
+
+        Args:
+            session: The Snowflake session object.
+            model_name: The name of the model in Snowflake.
+            version_name: The version name of the model. Defaults to None.
+            pip_requirements: Pip requirements for the model. Defaults to None.
+            conda_dependencies: Conda dependencies for the model. Defaults to None.
+            comment: Comment for the model. Defaults to None.
+            service_name: The name of the service to create.
+            service_compute_pool: The compute pool for the service.
+            image_repo: The name of the image repository.
+            image_build_compute_pool: The name of the compute pool used to build the model inference image. It uses
+            the service compute pool if None.
+            ingress_enabled: Whether ingress is enabled. Defaults to False.
+            max_instances: Maximum number of instances. Defaults to 1.
+            cpu_requests: CPU requests configuration. Defaults to None.
+            memory_requests: Memory requests configuration. Defaults to None.
+            gpu_requests: GPU requests configuration. Defaults to None.
+            num_workers: Number of workers. Defaults to None.
+            max_batch_rows: Maximum batch rows. Defaults to None.
+            force_rebuild: Whether to force rebuild the image. Defaults to False.
+            build_external_access_integrations: External access integrations for building the image. Defaults to None.
+            block: Whether to block the operation. Defaults to True.
+
+        Raises:
+            ValueError: if database and schema name is not provided and session doesn't have a
+            database and schema name.
+
+        Returns:
+            The service ID or an async job object.
+
+        .. # noqa: DAR003
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+
+        database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
+        session_database_name = session.get_current_database()
+        session_schema_name = session.get_current_schema()
+        if database_name_id is None:
+            if session_database_name is None:
+                raise ValueError("Either database needs to be provided or needs to be available in session.")
+            database_name_id = sql_identifier.SqlIdentifier(session_database_name)
+        if schema_name_id is None:
+            if session_schema_name is None:
+                raise ValueError("Either schema needs to be provided or needs to be available in session.")
+            schema_name_id = sql_identifier.SqlIdentifier(session_schema_name)
+
+        if version_name is None:
+            name_generator = hrid_generator.HRID16()
+            version_name = name_generator.generate()[1]
+
+        service_db_id, service_schema_id, service_id = sql_identifier.parse_fully_qualified_name(service_name)
+        image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
+
+        service_operator = service_ops.ServiceOperator(
+            session=session,
+            database_name=database_name_id,
+            schema_name=schema_name_id,
+        )
+        logger.info(f"A service job is going to register the hf model as: {model_name}.{version_name}")
+
+        return service_operator.create_service(
+            database_name=database_name_id,
+            schema_name=schema_name_id,
+            model_name=model_name_id,
+            version_name=sql_identifier.SqlIdentifier(version_name),
+            service_database_name=service_db_id,
+            service_schema_name=service_schema_id,
+            service_name=service_id,
+            image_build_compute_pool_name=(
+                sql_identifier.SqlIdentifier(image_build_compute_pool)
+                if image_build_compute_pool
+                else sql_identifier.SqlIdentifier(service_compute_pool)
+            ),
+            service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
+            image_repo_database_name=image_repo_db_id,
+            image_repo_schema_name=image_repo_schema_id,
+            image_repo_name=image_repo_id,
+            ingress_enabled=ingress_enabled,
+            max_instances=max_instances,
+            cpu_requests=cpu_requests,
+            memory_requests=memory_requests,
+            gpu_requests=gpu_requests,
+            num_workers=num_workers,
+            max_batch_rows=max_batch_rows,
+            force_rebuild=force_rebuild,
+            build_external_access_integrations=(
+                None
+                if build_external_access_integrations is None
+                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
+            ),
+            block=block,
+            statement_params=statement_params,
+            # hf model
+            hf_model_args=service_ops.HFModelArgs(
+                hf_model_name=self.model,
+                hf_task=self.task,
+                hf_tokenizer=self.tokenizer,
+                hf_revision=self.revision,
+                hf_token=self.token,
+                hf_trust_remote_code=bool(self.trust_remote_code),
+                hf_model_kwargs=self.model_kwargs,
+                pip_requirements=pip_requirements,
+                conda_dependencies=conda_dependencies,
+                comment=comment,
+                # TODO: remove warehouse in the next release
+                warehouse=session.get_current_warehouse(),
+            ),
+        )

snowflake/ml/modeling/framework/base.py

@@ -698,7 +698,7 @@ class BaseTransformer(BaseEstimator):
         self,
         attribute: Optional[Mapping[str, Union[int, float, str, Iterable[Union[int, float, str]]]]],
         dtype: Optional[type] = None,
-    ) -> Optional[npt.NDArray[Union[np.int_, np.float_, np.str_]]]:
+    ) -> Optional[npt.NDArray[Union[np.int_, np.float64, np.str_]]]:
         """
         Convert the attribute from dict to ndarray based on the order of `self.input_cols`.
 

snowflake/ml/modeling/metrics/classification.py

@@ -96,7 +96,7 @@ def confusion_matrix(
     labels: Optional[npt.ArrayLike] = None,
     sample_weight_col_name: Optional[str] = None,
     normalize: Optional[str] = None,
-) -> Union[npt.NDArray[np.int_], npt.NDArray[np.float_]]:
+) -> Union[npt.NDArray[np.int_], npt.NDArray[np.float64]]:
     """
     Compute confusion matrix to evaluate the accuracy of a classification.
 
@@ -320,7 +320,7 @@ def f1_score(
     average: Optional[str] = "binary",
     sample_weight_col_name: Optional[str] = None,
     zero_division: Union[str, int] = "warn",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Compute the F1 score, also known as balanced F-score or F-measure.
 
@@ -414,7 +414,7 @@ def fbeta_score(
     average: Optional[str] = "binary",
     sample_weight_col_name: Optional[str] = None,
     zero_division: Union[str, int] = "warn",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Compute the F-beta score.
 
@@ -696,7 +696,7 @@ def precision_recall_fscore_support(
     zero_division: Union[str, int] = "warn",
 ) -> Union[
     tuple[float, float, float, None],
-    tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]],
+    tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ]:
     """
     Compute precision, recall, F-measure and support for each class.
@@ -855,7 +855,7 @@ def precision_recall_fscore_support(
 
         res: Union[
             tuple[float, float, float, None],
-            tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]],
+            tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]],
         ] = result_object[:4]
         warning = result_object[-1]
         if warning:
@@ -1050,7 +1050,7 @@ def _register_multilabel_confusion_matrix_computer(
 
         def end_partition(
             self,
-        ) -> Iterable[tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]]]:
+        ) -> Iterable[tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]]:
             MCM = metrics.multilabel_confusion_matrix(
                 self._y_true,
                 self._y_pred,
@@ -1098,7 +1098,7 @@ def _binary_precision_score(
     pos_label: Union[str, int] = 1,
     sample_weight_col_name: Optional[str] = None,
     zero_division: Union[str, int] = "warn",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
 
     statement_params = telemetry.get_statement_params(_PROJECT, _SUBPROJECT)
 
@@ -1173,7 +1173,7 @@ def precision_score(
     average: Optional[str] = "binary",
     sample_weight_col_name: Optional[str] = None,
     zero_division: Union[str, int] = "warn",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Compute the precision.
 
@@ -1271,7 +1271,7 @@ def recall_score(
     average: Optional[str] = "binary",
     sample_weight_col_name: Optional[str] = None,
     zero_division: Union[str, int] = "warn",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Compute the recall.
 
@@ -1406,14 +1406,14 @@ def _check_binary_labels(
 
 
 def _prf_divide(
-    numerator: npt.NDArray[np.float_],
-    denominator: npt.NDArray[np.float_],
+    numerator: npt.NDArray[np.float64],
+    denominator: npt.NDArray[np.float64],
     metric: str,
     modifier: str,
     average: Optional[str] = None,
     warn_for: Union[tuple[str, ...], set[str]] = ("precision", "recall", "f-score"),
     zero_division: Union[str, int] = "warn",
-) -> npt.NDArray[np.float_]:
+) -> npt.NDArray[np.float64]:
     """Performs division and handles divide-by-zero.
 
     On zero-division, sets the corresponding result elements equal to
@@ -1436,7 +1436,7 @@ def _prf_divide(
             "warn", this acts as 0, but warnings are also raised.
 
     Returns:
-        npt.NDArray[np.float_]: Result of the division, an array of floats.
+        npt.NDArray[np.float64]: Result of the division, an array of floats.
     """
     mask = denominator == 0.0
     denominator = denominator.copy()
@@ -1522,7 +1522,7 @@ def _check_zero_division(zero_division: Union[int, float, str]) -> float:
         return np.nan
 
 
-def _nanaverage(a: npt.NDArray[np.float_], weights: Optional[npt.ArrayLike] = None) -> Any:
+def _nanaverage(a: npt.NDArray[np.float64], weights: Optional[npt.ArrayLike] = None) -> Any:
     """Compute the weighted average, ignoring NaNs.
 
     Args:

snowflake/ml/modeling/metrics/correlation.py

@@ -26,7 +26,7 @@ def correlation(*, df: snowpark.DataFrame, columns: Optional[Collection[str]] =
     The below steps explain how correlation matrix is computed in a distributed way:
     Let n = # of rows in the dataframe; sqrt_n = sqrt(n); X, Y are 2 columns in the dataframe
     Correlation(X, Y) = numerator/denominator where
-    numerator = dot(X/sqrt_n, Y/sqrt_n) - sum(X/n)*sum(X/n)
+    numerator = dot(X/sqrt_n, Y/sqrt_n) - sum(X/n)*sum(Y/n)
     denominator = std_dev(X)*std_dev(Y)
     std_dev(X) = sqrt(dot(X/sqrt_n, X/sqrt_n) - sum(X/n)*sum(X/n))
 
@@ -74,27 +74,38 @@ def correlation(*, df: snowpark.DataFrame, columns: Optional[Collection[str]] =
     # Pushing this to a udtf requires creating a temp udtf which takes about 20 secs, so it doesn't make sense
     # to have this in a udtf.
     n_cols = len(columns)
-    sum_arr = np.zeros(n_cols)
-    squared_sum_arr = np.zeros(n_cols)
+    column_means = np.zeros(n_cols)
+    mean_of_squares = np.zeros(n_cols)
     dot_prod = np.zeros((n_cols, n_cols))
     # Get sum, dot_prod and squared sum array from the results.
     for i in range(len(results)):
         x = results[i]
         if x[1] == "sum_by_count":
-            sum_arr = cloudpickle.loads(x[0])
+            column_means = cloudpickle.loads(x[0])
         else:
             row = int(x[1].strip("row_"))
             dot_prod[row, :] = cloudpickle.loads(x[0])
-            squared_sum_arr[row] = dot_prod[row, row]
+            mean_of_squares[row] = dot_prod[row, row]
 
     # sum(X/n)*sum(Y/n) is computed for all combinations of X,Y (columns in the dataframe)
-    exey_arr = np.einsum("t,m->tm", sum_arr, sum_arr, optimize="optimal")
+    exey_arr = np.einsum("t,m->tm", column_means, column_means, optimize="optimal")
     numerator_matrix = dot_prod - exey_arr
 
     # standard deviation for all columns in the dataframe
-    stddev_arr = np.sqrt(squared_sum_arr - np.einsum("i, i -> i", sum_arr, sum_arr, optimize="optimal"))
+    variance_arr = mean_of_squares - np.einsum("i, i -> i", column_means, column_means, optimize="optimal")
+    # ensure non-negative values from potential precision issues where variance might be slightly negative
+    variance_arr = np.maximum(variance_arr, 0)
+    stddev_arr = np.sqrt(variance_arr)
     # std_dev(X)*std_dev(Y) is computed for all combinations of X,Y (columns in the dataframe)
     denominator_matrix = np.einsum("t,m->tm", stddev_arr, stddev_arr, optimize="optimal")
-    corr_res = numerator_matrix / denominator_matrix
+
+    # Use np.divide to handle NaN cases
+    corr_res = np.divide(
+        numerator_matrix,
+        denominator_matrix,
+        out=np.full_like(numerator_matrix, np.nan),
+        where=(denominator_matrix != 0),
+    )
+
     correlation_matrix = pd.DataFrame(corr_res, columns=columns, index=columns)
     return correlation_matrix

snowflake/ml/modeling/metrics/ranking.py

@@ -26,7 +26,7 @@ def precision_recall_curve(
     probas_pred_col_name: str,
     pos_label: Optional[Union[str, int]] = None,
     sample_weight_col_name: Optional[str] = None,
-) -> tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]]:
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
     """
     Compute precision-recall pairs for different probability thresholds.
 
@@ -125,7 +125,7 @@ def precision_recall_curve(
 
     kwargs = telemetry.get_sproc_statement_params_kwargs(precision_recall_curve_anon_sproc, statement_params)
     result_object = result.deserialize(session, precision_recall_curve_anon_sproc(session, **kwargs))
-    res: tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]] = result_object
+    res: tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]] = result_object
     return res
 
 
@@ -140,7 +140,7 @@ def roc_auc_score(
     max_fpr: Optional[float] = None,
     multi_class: str = "raise",
     labels: Optional[npt.ArrayLike] = None,
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Compute Area Under the Receiver Operating Characteristic Curve (ROC AUC)
     from prediction scores.
@@ -276,7 +276,7 @@ def roc_auc_score(
 
     kwargs = telemetry.get_sproc_statement_params_kwargs(roc_auc_score_anon_sproc, statement_params)
     result_object = result.deserialize(session, roc_auc_score_anon_sproc(session, **kwargs))
-    auc: Union[float, npt.NDArray[np.float_]] = result_object
+    auc: Union[float, npt.NDArray[np.float64]] = result_object
     return auc
 
 
@@ -289,7 +289,7 @@ def roc_curve(
     pos_label: Optional[Union[str, int]] = None,
     sample_weight_col_name: Optional[str] = None,
     drop_intermediate: bool = True,
-) -> tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]]:
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
     """
     Compute Receiver operating characteristic (ROC).
 
@@ -380,6 +380,6 @@ def roc_curve(
     kwargs = telemetry.get_sproc_statement_params_kwargs(roc_curve_anon_sproc, statement_params)
     result_object = result.deserialize(session, roc_curve_anon_sproc(session, **kwargs))
 
-    res: tuple[npt.NDArray[np.float_], npt.NDArray[np.float_], npt.NDArray[np.float_]] = result_object
+    res: tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]] = result_object
 
     return res

snowflake/ml/modeling/metrics/regression.py

@@ -29,7 +29,7 @@ def d2_absolute_error_score(
     y_pred_col_names: Union[str, list[str]],
     sample_weight_col_name: Optional[str] = None,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     :math:`D^2` regression score function, \
     fraction of absolute error explained.
@@ -111,7 +111,7 @@ def d2_absolute_error_score(
 
     kwargs = telemetry.get_sproc_statement_params_kwargs(d2_absolute_error_score_anon_sproc, statement_params)
     result_object = result.deserialize(session, d2_absolute_error_score_anon_sproc(session, **kwargs))
-    score: Union[float, npt.NDArray[np.float_]] = result_object
+    score: Union[float, npt.NDArray[np.float64]] = result_object
     return score
 
 
@@ -124,7 +124,7 @@ def d2_pinball_score(
     sample_weight_col_name: Optional[str] = None,
     alpha: float = 0.5,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     :math:`D^2` regression score function, fraction of pinball loss explained.
 
@@ -211,7 +211,7 @@ def d2_pinball_score(
     kwargs = telemetry.get_sproc_statement_params_kwargs(d2_pinball_score_anon_sproc, statement_params)
     result_object = result.deserialize(session, d2_pinball_score_anon_sproc(session, **kwargs))
 
-    score: Union[float, npt.NDArray[np.float_]] = result_object
+    score: Union[float, npt.NDArray[np.float64]] = result_object
     return score
 
 
@@ -224,7 +224,7 @@ def explained_variance_score(
     sample_weight_col_name: Optional[str] = None,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
     force_finite: bool = True,
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Explained variance regression score function.
 
@@ -326,7 +326,7 @@ def explained_variance_score(
 
     kwargs = telemetry.get_sproc_statement_params_kwargs(explained_variance_score_anon_sproc, statement_params)
     result_object = result.deserialize(session, explained_variance_score_anon_sproc(session, **kwargs))
-    score: Union[float, npt.NDArray[np.float_]] = result_object
+    score: Union[float, npt.NDArray[np.float64]] = result_object
     return score
 
 
@@ -338,7 +338,7 @@ def mean_absolute_error(
     y_pred_col_names: Union[str, list[str]],
     sample_weight_col_name: Optional[str] = None,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Mean absolute error regression loss.
 
@@ -411,7 +411,7 @@ def mean_absolute_percentage_error(
     y_pred_col_names: Union[str, list[str]],
     sample_weight_col_name: Optional[str] = None,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Mean absolute percentage error (MAPE) regression loss.
 
@@ -495,7 +495,7 @@ def mean_squared_error(
     sample_weight_col_name: Optional[str] = None,
     multioutput: Union[str, npt.ArrayLike] = "uniform_average",
     squared: bool = True,
-) -> Union[float, npt.NDArray[np.float_]]:
+) -> Union[float, npt.NDArray[np.float64]]:
     """
     Mean squared error regression loss.
 

snowflake/ml/monitoring/explain_visualize.py

@@ -264,6 +264,7 @@ def plot_force(
 def plot_influence_sensitivity(
     shap_values: type_hints.SupportedDataType,
     feature_values: type_hints.SupportedDataType,
+    infer_is_categorical: bool = True,
     figsize: tuple[float, float] = DEFAULT_FIGSIZE,
 ) -> Any:
     """
@@ -274,6 +275,8 @@ def plot_influence_sensitivity(
     Args:
         shap_values: pandas Series or 2D array containing the SHAP values for a specific feature
         feature_values: pandas Series or 2D array containing the feature values for the same feature
+        infer_is_categorical: If True, the function will infer if the feature is categorical
+            based on the number of unique values.
         figsize: tuple of (width, height) for the plot
 
     Returns:
@@ -294,7 +297,7 @@ def plot_influence_sensitivity(
     elif feature_values_df.shape[0] != shap_values_df.shape[0]:
         raise ValueError("Feature values and SHAP values must have the same number of rows.")
 
-    scatter = _create_scatter_plot(feature_values, shap_values, figsize)
+    scatter = _create_scatter_plot(feature_values, shap_values, infer_is_categorical, figsize)
     return st.altair_chart(scatter) if use_streamlit else scatter
 
 
@@ -322,11 +325,13 @@ def _prepare_feature_values_for_streamlit(
     return feature_values, shap_values, st
 
 
-def _create_scatter_plot(feature_values: pd.Series, shap_values: pd.Series, figsize: tuple[float, float]) -> alt.Chart:
+def _create_scatter_plot(
+    feature_values: pd.Series, shap_values: pd.Series, infer_is_categorical: bool, figsize: tuple[float, float]
+) -> alt.Chart:
     unique_vals = np.sort(np.unique(feature_values.values))
     max_points_per_unique_value = float(np.max(np.bincount(np.searchsorted(unique_vals, feature_values.values))))
     points_per_value = len(feature_values.values) / len(unique_vals)
-    is_categorical = float(max(max_points_per_unique_value, points_per_value)) > 10
+    is_categorical = float(max(max_points_per_unique_value, points_per_value)) > 10 if infer_is_categorical else False
 
     kwargs = (
         {
@@ -403,9 +408,11 @@ def plot_violin(
         .transform_density(density="shap_value", groupby=["feature_name"], as_=["shap_value", "density"])
         .mark_area(orient="vertical")
         .encode(
-            y=alt.Y("density:Q", title=None).stack("center").impute(None).axis(labels=False, grid=False, ticks=True),
+            y=alt.Y("density:Q", title=None).stack("center").impute(None).axis(labels=False, grid=False, ticks=False),
             x=alt.X("shap_value:Q", title="SHAP Value"),
-            row=alt.Row("feature_name:N", sort=column_sort_order).spacing(0),
+            row=alt.Row(
+                "feature_name:N", sort=column_sort_order, header=alt.Header(labelAngle=0, labelAlign="left")
+            ).spacing(0),
             color=alt.Color("feature_name:N", legend=None),
             tooltip=["feature_name", "shap_value"],
         )

snowflake/ml/registry/_manager/model_manager.py

@@ -1,5 +1,5 @@
 from types import ModuleType
-from typing import Any, Optional, Protocol, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import pandas as pd
 from absl.logging import logging
@@ -17,15 +17,10 @@ from snowflake.ml.model._packager.model_meta import model_meta
 from snowflake.snowpark import exceptions as snowpark_exceptions, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
-logger = logging.getLogger(__name__)
-
+if TYPE_CHECKING:
+    from snowflake.ml.experiment._experiment_info import ExperimentInfo
 
-class EventHandler(Protocol):
-    """Protocol defining the interface for event handlers used during model operations."""
-
-    def update(self, message: str) -> None:
-        """Update with a progress message."""
-        ...
+logger = logging.getLogger(__name__)
 
 
 class ModelManager:
@@ -66,9 +61,10 @@ class ModelManager:
         code_paths: Optional[list[str]] = None,
         ext_modules: Optional[list[ModuleType]] = None,
         task: type_hints.Task = task.Task.UNKNOWN,
+        experiment_info: Optional["ExperimentInfo"] = None,
         options: Optional[type_hints.ModelSaveOption] = None,
         statement_params: Optional[dict[str, Any]] = None,
-        event_handler: EventHandler,
+        progress_status: Optional[Any] = None,
     ) -> model_version_impl.ModelVersion:
 
         database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
@@ -150,9 +146,10 @@ class ModelManager:
             code_paths=code_paths,
             ext_modules=ext_modules,
             task=task,
+            experiment_info=experiment_info,
             options=options,
             statement_params=statement_params,
-            event_handler=event_handler,
+            progress_status=progress_status,
         )
 
     def _log_model(
@@ -175,9 +172,10 @@ class ModelManager:
         code_paths: Optional[list[str]] = None,
         ext_modules: Optional[list[ModuleType]] = None,
         task: type_hints.Task = task.Task.UNKNOWN,
+        experiment_info: Optional["ExperimentInfo"] = None,
         options: Optional[type_hints.ModelSaveOption] = None,
         statement_params: Optional[dict[str, Any]] = None,
-        event_handler: EventHandler,
+        progress_status: Optional[Any] = None,
     ) -> model_version_impl.ModelVersion:
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
         version_name_id = sql_identifier.SqlIdentifier(version_name)
@@ -265,7 +263,9 @@ class ModelManager:
                 )
 
         logger.info("Start packaging and uploading your model. It might take some time based on the size of the model.")
-        event_handler.update("📦 Packaging model...")
+        if progress_status:
+            progress_status.update("packaging model...")
+            progress_status.increment()
 
         # Extract save_location from options if present
         save_location = None
@@ -279,6 +279,11 @@ class ModelManager:
             statement_params=statement_params,
             save_location=save_location,
         )
+
+        if progress_status:
+            progress_status.update("creating model manifest...")
+            progress_status.increment()
+
         model_metadata: model_meta.ModelMetadata = mc.save(
             name=model_name_id.resolved(),
             model=model,
@@ -295,7 +300,12 @@ class ModelManager:
             ext_modules=ext_modules,
             options=options,
             task=task,
+            experiment_info=experiment_info,
         )
+
+        if progress_status:
+            progress_status.update("uploading model files...")
+            progress_status.increment()
         statement_params = telemetry.add_statement_params_custom_tags(
             statement_params, model_metadata.telemetry_metadata()
         )
@@ -304,7 +314,9 @@ class ModelManager:
         )
 
         logger.info("Start creating MODEL object for you in the Snowflake.")
-        event_handler.update("🏗️ Creating model object in Snowflake...")
+        if progress_status:
+            progress_status.update("creating model object in Snowflake...")
+            progress_status.increment()
 
         self._model_ops.create_from_stage(
             composed_model=mc,
@@ -331,6 +343,10 @@ class ModelManager:
             version_name=version_name_id,
         )
 
+        if progress_status:
+            progress_status.update("setting model metadata...")
+            progress_status.increment()
+
         if comment:
             mv.comment = comment
 
@@ -344,7 +360,8 @@ class ModelManager:
                 statement_params=statement_params,
             )
 
-        event_handler.update("✅ Model logged successfully!")
+        if progress_status:
+            progress_status.update("model logged successfully!")
 
         return mv