snowflake-ml-python 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

snowflake/ml/fileset/stage_fs.py

@@ -75,6 +75,12 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         -74.00688,40.73049,-74.00563,40.70676,2\n'
     """
 
+    # Cached state marking whether we should use Snowpark file download instead of pre-signed URLs:
+    #   None -> Try pre-signed URL access, fall back to file download
+    #   True -> Use file download path without trying pre-signed URL access
+    #   False -> Use pre-signed URL access, skip download fallback on failure
+    _USE_FALLBACK_FILE_ACCESS = None
+
     def __init__(
         self,
         *,
@@ -227,6 +233,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
             SnowflakeMLException: An error occurred when the given path points to a file that cannot be found.
         """
         path = path.lstrip("/")
+        if self._USE_FALLBACK_FILE_ACCESS:
+            return self._session.file.get_stream(f"{self.stage_name}/{path}")
         cached_presigned_url = self._url_cache.get(path, None)
         if not cached_presigned_url:
             res = self._fetch_presigned_urls([path])
@@ -242,6 +250,16 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         try:
             return self._fs._open(url, mode=mode, **kwargs)
         except FileNotFoundError:
+            # Enable fallback if _USE_FALLBACK_FILE_ACCESS is True or None; set to False to disable
+            if self._USE_FALLBACK_FILE_ACCESS != False:  # noqa: E712
+                try:
+                    # Try falling back to Snowpark file read if presigned URL access failed
+                    # If fallback is successful, most likely in sproc without External Access Integration
+                    content = self._session.file.get_stream(f"{self.stage_name}/{path}")
+                    self._USE_FALLBACK_FILE_ACCESS = True
+                    return content
+                except Exception:
+                    pass
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.SNOWML_NOT_FOUND,
                 original_exception=fileset_errors.StageFileNotFoundError(f"Stage file {path} doesn't exist."),

snowflake/ml/model/_api.py

@@ -108,7 +108,7 @@ def save_model(
     model: model_types.SupportedRequireSignatureModelType,
     session: Session,
     stage_path: str,
-    sample_input: model_types.SupportedDataType,
+    sample_input_data: model_types.SupportedDataType,
     metadata: Optional[Dict[str, str]] = None,
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
@@ -125,7 +125,7 @@ def save_model(
         model: Model object.
         session: Snowpark connection session.
         stage_path: Path to the stage where model will be saved.
-        sample_input: Sample input data to infer the model signatures from.
+        sample_input_data: Sample input data to infer the model signatures from.
         metadata: Model metadata.
         conda_dependencies: List of Conda package specs. Use "[channel::]package [operator version]" syntax to specify
             a dependency. It is a recommended way to specify your dependencies using conda. When channel is not
@@ -149,7 +149,7 @@ def save_model(
     session: Session,
     stage_path: str,
     signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-    sample_input: Optional[model_types.SupportedDataType] = None,
+    sample_input_data: Optional[model_types.SupportedDataType] = None,
     metadata: Optional[Dict[str, str]] = None,
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
@@ -165,11 +165,11 @@ def save_model(
         model: Model object.
         session: Snowpark connection session.
         stage_path: Path to the stage where model will be saved.
-        signatures: Model data signatures for inputs and output for every target methods. If it is None, sample_input
-            would be used to infer the signatures if it is a local (non-SnowML modeling model).
-            If not None, sample_input should not be specified. Defaults to None.
-        sample_input: Sample input data to infer the model signatures from. If it is None, signatures must be specified
-            if it is a local (non-SnowML modeling model). If not None, signatures should not be specified.
+        signatures: Model data signatures for inputs and output for every target methods. If it is None,
+            sample_input_data would be used to infer the signatures if it is a local (non-SnowML modeling model).
+            If not None, sample_input_data should not be specified. Defaults to None.
+        sample_input_data: Sample input data to infer the model signatures from. If it is None, signatures must be
+            specified if it is a local (non-SnowML modeling model). If not None, signatures should not be specified.
             Defaults to None.
         metadata: Model metadata.
         conda_dependencies: List of Conda package specs. Use "[channel::]package [operator version]" syntax to specify
@@ -196,7 +196,7 @@ def save_model(
         name=name,
         model=model,
         signatures=signatures,
-        sample_input=sample_input,
+        sample_input_data=sample_input_data,
         metadata=metadata,
         conda_dependencies=conda_dependencies,
         pip_requirements=pip_requirements,

snowflake/ml/model/_client/model/model_version_impl.py

@@ -21,6 +21,7 @@ class ModelVersion:
     _model_ops: model_ops.ModelOperator
     _model_name: sql_identifier.SqlIdentifier
     _version_name: sql_identifier.SqlIdentifier
+    _functions: List[model_manifest_schema.ModelFunctionInfo]
 
     def __init__(self) -> None:
         raise RuntimeError("ModelVersion's initializer is not meant to be used. Use `version` from model instead.")
@@ -37,6 +38,7 @@ class ModelVersion:
         self._model_ops = model_ops
         self._model_name = model_name
         self._version_name = version_name
+        self._functions = self._get_functions()
         return self
 
     def __eq__(self, __value: object) -> bool:
@@ -239,20 +241,7 @@ class ModelVersion:
             return_functions_info.append(fi)
         return return_functions_info
 
-    @telemetry.send_api_usage_telemetry(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
-    )
-    def show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
-        """Show all functions information in a model version that is callable.
-
-        Returns:
-            A list of ModelFunctionInfo objects containing the following information:
-
-            - name: The name of the function to be called (both in SQL and in Python SDK).
-            - target_method: The original method name in the logged Python object.
-            - signature: Python signature of the original method.
-        """
+    def _get_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
@@ -274,6 +263,22 @@ class ModelVersion:
         except (NotImplementedError, ValueError, connector.DataError):
             return self._legacy_show_functions()
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
+        """Show all functions information in a model version that is callable.
+
+        Returns:
+            A list of ModelFunctionInfo objects containing the following information:
+
+            - name: The name of the function to be called (both in SQL and in Python SDK).
+            - target_method: The original method name in the logged Python object.
+            - signature: Python signature of the original method.
+        """
+        return self._functions
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -306,7 +311,7 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
-        functions: List[model_manifest_schema.ModelFunctionInfo] = self.show_functions()
+        functions: List[model_manifest_schema.ModelFunctionInfo] = self._functions
         if function_name:
             req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
             find_method: Callable[[model_manifest_schema.ModelFunctionInfo], bool] = (

snowflake/ml/model/_deploy_client/image_builds/docker_context.py

@@ -52,9 +52,7 @@ class DockerContext:
 
     def _copy_entrypoint_script_to_docker_context(self) -> None:
         """Copy gunicorn_run.sh entrypoint to docker context directory."""
-        script_path = importlib_resources.files(image_builds).joinpath(  # type: ignore[no-untyped-call]
-            constants.ENTRYPOINT_SCRIPT
-        )
+        script_path = importlib_resources.files(image_builds).joinpath(constants.ENTRYPOINT_SCRIPT)
         target_path = os.path.join(self.context_dir, constants.ENTRYPOINT_SCRIPT)
 
         with open(script_path, encoding="utf-8") as source_file, file_utils.open_file(target_path, "w") as target_file:
@@ -72,9 +70,7 @@ class DockerContext:
         """
         docker_file_path = os.path.join(self.context_dir, "Dockerfile")
         docker_file_template = (
-            importlib_resources.files(image_builds)
-            .joinpath("templates/dockerfile_template")  # type: ignore[no-untyped-call]
-            .read_text("utf-8")
+            importlib_resources.files(image_builds).joinpath("templates/dockerfile_template").read_text("utf-8")
         )
 
         if self.model_zip_stage_path is not None:
@@ -123,9 +119,7 @@ class DockerContext:
         server code.
         """
         with importlib_resources.as_file(
-            importlib_resources.files(image_builds).joinpath(  # type: ignore[no-untyped-call]
-                constants.INFERENCE_SERVER_DIR
-            )
+            importlib_resources.files(image_builds).joinpath(constants.INFERENCE_SERVER_DIR)
         ) as inference_server_folder_path:
             destination_folder_path = os.path.join(self.context_dir, constants.INFERENCE_SERVER_DIR)
             ignore_patterns = shutil.ignore_patterns("BUILD.bazel", "*test.py", "*.\\.*", "__pycache__")

snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py

@@ -114,7 +114,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         """
         kaniko_shell_script_template = (
             importlib_resources.files(image_builds)
-            .joinpath(f"templates/{constants.KANIKO_SHELL_SCRIPT_TEMPLATE}")  # type: ignore[no-untyped-call]
+            .joinpath(f"templates/{constants.KANIKO_SHELL_SCRIPT_TEMPLATE}")
             .read_text("utf-8")
         )
 
@@ -170,13 +170,11 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
 
         spec_template = (
             importlib_resources.files(image_builds)
-            .joinpath(f"templates/{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}")  # type: ignore[no-untyped-call]
+            .joinpath(f"templates/{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}")
             .read_text("utf-8")
         )
 
-        spec_file_path = os.path.join(
-            os.path.dirname(self.context_dir), f"{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}.yaml"
-        )
+        spec_file_path = os.path.join(self.context_dir, f"{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}.yaml")
 
         with file_utils.open_file(spec_file_path, "w+") as spec_file:
             assert self.artifact_stage_location.startswith("@")

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

@@ -239,14 +239,17 @@ def _validate_compute_pool(session: Session, *, options: deploy_options.SnowServ
                 ),
             )
 
-    elif state not in ["ACTIVE", "IDLE"]:
+    elif state not in ["STARTING", "ACTIVE", "IDLE"]:
         raise snowml_exceptions.SnowflakeMLException(
             error_code=error_codes.INVALID_SNOWPARK_COMPUTE_POOL,
             original_exception=RuntimeError(
-                "The compute pool you are requesting to use is not in the ACTIVE/IDLE status."
+                "The compute pool you are requesting to use is not in the ACTIVE/IDLE/STARTING status."
             ),
         )
 
+    if state in ["SUSPENDED", "STARTING"]:
+        logger.warning(f"The compute pool you are requesting is in {state} state. We are waiting it to be ready.")
+
     if options.use_gpu:
         assert options.num_gpus is not None
         request_gpus = options.num_gpus
@@ -482,14 +485,12 @@ class SnowServiceDeployment:
         if self.options.model_in_image:
             spec_template = (
                 importlib_resources.files(snowservice)
-                .joinpath("templates/service_spec_template_with_model")  # type: ignore[no-untyped-call]
+                .joinpath("templates/service_spec_template_with_model")
                 .read_text("utf-8")
             )
         else:
             spec_template = (
-                importlib_resources.files(snowservice)
-                .joinpath("templates/service_spec_template")  # type: ignore[no-untyped-call]
-                .read_text("utf-8")
+                importlib_resources.files(snowservice).joinpath("templates/service_spec_template").read_text("utf-8")
             )
 
         with _debug_aware_tmp_directory(self.debug_dir) as dir_path:

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,6 +1,7 @@
 import glob
 import pathlib
 import tempfile
+import uuid
 import zipfile
 from types import ModuleType
 from typing import Any, Dict, List, Optional
@@ -28,10 +29,9 @@ class ModelComposer:
         packager: A ModelPackager object managing the (un)packaging of a Snowflake Native Model in the MODEL object.
 
         _packager_workspace_path: A local path created from packager where it will dump all files there and ModelModel
-        will zip it. This would not required if we make directory import work.
+        will zip it. This would not be required if we make directory import work.
     """
 
-    MODEL_FILE_REL_PATH = "model.zip"
     MODEL_DIR_REL_PATH = "model"
 
     def __init__(
@@ -50,6 +50,8 @@ class ModelComposer:
         self.packager = model_packager.ModelPackager(local_dir_path=str(self._packager_workspace_path))
         self.manifest = model_manifest.ModelManifest(workspace_path=self.workspace_path)
 
+        self.model_file_rel_path = f"model-{uuid.uuid4().hex}.zip"
+
         self._statement_params = statement_params
 
     def __del__(self) -> None:
@@ -66,11 +68,11 @@ class ModelComposer:
 
     @property
     def model_stage_path(self) -> str:
-        return (self.stage_path / ModelComposer.MODEL_FILE_REL_PATH).as_posix()
+        return (self.stage_path / self.model_file_rel_path).as_posix()
 
     @property
     def model_local_path(self) -> str:
-        return str(self.workspace_path / ModelComposer.MODEL_FILE_REL_PATH)
+        return str(self.workspace_path / self.model_file_rel_path)
 
     def save(
         self,
@@ -78,7 +80,7 @@ class ModelComposer:
         name: str,
         model: model_types.SupportedModelType,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         metadata: Optional[Dict[str, str]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
@@ -106,7 +108,7 @@ class ModelComposer:
             name=name,
             model=model,
             signatures=signatures,
-            sample_input=sample_input,
+            sample_input_data=sample_input_data,
             metadata=metadata,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
@@ -130,7 +132,7 @@ class ModelComposer:
         self.manifest.save(
             session=self.session,
             model_meta=self.packager.meta,
-            model_file_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_FILE_REL_PATH),
+            model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
         )
 
@@ -156,7 +158,7 @@ class ModelComposer:
 
         # TODO (Server-side Model Rollout): Remove this section.
         model_zip_path = pathlib.Path(glob.glob(str(self.workspace_path / "*.zip"))[0])
-        ModelComposer.MODEL_FILE_REL_PATH = str(model_zip_path.relative_to(self.workspace_path))
+        self.model_file_rel_path = str(model_zip_path.relative_to(self.workspace_path))
 
         with zipfile.ZipFile(self.model_local_path, mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
             zf.extractall(path=self._packager_workspace_path)

snowflake/ml/model/_model_composer/model_method/function_generator.py

@@ -46,7 +46,7 @@ class FunctionGenerator:
 
         function_template = (
             importlib_resources.files("snowflake.ml.model._model_composer.model_method")
-            .joinpath(template_filename)  # type: ignore[no-untyped-call]
+            .joinpath(template_filename)
             .read_text()
         )
 

snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template

@@ -72,5 +72,6 @@ dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 
 # Actual table function
 class {function_name}:
-    def end_partition(df: pd.DataFrame) -> pd.DataFrame:
+    @vectorized(input=pd.DataFrame)
+    def end_partition(self, df: pd.DataFrame) -> pd.DataFrame:
         return runner(df)

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -33,8 +33,8 @@ def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     method_option = options.get("method_options", {}).get(target_method, {})
-
-    function_type = method_option.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+    global_function_type = options.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+    function_type = method_option.get("function_type", global_function_type)
     if function_type not in [function_type.value for function_type in ModelMethodFunctionTypes]:
         raise NotImplementedError
 

snowflake/ml/model/_model_composer/model_runtime/_runtime_requirements.py

@@ -5,6 +5,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<2",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.8.0,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
     "typing-extensions>=4.1.0,<5"
 ]

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

@@ -48,7 +48,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
         model: model_types._ModelType,
         model_meta: model_meta.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.BaseModelSaveOption],
     ) -> None:
@@ -59,7 +59,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
             model: The model object.
             model_meta: The model metadata.
             model_blobs_dir_path: Directory path to the model.
-            sample_input: Sample input to infer the signatures from.
+            sample_input_data: Sample input to infer the signatures from.
             is_sub_model: Flag to show if it is a sub model, a sub model does not need signature.
             kwargs: Additional saving options.
 

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

@@ -14,19 +14,19 @@ def validate_signature(
     model: model_types.SupportedRequireSignatureModelType,
     model_meta: model_meta.ModelMetadata,
     target_methods: Iterable[str],
-    sample_input: Optional[model_types.SupportedDataType],
+    sample_input_data: Optional[model_types.SupportedDataType],
     get_prediction_fn: Callable[[str, model_types.SupportedLocalDataType], model_types.SupportedLocalDataType],
 ) -> model_meta.ModelMetadata:
     if model_meta.signatures:
         validate_target_methods(model, list(model_meta.signatures.keys()))
         return model_meta
 
-    # In this case sample_input should be available, because of the check in save_model.
+    # In this case sample_input_data should be available, because of the check in save_model.
     assert (
-        sample_input is not None
+        sample_input_data is not None
     ), "Model signature and sample input are None at the same time. This should not happen with local model."
-    trunc_sample_input = model_signature._truncate_data(sample_input)
-    if isinstance(sample_input, SnowparkDataFrame):
+    trunc_sample_input = model_signature._truncate_data(sample_input_data)
+    if isinstance(sample_input_data, SnowparkDataFrame):
         # Added because of Any from missing stubs.
         trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
         local_sample_input = snowpark_handler.SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)

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

@@ -45,27 +45,27 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         model: "custom_model.CustomModel",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.CustomModelSaveOption],
     ) -> None:
         assert isinstance(model, custom_model.CustomModel)
 
         def get_prediction(
-            target_method_name: str, sample_input: model_types.SupportedLocalDataType
+            target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
         ) -> model_types.SupportedLocalDataType:
             target_method = getattr(model, target_method_name, None)
             assert callable(target_method) and inspect.ismethod(target_method)
             target_method = target_method.__func__
 
-            if not isinstance(sample_input, pd.DataFrame):
-                sample_input = model_signature._convert_local_data_to_df(sample_input)
+            if not isinstance(sample_input_data, pd.DataFrame):
+                sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
 
             if inspect.iscoroutinefunction(target_method):
                 with anyio.start_blocking_portal() as portal:
-                    predictions_df = portal.call(target_method, model, sample_input)
+                    predictions_df = portal.call(target_method, model, sample_input_data)
             else:
-                predictions_df = target_method(model, sample_input)
+                predictions_df = target_method(model, sample_input_data)
             return predictions_df
 
         if not is_sub_model:
@@ -73,7 +73,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 model=model,
                 model_meta=model_meta,
                 target_methods=[method.__name__ for method in model._get_infer_methods()],
-                sample_input=sample_input,
+                sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
 

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

@@ -128,7 +128,7 @@ class HuggingFacePipelineHandler(
         model: Union[huggingface_pipeline.HuggingFacePipelineModel, "transformers.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.HuggingFaceSaveOptions],
     ) -> None:
@@ -174,7 +174,7 @@ class HuggingFacePipelineHandler(
                     )
             else:
                 handlers_utils.validate_target_methods(model, target_methods)
-                if sample_input is not None:
+                if sample_input_data is not None:
                     warnings.warn(
                         "Inferring model signature from sample input for hugggingface pipeline is not supported. "
                         + "Model signature will automatically be inferred from pipeline task. "

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

@@ -54,7 +54,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
         model: llm.LLM,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.LLMSaveOptions],
     ) -> None:

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

@@ -92,7 +92,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         model: "mlflow.pyfunc.PyFuncModel",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.MLFlowSaveOptions],
     ) -> None:

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

@@ -21,6 +21,7 @@ from snowflake.ml.model._signatures import (
 )
 
 if TYPE_CHECKING:
+    import sentence_transformers  # noqa: F401
     import torch
 
 
@@ -44,9 +45,11 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         cls,
         model: model_types.SupportedModelType,
     ) -> TypeGuard["torch.nn.Module"]:
-        return type_utils.LazyType("torch.nn.Module").isinstance(model) and not type_utils.LazyType(
-            "torch.jit.ScriptModule"
-        ).isinstance(model)
+        return (
+            type_utils.LazyType("torch.nn.Module").isinstance(model)
+            and not type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
+            and not type_utils.LazyType("sentence_transformers.SentenceTransformer").isinstance(model)
+        )
 
     @classmethod
     def cast_model(
@@ -66,7 +69,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         model: "torch.nn.Module",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        sample_input: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.PyTorchSaveOptions],
     ) -> None:
@@ -82,18 +85,18 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
             )
 
             def get_prediction(
-                target_method_name: str, sample_input: "model_types.SupportedLocalDataType"
+                target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
             ) -> model_types.SupportedLocalDataType:
-                if not pytorch_handler.SeqOfPyTorchTensorHandler.can_handle(sample_input):
-                    sample_input = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(
-                        model_signature._convert_local_data_to_df(sample_input)
+                if not pytorch_handler.SeqOfPyTorchTensorHandler.can_handle(sample_input_data):
+                    sample_input_data = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(
+                        model_signature._convert_local_data_to_df(sample_input_data)
                     )
 
                 model.eval()
                 target_method = getattr(model, target_method_name, None)
                 assert callable(target_method)
                 with torch.no_grad():
-                    predictions_df = target_method(*sample_input)
+                    predictions_df = target_method(*sample_input_data)
 
                 if isinstance(predictions_df, torch.Tensor):
                     predictions_df = [predictions_df]
@@ -103,7 +106,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
                 model=model,
                 model_meta=model_meta,
                 target_methods=target_methods,
-                sample_input=sample_input,
+                sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )