snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -280,7 +280,7 @@ def _get_or_create_image_repo(session: Session, *, service_func_name: str, image
         conn = session._conn._conn
         # We try to use the same db and schema as the service function locates, as we could retrieve those information
         # if that is a fully qualified one. If not we use the current session one.
-        (_db, _schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)
+        (_db, _schema, _) = identifier.parse_schema_level_object_identifier(service_func_name)
         db = _db if _db is not None else conn._database
         schema = _schema if _schema is not None else conn._schema
         assert isinstance(db, str) and isinstance(schema, str)
@@ -343,7 +343,7 @@ class SnowServiceDeployment:
         self.model_zip_stage_path = model_zip_stage_path
         self.options = options
         self.target_method = target_method
-        (db, schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)
+        (db, schema, _) = identifier.parse_schema_level_object_identifier(service_func_name)
 
         self._service_name = identifier.get_schema_level_object_identifier(db, schema, f"service_{model_id}")
         self._job_name = identifier.get_schema_level_object_identifier(db, schema, f"build_{model_id}")
@@ -503,7 +503,7 @@ class SnowServiceDeployment:
                 norm_stage_path = posixpath.normpath(identifier.remove_prefix(self.model_zip_stage_path, "@"))
                 # Ensure model stage path has root prefix as stage mount will it mount it to root.
                 absolute_model_stage_path = os.path.join("/", norm_stage_path)
-                (db, schema, stage, path) = identifier.parse_schema_level_object_identifier(norm_stage_path)
+                (db, schema, stage, path) = identifier.parse_snowflake_stage_path(norm_stage_path)
                 substitutes = {
                     "image": image,
                     "predict_endpoint_name": constants.PREDICT,

snowflake/ml/model/_model_composer/model_composer.py

@@ -10,6 +10,7 @@ from absl import logging
 from packaging import requirements
 from typing_extensions import deprecated
 
+from snowflake import snowpark
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml.data import data_source
@@ -91,6 +92,7 @@ class ModelComposer:
         python_version: Optional[str] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
+        model_objective: model_types.ModelObjective = model_types.ModelObjective.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
         if not options:
@@ -119,6 +121,7 @@ class ModelComposer:
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,
+            model_objective=model_objective,
             options=options,
         )
         assert self.packager.meta is not None
@@ -185,4 +188,6 @@ class ModelComposer:
         data_sources = lineage_utils.get_data_sources(model)
         if not data_sources and sample_input_data is not None:
             data_sources = lineage_utils.get_data_sources(sample_input_data)
+            if not data_sources and isinstance(sample_input_data, snowpark.DataFrame):
+                data_sources = [data_source.DataFrameInfo(sample_input_data.queries["queries"][-1])]
         return data_sources

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

@@ -1,11 +1,11 @@
 import collections
 import copy
 import pathlib
-import warnings
 from typing import List, Optional, cast
 
 import yaml
 
+from snowflake.ml._internal import env_utils
 from snowflake.ml.data import data_source
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -47,7 +47,9 @@ class ModelManifest:
         runtime_to_use = copy.deepcopy(model_meta.runtimes["cpu"])
         runtime_to_use.name = self._DEFAULT_RUNTIME_NAME
         runtime_to_use.imports.append(str(model_rel_path) + "/")
-        runtime_dict = runtime_to_use.save(self.workspace_path)
+        runtime_dict = runtime_to_use.save(
+            self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+        )
 
         self.function_generator = function_generator.FunctionGenerator(model_dir_rel_path=model_rel_path)
         self.methods: List[model_method.ModelMethod] = []
@@ -75,13 +77,9 @@ class ModelManifest:
             )
 
         dependencies = model_manifest_schema.ModelRuntimeDependenciesDict(conda=runtime_dict["dependencies"]["conda"])
-        if options.get("include_pip_dependencies"):
-            warnings.warn(
-                "`include_pip_dependencies` specified as True: pip dependencies will be included and may not"
-                "be warehouse-compabible. The model may need to be run in SPCS.",
-                category=UserWarning,
-                stacklevel=1,
-            )
+
+        # We only want to include pip dependencies file if there are any pip requirements.
+        if len(model_meta.env.pip_requirements) > 0:
             dependencies["pip"] = runtime_dict["dependencies"]["pip"]
 
         manifest_dict = model_manifest_schema.ModelManifestDict(
@@ -137,10 +135,15 @@ class ModelManifest:
                 if isinstance(source, data_source.DatasetInfo):
                     result.append(
                         model_manifest_schema.LineageSourceDict(
-                            # Currently, we only support lineage from Dataset.
                             type=model_manifest_schema.LineageSourceTypes.DATASET.value,
                             entity=source.fully_qualified_name,
                             version=source.version,
                         )
                     )
+                elif isinstance(source, data_source.DataFrameInfo):
+                    result.append(
+                        model_manifest_schema.LineageSourceDict(
+                            type=model_manifest_schema.LineageSourceTypes.QUERY.value, entity=source.sql
+                        )
+                    )
         return result

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

@@ -57,12 +57,14 @@ class ModelFunctionInfo(TypedDict):
         target_method: actual target method name to be called.
         target_method_function_type: target method function type (FUNCTION or TABLE_FUNCTION).
         signature: The signature of the model method.
+        is_partitioned: Whether the function is partitioned.
     """
 
     name: Required[str]
     target_method: Required[str]
     target_method_function_type: Required[str]
     signature: Required[model_signature.ModelSignature]
+    is_partitioned: Required[bool]
 
 
 class ModelFunctionInfoDict(TypedDict):
@@ -78,6 +80,7 @@ class SnowparkMLDataDict(TypedDict):
 
 class LineageSourceTypes(enum.Enum):
     DATASET = "DATASET"
+    QUERY = "QUERY"
 
 
 class LineageSourceDict(TypedDict):

snowflake/ml/model/_packager/model_env/model_env.py

@@ -363,9 +363,14 @@ class ModelEnv:
         self.cuda_version = env_dict.get("cuda_version", None)
         self.snowpark_ml_version = env_dict["snowpark_ml_version"]
 
-    def save_as_dict(self, base_dir: pathlib.Path) -> model_meta_schema.ModelEnvDict:
+    def save_as_dict(
+        self, base_dir: pathlib.Path, default_channel_override: str = env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+    ) -> model_meta_schema.ModelEnvDict:
         env_utils.save_conda_env_file(
-            pathlib.Path(base_dir / self.conda_env_rel_path), self._conda_dependencies, self.python_version
+            pathlib.Path(base_dir / self.conda_env_rel_path),
+            self._conda_dependencies,
+            self.python_version,
+            default_channel_override=default_channel_override,
         )
         env_utils.save_requirements_file(
             pathlib.Path(base_dir / self.pip_requirements_rel_path), self._pip_requirements

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

@@ -1,7 +1,8 @@
+import os
 from abc import abstractmethod
-from enum import Enum
 from typing import Dict, Generic, Optional, Protocol, Type, final
 
+import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml.model import custom_model, type_hints as model_types
@@ -9,15 +10,6 @@ from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import model_meta
 
 
-class ModelObjective(Enum):
-    # This is not getting stored anywhere as metadata yet so it should be fine to slowly extend it for better coverage
-    UNKNOWN = "unknown"
-    BINARY_CLASSIFICATION = "binary_classification"
-    MULTI_CLASSIFICATION = "multi_classification"
-    REGRESSION = "regression"
-    RANKING = "ranking"
-
-
 class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
     HANDLER_TYPE: model_types.SupportedModelHandlerType
     HANDLER_VERSION: str
@@ -106,6 +98,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
         cls,
         raw_model: model_types._ModelType,
         model_meta: model_meta.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.BaseModelLoadOption],
     ) -> custom_model.CustomModel:
         """Create a custom model class wrap for unified interface when being deployed. The predict method will be
@@ -114,6 +107,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
         Args:
             raw_model: original model object,
             model_meta: The model metadata.
+            background_data: The background data used for the model explanations.
             kwargs: Options when converting the model.
 
         Raises:
@@ -131,7 +125,8 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
     _MIN_SNOWPARK_ML_VERSION: The minimal version of Snowpark ML library to use the current handler.
     _HANDLER_MIGRATOR_PLANS: Dict holding handler migrator plans.
 
-    MODELE_BLOB_FILE_OR_DIR: Relative path of the model blob file in the model subdir. Default to "model.pkl".
+    MODEL_BLOB_FILE_OR_DIR: Relative path of the model blob file in the model subdir. Default to "model.pkl".
+    BG_DATA_FILE_SUFFIX: Suffix of the background data file. Default to "_background_data.pqt".
     MODEL_ARTIFACTS_DIR: Relative path of the model artifacts dir in the model subdir. Default to "artifacts"
     DEFAULT_TARGET_METHODS: Default target methods to be logged if not specified in this kind of model. Default to
         ["predict"]
@@ -139,8 +134,10 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
         inputting sample data or model signature. Default to False.
     """
 
-    MODELE_BLOB_FILE_OR_DIR = "model.pkl"
+    MODEL_BLOB_FILE_OR_DIR = "model.pkl"
+    BG_DATA_FILE_SUFFIX = "_background_data.pqt"
     MODEL_ARTIFACTS_DIR = "artifacts"
+    EXPLAIN_ARTIFACTS_DIR = "explain_artifacts"
     DEFAULT_TARGET_METHODS = ["predict"]
     IS_AUTO_SIGNATURE = False
 
@@ -169,3 +166,23 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
                 model_meta=model_meta,
                 model_blobs_dir_path=model_blobs_dir_path,
             )
+
+    @classmethod
+    @final
+    def load_background_data(cls, name: str, model_blobs_dir_path: str) -> Optional[pd.DataFrame]:
+        """Load the model into memory.
+
+        Args:
+            name: Name of the model.
+            model_blobs_dir_path: Directory path to the whole model.
+
+        Returns:
+            Optional[pd.DataFrame], background data as pandas DataFrame, if exists.
+        """
+        data_blob_path = os.path.join(model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, name + cls.BG_DATA_FILE_SUFFIX)
+        if not os.path.exists(model_blobs_dir_path) or not os.path.isfile(data_blob_path):
+            return None
+        with open(data_blob_path, "rb") as f:
+            background_data = pd.read_parquet(f)
+
+        return background_data

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

@@ -1,9 +1,11 @@
 import json
+import warnings
 from typing import Any, Callable, Iterable, Optional, Sequence, cast
 
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
+from absl import logging
 
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_meta import model_meta
@@ -11,6 +13,17 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
 
 
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)
+
+
 def _is_callable(model: model_types.SupportedModelType, method_name: str) -> bool:
     return callable(getattr(model, method_name, None))
 
@@ -93,23 +106,42 @@ def convert_explanations_to_2D_df(
         return pd.DataFrame(explanations)
 
     if hasattr(model, "classes_"):
-        classes_list = [cl for cl in model.classes_]  # type:ignore[union-attr]
+        classes_list = [str(cl) for cl in model.classes_]  # type:ignore[union-attr]
         len_classes = len(classes_list)
         if explanations.shape[2] != len_classes:
             raise ValueError(f"Model has {len_classes} classes but explanations have {explanations.shape[2]}")
     else:
-        classes_list = [i for i in range(explanations.shape[2])]
-    exp_2d = []
-    # TODO (SNOW-1549044): Optimize this
-    for row in explanations:
-        col_list = []
-        for column in row:
-            class_explanations = {}
-            for cl, cl_exp in zip(classes_list, column):
-                if isinstance(cl, (int, np.integer)):
-                    cl = int(cl)
-                class_explanations[cl] = cl_exp
-            col_list.append(json.dumps(class_explanations))
-        exp_2d.append(col_list)
+        classes_list = [str(i) for i in range(explanations.shape[2])]
+
+    def row_to_dict(row: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        """Converts a single row to a dictionary."""
+        # convert to object or numpy creates strings of fixed length
+        return np.asarray(json.dumps(dict(zip(classes_list, row)), cls=NumpyEncoder), dtype=object)
+
+    exp_2d = np.apply_along_axis(row_to_dict, -1, explanations)
 
     return pd.DataFrame(exp_2d)
+
+
+def validate_model_objective(
+    passed_model_objective: model_types.ModelObjective, inferred_model_objective: model_types.ModelObjective
+) -> model_types.ModelObjective:
+    if (
+        passed_model_objective != model_types.ModelObjective.UNKNOWN
+        and inferred_model_objective != model_types.ModelObjective.UNKNOWN
+    ):
+        if passed_model_objective != inferred_model_objective:
+            warnings.warn(
+                f"Inferred ModelObjective: {inferred_model_objective.name} is used as model objective for this model "
+                f"version and passed argument ModelObjective: {passed_model_objective.name} is ignored",
+                category=UserWarning,
+                stacklevel=1,
+            )
+        return inferred_model_objective
+    elif inferred_model_objective != model_types.ModelObjective.UNKNOWN:
+        logging.info(
+            f"Inferred ModelObjective: {inferred_model_objective.name} is used as model objective for this model "
+            f"version"
+        )
+        return inferred_model_objective
+    return passed_model_objective

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

@@ -30,24 +30,24 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model.bin"
+    MODEL_BLOB_FILE_OR_DIR = "model.bin"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
-    def get_model_objective(cls, model: "catboost.CatBoost") -> _base.ModelObjective:
+    def get_model_objective_and_output_type(cls, model: "catboost.CatBoost") -> model_types.ModelObjective:
         import catboost
 
         if isinstance(model, catboost.CatBoostClassifier):
             num_classes = handlers_utils.get_num_classes_if_exists(model)
             if num_classes == 2:
-                return _base.ModelObjective.BINARY_CLASSIFICATION
-            return _base.ModelObjective.MULTI_CLASSIFICATION
+                return model_types.ModelObjective.BINARY_CLASSIFICATION
+            return model_types.ModelObjective.MULTI_CLASSIFICATION
         if isinstance(model, catboost.CatBoostRanker):
-            return _base.ModelObjective.RANKING
+            return model_types.ModelObjective.RANKING
         if isinstance(model, catboost.CatBoostRegressor):
-            return _base.ModelObjective.REGRESSION
+            return model_types.ModelObjective.REGRESSION
         # TODO: Find out model type from the generic Catboost Model
-        return _base.ModelObjective.UNKNOWN
+        return model_types.ModelObjective.UNKNOWN
 
     @classmethod
     def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard["catboost.CatBoost"]:
@@ -77,6 +77,8 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.CatBoostModelSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", True)
+
         import catboost
 
         assert isinstance(model, catboost.CatBoost)
@@ -105,9 +107,14 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            if kwargs.get("enable_explainability", False):
+            inferred_model_objective = cls.get_model_objective_and_output_type(model)
+            model_meta.model_objective = handlers_utils.validate_model_objective(
+                model_meta.model_objective, inferred_model_objective
+            )
+            model_objective = model_meta.model_objective
+            if enable_explainability:
                 output_type = model_signature.DataType.DOUBLE
-                if cls.get_model_objective(model) == _base.ModelObjective.MULTI_CLASSIFICATION:
+                if model_objective == model_types.ModelObjective.MULTI_CLASSIFICATION:
                     output_type = model_signature.DataType.STRING
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
@@ -115,10 +122,13 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                     target_method="predict",
                     output_return_type=output_type,
                 )
+                model_meta.function_properties = {
+                    "explain": {model_meta_schema.FunctionProperties.PARTITIONED.value: False}
+                }
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        model_save_path = os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR)
+        model_save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
 
         model.save_model(model_save_path)
 
@@ -126,7 +136,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.CatBoostModelBlobOptions({"catboost_estimator_type": model.__class__.__name__}),
         )
         model_meta.models[name] = base_meta
@@ -138,11 +148,9 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
             ],
             check_local_version=True,
         )
-        if kwargs.get("enable_explainability", False):
-            model_meta.env.include_if_absent(
-                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
-                check_local_version=True,
-            )
+        if enable_explainability:
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
         return None
@@ -188,6 +196,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
         cls,
         raw_model: "catboost.CatBoost",
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.CatBoostModelLoadOptions],
     ) -> custom_model.CustomModel:
         import catboost

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

@@ -51,6 +51,9 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         **kwargs: Unpack[model_types.CustomModelSaveOption],
     ) -> None:
         assert isinstance(model, custom_model.CustomModel)
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for custom model.")
 
         def get_prediction(
             target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
@@ -108,13 +111,13 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         # Make sure that the module where the model is defined get pickled by value as well.
         cloudpickle.register_pickle_by_value(sys.modules[model.__module__])
         pickled_obj = (model.__class__, model.context)
-        with open(os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR), "wb") as f:
+        with open(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
             cloudpickle.dump(pickled_obj, f)
         # model meta will be saved by the context manager
         model_meta.models[name] = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             handler_version=cls.HANDLER_VERSION,
             function_properties=model_meta.function_properties,
             artifacts={
@@ -183,6 +186,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         cls,
         raw_model: custom_model.CustomModel,
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.CustomModelLoadOption],
     ) -> custom_model.CustomModel:
         return raw_model

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

@@ -89,7 +89,7 @@ class HuggingFacePipelineHandler(
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model"
+    MODEL_BLOB_FILE_OR_DIR = "model"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
     DEFAULT_TARGET_METHODS = ["__call__"]
     IS_AUTO_SIGNATURE = True
@@ -133,6 +133,9 @@ class HuggingFacePipelineHandler(
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.HuggingFaceSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for huggingface model.")
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             task = model.task  # type:ignore[attr-defined]
             framework = model.framework  # type:ignore[attr-defined]
@@ -193,7 +196,7 @@ class HuggingFacePipelineHandler(
 
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             model.save_pretrained(  # type:ignore[attr-defined]
-                os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR)
+                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             )
             pipeline_params = {
                 "_batch_size": model._batch_size,  # type:ignore[attr-defined]
@@ -205,7 +208,7 @@ class HuggingFacePipelineHandler(
             with open(
                 os.path.join(
                     model_blob_path,
-                    cls.MODELE_BLOB_FILE_OR_DIR,
+                    cls.MODEL_BLOB_FILE_OR_DIR,
                     cls.ADDITIONAL_CONFIG_FILE,
                 ),
                 "wb",
@@ -213,7 +216,7 @@ class HuggingFacePipelineHandler(
                 cloudpickle.dump(pipeline_params, f)
         else:
             with open(
-                os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR),
+                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR),
                 "wb",
             ) as f:
                 cloudpickle.dump(model, f)
@@ -222,7 +225,7 @@ class HuggingFacePipelineHandler(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.HuggingFacePipelineModelBlobOptions(
                 {
                     "task": task,
@@ -329,6 +332,7 @@ class HuggingFacePipelineHandler(
         cls,
         raw_model: Union[huggingface_pipeline.HuggingFacePipelineModel, "transformers.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
     ) -> custom_model.CustomModel:
         import transformers
@@ -365,7 +369,9 @@ class HuggingFacePipelineHandler(
                     else:
                         # For others, we could offer the whole dataframe as a list.
                         # Some of them may need some conversion
-                        if isinstance(raw_model, transformers.ConversationalPipeline):
+                        if hasattr(transformers, "ConversationalPipeline") and isinstance(
+                            raw_model, transformers.ConversationalPipeline
+                        ):
                             input_data = [
                                 transformers.Conversation(
                                     text=conv_data["user_inputs"][0],
@@ -387,27 +393,33 @@ class HuggingFacePipelineHandler(
                     # Making it not aligned with the auto-inferred signature.
                     # If the output is a dict, we could blindly create a list containing that.
                     # Otherwise, creating pandas DataFrame won't succeed.
-                    if isinstance(temp_res, (dict, transformers.Conversation)) or (
-                        # For some pipeline that is expected to generate a list of dict per input
-                        # When it omit outer list, it becomes list of dict instead of list of list of dict.
-                        # We need to distinguish them from those pipelines that designed to output a dict per input
-                        # So we need to check the pipeline type.
-                        isinstance(
-                            raw_model,
-                            (
-                                transformers.FillMaskPipeline,
-                                transformers.QuestionAnsweringPipeline,
-                            ),
+                    if (
+                        (hasattr(transformers, "Conversation") and isinstance(temp_res, transformers.Conversation))
+                        or isinstance(temp_res, dict)
+                        or (
+                            # For some pipeline that is expected to generate a list of dict per input
+                            # When it omit outer list, it becomes list of dict instead of list of list of dict.
+                            # We need to distinguish them from those pipelines that designed to output a dict per input
+                            # So we need to check the pipeline type.
+                            isinstance(
+                                raw_model,
+                                (
+                                    transformers.FillMaskPipeline,
+                                    transformers.QuestionAnsweringPipeline,
+                                ),
+                            )
+                            and X.shape[0] == 1
+                            and isinstance(temp_res[0], dict)
                         )
-                        and X.shape[0] == 1
-                        and isinstance(temp_res[0], dict)
                     ):
                         temp_res = [temp_res]
 
                     if len(temp_res) == 0:
                         return pd.DataFrame()
 
-                    if isinstance(raw_model, transformers.ConversationalPipeline):
+                    if hasattr(transformers, "ConversationalPipeline") and isinstance(
+                        raw_model, transformers.ConversationalPipeline
+                    ):
                         temp_res = [[conv.generated_responses] for conv in temp_res]
 
                     # To concat those who outputs a list with one input.