snowflake-ml-python 1.7.5__py3-none-any.whl → 1.8.0__py3-none-any.whl

snowflake/ml/model/_client/ops/model_ops.py

@@ -1,3 +1,4 @@
+import enum
 import json
 import os
 import pathlib
@@ -31,6 +32,12 @@ from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
+# An enum class to represent Create Or Alter Model SQL command.
+class ModelAction(enum.Enum):
+    CREATE = "CREATE"
+    ALTER = "ALTER"
+
+
 class ServiceInfo(TypedDict):
     name: str
     status: str
@@ -92,7 +99,7 @@ class ModelOperator:
             and self._model_version_client == __value._model_version_client
         )
 
-    def prepare_model_stage_path(
+    def prepare_model_temp_stage_path(
         self,
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
@@ -110,17 +117,28 @@ class ModelOperator:
         )
         return f"@{self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)}/model"
 
-    def create_from_stage(
+    def get_model_version_stage_path(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+    ) -> str:
+        return (
+            f"snow://model/{self._stage_client.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f"/versions/{version_name}/"
+        )
+
+    def get_model_action_from_model_name_and_version(
         self,
-        composed_model: model_composer.ModelComposer,
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> None:
-        stage_path = str(composed_model.stage_path)
+    ) -> ModelAction:
         if self.validate_existence(
             database_name=database_name,
             schema_name=schema_name,
@@ -140,6 +158,79 @@ class ModelOperator:
                     f" version {version_name} already existed."
                 )
             else:
+                return ModelAction.ALTER
+        else:
+            return ModelAction.CREATE
+
+    def add_or_create_live_version(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        model_action = self.get_model_action_from_model_name_and_version(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
+        if model_action == ModelAction.CREATE:
+            self._model_version_client.create_live_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        elif model_action == ModelAction.ALTER:
+            self._model_version_client.add_live_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            raise AssertionError(f"The model_action is {model_action}. Expected CREATE or ALTER.")
+
+    def create_from_stage(
+        self,
+        composed_model: model_composer.ModelComposer,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+        use_live_commit: Optional[bool] = False,
+    ) -> None:
+
+        if use_live_commit:
+            # if the model version is live, we can only commit the version
+            self._model_version_client.commit_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+        else:
+            stage_path = str(composed_model.stage_path)
+            # if the model version is not live,
+            # find whether the model exists and whether the version exists
+            # and then decide whether to create or alter the model
+            model_action = self.get_model_action_from_model_name_and_version(
+                database_name=database_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
+            if model_action == ModelAction.ALTER:
                 self._model_version_client.add_version_from_stage(
                     database_name=database_name,
                     schema_name=schema_name,
@@ -148,15 +239,17 @@ class ModelOperator:
                     version_name=version_name,
                     statement_params=statement_params,
                 )
-        else:
-            self._model_version_client.create_from_stage(
-                database_name=database_name,
-                schema_name=schema_name,
-                stage_path=stage_path,
-                model_name=model_name,
-                version_name=version_name,
-                statement_params=statement_params,
-            )
+            elif model_action == ModelAction.CREATE:
+                self._model_version_client.create_from_stage(
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    stage_path=stage_path,
+                    model_name=model_name,
+                    version_name=version_name,
+                    statement_params=statement_params,
+                )
+            else:
+                raise AssertionError(f"The model_action is {model_action}. Expected CREATE or ALTER.")
 
     def create_from_model_version(
         self,

snowflake/ml/model/_client/ops/service_ops.py

@@ -100,7 +100,7 @@ class ServiceOperator:
         max_instances: int,
         cpu_requests: Optional[str],
         memory_requests: Optional[str],
-        gpu_requests: Optional[str],
+        gpu_requests: Optional[Union[int, str]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,

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

@@ -1,5 +1,5 @@
 import pathlib
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import yaml
 
@@ -38,7 +38,7 @@ class ModelDeploymentSpec:
         max_instances: int,
         cpu: Optional[str],
         memory: Optional[str],
-        gpu: Optional[str],
+        gpu: Optional[Union[str, int]],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
@@ -86,7 +86,11 @@ class ModelDeploymentSpec:
             service_dict["memory"] = memory
 
         if gpu:
-            service_dict["gpu"] = gpu
+            if isinstance(gpu, int):
+                gpu_str = str(gpu)
+            else:
+                gpu_str = gpu
+            service_dict["gpu"] = gpu_str
 
         if num_workers:
             service_dict["num_workers"] = num_workers

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

@@ -71,6 +71,64 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def create_live_version(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        sql = (
+            f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" WITH LIVE VERSION {version_name.identifier()}"
+        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def add_live_version(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        sql = (
+            f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" ADD LIVE VERSION {version_name.identifier()}"
+        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def commit_version(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        sql = (
+            f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" COMMIT VERSION {version_name.identifier()}"
+        )
+
+        query_result_checker.SqlResultValidator(
+            self._session,
+            sql,
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
     # TODO(SNOW-987381): Merge with above when we have `create or alter module m [with] version v1 ...`
     def add_version_from_stage(
         self,

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

@@ -1,7 +1,7 @@
 import enum
 import json
 import textwrap
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import platform_capabilities
@@ -11,6 +11,7 @@ from snowflake.ml._internal.utils import (
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
+from snowflake.ml.model._model_composer.model_method import constants
 from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -41,7 +42,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
         image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
         image_repo_name: sql_identifier.SqlIdentifier,
-        gpu: Optional[str],
+        gpu: Optional[Union[str, int]],
         force_rebuild: bool,
         external_access_integration: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -121,6 +122,11 @@ class ServiceSQLClient(_base._BaseSQLClient):
             args_sql_list.append(input_arg_value)
         args_sql = ", ".join(args_sql_list)
 
+        wide_input = len(input_args) > constants.SNOWPARK_UDF_INPUT_COL_LIMIT
+        if wide_input:
+            input_args_sql = ", ".join(f"'{arg}', {arg.identifier()}" for arg in input_args)
+            args_sql = f"object_construct_keep_null({input_args_sql})"
+
         if platform_capabilities.PlatformCapabilities.get_instance().is_nested_function_enabled():
             fully_qualified_service_name = self.fully_qualified_object_name(
                 actual_database_name, actual_schema_name, service_name

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,8 +1,10 @@
 import pathlib
 import tempfile
 import uuid
+import warnings
 from types import ModuleType
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
+from urllib import parse
 
 from absl import logging
 from packaging import requirements
@@ -44,7 +46,13 @@ class ModelComposer:
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         self.session = session
-        self.stage_path = pathlib.PurePosixPath(stage_path)
+        self.stage_path: Union[pathlib.PurePosixPath, parse.ParseResult] = None  # type: ignore[assignment]
+        if stage_path.startswith("snow://"):
+            # The stage path is a snowflake internal stage path
+            self.stage_path = parse.urlparse(stage_path)
+        else:
+            # The stage path is a user stage path
+            self.stage_path = pathlib.PurePosixPath(stage_path)
 
         self._workspace = tempfile.TemporaryDirectory()
         self._packager_workspace = tempfile.TemporaryDirectory()
@@ -70,7 +78,20 @@ class ModelComposer:
 
     @property
     def model_stage_path(self) -> str:
-        return (self.stage_path / self.model_file_rel_path).as_posix()
+        if isinstance(self.stage_path, parse.ParseResult):
+            model_file_path = (pathlib.PosixPath(self.stage_path.path) / self.model_file_rel_path).as_posix()
+            new_url = parse.ParseResult(
+                scheme=self.stage_path.scheme,
+                netloc=self.stage_path.netloc,
+                path=str(model_file_path),
+                params=self.stage_path.params,
+                query=self.stage_path.query,
+                fragment=self.stage_path.fragment,
+            )
+            return str(parse.urlunparse(new_url))
+        else:
+            assert isinstance(self.stage_path, pathlib.PurePosixPath)
+            return (self.stage_path / self.model_file_rel_path).as_posix()
 
     @property
     def model_local_path(self) -> str:
@@ -86,6 +107,7 @@ class ModelComposer:
         metadata: Optional[Dict[str, str]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         target_platforms: Optional[List[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
         user_files: Optional[Dict[str, List[str]]] = None,
@@ -94,8 +116,32 @@ class ModelComposer:
         task: model_types.Task = model_types.Task.UNKNOWN,
         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
+        conda_dep_dict = env_utils.validate_conda_dependency_string_list(
+            conda_dependencies if conda_dependencies else []
+        )
+        is_warehouse_runnable = (
+            not conda_dep_dict
+            or all(
+                chan == env_utils.DEFAULT_CHANNEL_NAME or chan == env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+                for chan in conda_dep_dict
+            )
+        ) and (not pip_requirements)
+        disable_explainability = (
+            target_platforms and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
+        ) or (not is_warehouse_runnable)
+
+        if disable_explainability and options and options.get("enable_explainability", False):
+            warnings.warn(
+                ("The model can be deployed to Snowpark Container Services only if `enable_explainability=False`."),
+                category=UserWarning,
+                stacklevel=2,
+            )
+
         if not options:
             options = model_types.BaseModelSaveOption()
+            if disable_explainability:
+                options["enable_explainability"] = False
 
         if not snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             snowml_matched_versions = env_utils.get_matched_package_versions_in_information_schema(
@@ -120,6 +166,7 @@ class ModelComposer:
             metadata=metadata,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,

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

@@ -78,6 +78,7 @@ class ModelManifest:
             logger.info("Relaxing version constraints for dependencies in the model.")
             logger.info(f"Conda dependencies: {runtime_to_use.runtime_env.conda_dependencies}")
             logger.info(f"Pip requirements: {runtime_to_use.runtime_env.pip_requirements}")
+            logger.info(f"artifact_repository_map: {runtime_to_use.runtime_env.artifact_repository_map}")
         runtime_dict = runtime_to_use.save(
             self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
         )
@@ -124,6 +125,9 @@ class ModelManifest:
         if len(model_meta.env.pip_requirements) > 0:
             dependencies["pip"] = runtime_dict["dependencies"]["pip"]
 
+        if model_meta.env.artifact_repository_map:
+            dependencies["artifact_repository_map"] = runtime_dict["dependencies"]["artifact_repository_map"]
+
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
             runtimes={

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

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
 import enum
-from typing import Any, Dict, List, Literal, TypedDict, Union
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -20,6 +20,7 @@ class ModelMethodFunctionTypes(enum.Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: NotRequired[str]
     pip: NotRequired[str]
+    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):

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

@@ -98,7 +98,6 @@ class ModelMethod:
     def _get_method_arg_from_feature(
         feature: model_signature.BaseFeatureSpec, case_sensitive: bool = False
     ) -> model_manifest_schema.ModelMethodSignatureFieldWithName:
-        assert isinstance(feature, model_signature.FeatureSpec), "FeatureGroupSpec is not supported."
         try:
             feature_name = sql_identifier.SqlIdentifier(feature.name, case_sensitive=case_sensitive)
         except ValueError as e:

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

@@ -3,7 +3,7 @@ import itertools
 import os
 import pathlib
 import warnings
-from typing import DefaultDict, List, Optional
+from typing import DefaultDict, Dict, List, Optional
 
 from packaging import requirements, version
 
@@ -36,6 +36,7 @@ class ModelEnv:
             pip_requirements_rel_path = os.path.join(_DEFAULT_ENV_DIR, _DEFAULT_PIP_REQUIREMENTS_FILENAME)
         self.conda_env_rel_path = pathlib.PurePosixPath(pathlib.Path(conda_env_rel_path).as_posix())
         self.pip_requirements_rel_path = pathlib.PurePosixPath(pathlib.Path(pip_requirements_rel_path).as_posix())
+        self.artifact_repository_map: Optional[Dict[str, str]] = None
         self._conda_dependencies: DefaultDict[str, List[requirements.Requirement]] = collections.defaultdict(list)
         self._pip_requirements: List[requirements.Requirement] = []
         self._python_version: version.Version = version.parse(snowml_env.PYTHON_VERSION)
@@ -345,6 +346,7 @@ class ModelEnv:
     def load_from_dict(self, base_dir: pathlib.Path, env_dict: model_meta_schema.ModelEnvDict) -> None:
         self.conda_env_rel_path = pathlib.PurePosixPath(env_dict["conda"])
         self.pip_requirements_rel_path = pathlib.PurePosixPath(env_dict["pip"])
+        self.artifact_repository_map = env_dict.get("artifact_repository_map", None)
 
         self.load_from_conda_file(base_dir / self.conda_env_rel_path)
         self.load_from_pip_file(base_dir / self.pip_requirements_rel_path)
@@ -373,6 +375,7 @@ class ModelEnv:
         return {
             "conda": self.conda_env_rel_path.as_posix(),
             "pip": self.pip_requirements_rel_path.as_posix(),
+            "artifact_repository_map": self.artifact_repository_map if self.artifact_repository_map is not None else {},
             "python_version": self.python_version,
             "cuda_version": self.cuda_version,
             "snowpark_ml_version": self.snowpark_ml_version,

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

@@ -30,10 +30,7 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta as model_meta_api,
     model_meta_schema,
 )
-from snowflake.ml.model._signatures import (
-    builtins_handler,
-    utils as model_signature_utils,
-)
+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
 
@@ -66,16 +63,16 @@ def get_requirements_from_task(task: str, spcs_only: bool = False) -> List[model
     return []
 
 
-class NumpyEncoder(json.JSONEncoder):
-    # This is a JSON encoder class to ensure the output from Huggingface pipeline is JSON serializable.
-    # What it covers is numpy object.
-    def default(self, z: object) -> object:
-        if isinstance(z, np.number):
-            if np.can_cast(z, np.int64, casting="safe"):
-                return int(z)
-            elif np.can_cast(z, np.float64, casting="safe"):
-                return z.astype(np.float64)
-        return super().default(z)
+def sanitize_output(data: Any) -> Any:
+    if isinstance(data, np.number):
+        return data.item()
+    if isinstance(data, np.ndarray):
+        return sanitize_output(data.tolist())
+    if isinstance(data, list):
+        return [sanitize_output(x) for x in data]
+    if isinstance(data, dict):
+        return {k: sanitize_output(v) for k, v in data.items()}
+    return data
 
 
 @final
@@ -410,13 +407,17 @@ class HuggingFacePipelineHandler(
                                 )
                                 for conv_data in X.to_dict("records")
                             ]
-                        elif len(signature.inputs) == 1:
-                            input_data = X.to_dict("list")[signature.inputs[0].name]
                         else:
                             if isinstance(raw_model, transformers.TableQuestionAnsweringPipeline):
                                 X["table"] = X["table"].apply(json.loads)
 
-                            input_data = X.to_dict("records")
+                            # Most pipelines if it is expecting more than one arguments,
+                            # it is expecting a list of dict, where each dict has keys corresponding to the argument.
+                            if len(signature.inputs) > 1:
+                                input_data = X.to_dict("records")
+                            # If it is only expecting one argument, Then it is expecting a list of something.
+                            else:
+                                input_data = X[signature.inputs[0].name].to_list()
                         temp_res = getattr(raw_model, target_method)(input_data)
 
                     # Some huggingface pipeline will omit the outer list when there is only 1 input.
@@ -439,7 +440,6 @@ class HuggingFacePipelineHandler(
                                 ),
                             )
                             and X.shape[0] == 1
-                            and isinstance(temp_res[0], dict)
                         )
                     ):
                         temp_res = [temp_res]
@@ -453,14 +453,18 @@ class HuggingFacePipelineHandler(
                         temp_res = [[conv.generated_responses] for conv in temp_res]
 
                     # To concat those who outputs a list with one input.
-                    if builtins_handler.ListOfBuiltinHandler.can_handle(temp_res):
-                        res = builtins_handler.ListOfBuiltinHandler.convert_to_df(temp_res)
-                    elif isinstance(temp_res[0], dict):
+                    if isinstance(temp_res[0], list):
+                        if isinstance(temp_res[0][0], dict):
+                            res = pd.DataFrame({0: temp_res})
+                        else:
+                            res = pd.DataFrame(temp_res)
+                    else:
                         res = pd.DataFrame(temp_res)
-                    elif isinstance(temp_res[0], list):
-                        res = pd.DataFrame([json.dumps(output, cls=NumpyEncoder) for output in temp_res])
+
+                    if hasattr(res, "map"):
+                        res = res.map(sanitize_output)
                     else:
-                        raise ValueError(f"Cannot parse output {temp_res} from pipeline object")
+                        res = res.applymap(sanitize_output)
 
                     return model_signature_utils.rename_pandas_df(data=res, features=signature.outputs)
 

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

@@ -191,11 +191,7 @@ class KerasHandler(_base.BaseModelHandler["keras.Model"]):
                 signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
-                dtype_map = {
-                    spec.name: spec.as_dtype(force_numpy_dtype=True)
-                    for spec in signature.inputs
-                    if isinstance(spec, model_signature.FeatureSpec)
-                }
+                dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs}
 
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame: