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

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)
@@ -113,7 +114,33 @@ class ModelEnv:
             self._snowpark_ml_version = version.parse(snowpark_ml_version)
 
     def include_if_absent(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
-        """Append requirements into model env if absent.
+        """Append requirements into model env if absent. Depending on the environment, requirements may be added
+        to either the pip requirements or conda dependencies.
+
+        Args:
+            pkgs: A list of ModelDependency namedtuple to be appended.
+            check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
+        """
+        if self.pip_requirements and not self.conda_dependencies and pkgs:
+            pip_pkg_reqs: List[str] = []
+            warnings.warn(
+                (
+                    "Dependencies specified from pip requirements."
+                    " This may prevent model deploying to Snowflake Warehouse."
+                ),
+                category=UserWarning,
+                stacklevel=2,
+            )
+            for conda_req_str, pip_name in pkgs:
+                _, conda_req = env_utils._validate_conda_dependency_string(conda_req_str)
+                pip_req = requirements.Requirement(f"{pip_name}{conda_req.specifier}")
+                pip_pkg_reqs.append(str(pip_req))
+            self._include_if_absent_pip(pip_pkg_reqs, check_local_version)
+        else:
+            self._include_if_absent_conda(pkgs, check_local_version)
+
+    def _include_if_absent_conda(self, pkgs: List[ModelDependency], check_local_version: bool = False) -> None:
+        """Append requirements into model env conda dependencies if absent.
 
         Args:
             pkgs: A list of ModelDependency namedtuple to be appended.
@@ -134,8 +161,8 @@ class ModelEnv:
                 if show_warning_message:
                     warnings.warn(
                         (
-                            f"Basic dependency {req_to_add.name} specified from PIP requirements."
-                            + " This may prevent model deploying to Snowflake Warehouse."
+                            f"Basic dependency {req_to_add.name} specified from pip requirements."
+                            " This may prevent model deploying to Snowflake Warehouse."
                         ),
                         category=UserWarning,
                         stacklevel=2,
@@ -157,11 +184,11 @@ class ModelEnv:
                         stacklevel=2,
                     )
 
-    def include_if_absent_pip(self, pkgs: List[str], check_local_version: bool = False) -> None:
-        """Append pip requirements into model env if absent.
+    def _include_if_absent_pip(self, pkgs: List[str], check_local_version: bool = False) -> None:
+        """Append pip requirements into model env pip requirements if absent.
 
         Args:
-            pkgs: A list of string to be appended in pip requirement.
+            pkgs: A list of strings to be appended to pip environment.
             check_local_version: Flag to indicate if it is required to pin to local version. Defaults to False.
         """
 
@@ -187,25 +214,6 @@ class ModelEnv:
                 self._conda_dependencies[channel].remove(spec)
 
     def generate_env_for_cuda(self) -> None:
-        if self.cuda_version is None:
-            return
-
-        cuda_spec = env_utils.find_dep_spec(
-            self._conda_dependencies, self._pip_requirements, conda_pkg_name="cuda", remove_spec=False
-        )
-        if cuda_spec and not cuda_spec.specifier.contains(self.cuda_version):
-            raise ValueError(
-                "The CUDA requirement you specified in your conda dependencies or pip requirements is"
-                " conflicting with CUDA version required. Please do not specify CUDA dependency using conda"
-                " dependencies or pip requirements."
-            )
-
-        if not cuda_spec:
-            self.include_if_absent(
-                [ModelDependency(requirement=f"nvidia::cuda=={self.cuda_version}.*", pip_name="cuda")],
-                check_local_version=False,
-            )
-
         xgboost_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=True
         )
@@ -236,7 +244,7 @@ class ModelEnv:
                 check_local_version=False,
             )
 
-            self.include_if_absent_pip(["bitsandbytes>=0.41.0"], check_local_version=False)
+            self._include_if_absent_pip(["bitsandbytes>=0.41.0"], check_local_version=False)
 
     def relax_version(self) -> None:
         """Relax the version requirements for both conda dependencies and pip requirements.
@@ -252,7 +260,9 @@ class ModelEnv:
         self._pip_requirements = list(map(env_utils.relax_requirement_version, self._pip_requirements))
 
     def load_from_conda_file(self, conda_env_path: pathlib.Path) -> None:
-        conda_dependencies_dict, pip_requirements_list, python_version = env_utils.load_conda_env_file(conda_env_path)
+        conda_dependencies_dict, pip_requirements_list, python_version, cuda_version = env_utils.load_conda_env_file(
+            conda_env_path
+        )
 
         for channel, channel_dependencies in conda_dependencies_dict.items():
             if channel != env_utils.DEFAULT_CHANNEL_NAME:
@@ -310,6 +320,9 @@ class ModelEnv:
         if python_version:
             self.python_version = python_version
 
+        if cuda_version:
+            self.cuda_version = cuda_version
+
     def load_from_pip_file(self, pip_requirements_path: pathlib.Path) -> None:
         pip_requirements_list = env_utils.load_requirements_file(pip_requirements_path)
 
@@ -333,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)
@@ -342,12 +356,17 @@ class ModelEnv:
         self.snowpark_ml_version = env_dict["snowpark_ml_version"]
 
     def save_as_dict(
-        self, base_dir: pathlib.Path, default_channel_override: str = env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
+        self,
+        base_dir: pathlib.Path,
+        default_channel_override: str = env_utils.SNOWFLAKE_CONDA_CHANNEL_URL,
+        is_gpu: Optional[bool] = False,
     ) -> model_meta_schema.ModelEnvDict:
+        cuda_version = self.cuda_version if is_gpu else None
         env_utils.save_conda_env_file(
             pathlib.Path(base_dir / self.conda_env_rel_path),
             self._conda_dependencies,
             self.python_version,
+            cuda_version,
             default_channel_override=default_channel_override,
         )
         env_utils.save_requirements_file(
@@ -356,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/_utils.py

@@ -39,7 +39,7 @@ def _is_callable(model: model_types.SupportedModelType, method_name: str) -> boo
 
 
 def get_truncated_sample_data(
-    sample_input_data: model_types.SupportedDataType, length: int = 100
+    sample_input_data: model_types.SupportedDataType, length: int = 100, is_for_modeling_model: bool = False
 ) -> model_types.SupportedLocalDataType:
     trunc_sample_input = model_signature._truncate_data(sample_input_data, length=length)
     local_sample_input: model_types.SupportedLocalDataType = None
@@ -47,6 +47,8 @@ def get_truncated_sample_data(
         # 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)
+        if is_for_modeling_model:
+            local_sample_input.columns = trunc_sample_input.columns
     else:
         local_sample_input = trunc_sample_input
     return local_sample_input
@@ -58,13 +60,15 @@ def validate_signature(
     target_methods: Iterable[str],
     sample_input_data: Optional[model_types.SupportedDataType],
     get_prediction_fn: Callable[[str, model_types.SupportedLocalDataType], model_types.SupportedLocalDataType],
+    is_for_modeling_model: bool = False,
 ) -> model_meta.ModelMetadata:
     if model_meta.signatures:
         validate_target_methods(model, list(model_meta.signatures.keys()))
         if sample_input_data is not None:
-            local_sample_input = get_truncated_sample_data(sample_input_data)
+            local_sample_input = get_truncated_sample_data(
+                sample_input_data, is_for_modeling_model=is_for_modeling_model
+            )
             for target_method in model_meta.signatures.keys():
-
                 model_signature_inst = model_meta.signatures.get(target_method)
                 if model_signature_inst is not None:
                     # strict validation the input signature
@@ -77,7 +81,7 @@ def validate_signature(
     assert (
         sample_input_data is not None
     ), "Model signature and sample input are None at the same time. This should not happen with local model."
-    local_sample_input = get_truncated_sample_data(sample_input_data)
+    local_sample_input = get_truncated_sample_data(sample_input_data, is_for_modeling_model=is_for_modeling_model)
     for target_method in target_methods:
         predictions_df = get_prediction_fn(target_method, local_sample_input)
         sig = model_signature.infer_signature(

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
@@ -146,6 +143,10 @@ class HuggingFacePipelineHandler(
             framework = getattr(model, "framework", None)
             batch_size = getattr(model, "batch_size", None)
 
+        has_tokenizer = getattr(model, "tokenizer", None) is not None
+        has_feature_extractor = getattr(model, "feature_extractor", None) is not None
+        has_image_preprocessor = getattr(model, "image_preprocessor", None) is not None
+
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             params = {
                 **model._preprocess_params,  # type:ignore[attr-defined]
@@ -234,6 +235,9 @@ class HuggingFacePipelineHandler(
                 {
                     "task": task,
                     "batch_size": batch_size if batch_size is not None else 1,
+                    "has_tokenizer": has_tokenizer,
+                    "has_feature_extractor": has_feature_extractor,
+                    "has_image_preprocessor": has_image_preprocessor,
                 }
             ),
         )
@@ -308,6 +312,14 @@ class HuggingFacePipelineHandler(
         if os.path.isdir(model_blob_file_or_dir_path):
             import transformers
 
+            additional_pipeline_params = {}
+            if model_blob_options.get("has_tokenizer", False):
+                additional_pipeline_params["tokenizer"] = model_blob_file_or_dir_path
+            if model_blob_options.get("has_feature_extractor", False):
+                additional_pipeline_params["feature_extractor"] = model_blob_file_or_dir_path
+            if model_blob_options.get("has_image_preprocessor", False):
+                additional_pipeline_params["image_preprocessor"] = model_blob_file_or_dir_path
+
             with open(
                 os.path.join(
                     model_blob_file_or_dir_path,
@@ -324,6 +336,7 @@ class HuggingFacePipelineHandler(
                 model=model_blob_file_or_dir_path,
                 trust_remote_code=True,
                 torch_dtype="auto",
+                **additional_pipeline_params,
                 **device_config,
             )
 
@@ -394,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.
@@ -423,7 +440,6 @@ class HuggingFacePipelineHandler(
                                 ),
                             )
                             and X.shape[0] == 1
-                            and isinstance(temp_res[0], dict)
                         )
                     ):
                         temp_res = [temp_res]
@@ -437,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

@@ -0,0 +1,226 @@
+import os
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+
+import cloudpickle
+import numpy as np
+import pandas as pd
+from packaging import version
+from typing_extensions import TypeGuard, Unpack
+
+from snowflake.ml._internal import type_utils
+from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
+from snowflake.ml.model._packager.model_env import model_env
+from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
+from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
+from snowflake.ml.model._packager.model_meta import (
+    model_blob_meta,
+    model_meta as model_meta_api,
+)
+from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
+
+if TYPE_CHECKING:
+    import keras
+
+
+@final
+class KerasHandler(_base.BaseModelHandler["keras.Model"]):
+    """Handler for Keras v3 model.
+
+    Currently keras.Model based classes are supported.
+    """
+
+    HANDLER_TYPE = "keras"
+    HANDLER_VERSION = "2025-01-01"
+    _MIN_SNOWPARK_ML_VERSION = "1.7.5"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+
+    MODEL_BLOB_FILE_OR_DIR = "model.keras"
+    CUSTOM_OBJECT_SAVE_PATH = "custom_objects.pkl"
+    DEFAULT_TARGET_METHODS = ["predict"]
+
+    @classmethod
+    def can_handle(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> TypeGuard["keras.Model"]:
+        if not type_utils.LazyType("keras.Model").isinstance(model):
+            return False
+        import keras
+
+        return version.parse(keras.__version__) >= version.parse("3.0.0")
+
+    @classmethod
+    def cast_model(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> "keras.Model":
+        import keras
+
+        assert isinstance(model, keras.Model)
+
+        return cast(keras.Model, model)
+
+    @classmethod
+    def save_model(
+        cls,
+        name: str,
+        model: "keras.Model",
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.TensorflowSaveOptions],
+    ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for Tensorflow model.")
+
+        import keras
+
+        assert isinstance(model, keras.Model)
+
+        if not is_sub_model:
+            target_methods = handlers_utils.get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=cls.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                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)
+                predictions_df = target_method(sample_input_data)
+
+                if (
+                    type_utils.LazyType("tensorflow.Tensor").isinstance(predictions_df)
+                    or type_utils.LazyType("tensorflow.Variable").isinstance(predictions_df)
+                    or type_utils.LazyType("torch.Tensor").isinstance(predictions_df)
+                ):
+                    predictions_df = [predictions_df]
+
+                return predictions_df
+
+            model_meta = handlers_utils.validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input_data=sample_input_data,
+                get_prediction_fn=get_prediction,
+            )
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        os.makedirs(model_blob_path, exist_ok=True)
+        save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+        model.save(save_path)
+
+        custom_object_save_path = os.path.join(model_blob_path, cls.CUSTOM_OBJECT_SAVE_PATH)
+        custom_objects = keras.saving.get_custom_objects()
+        with open(custom_object_save_path, "wb") as f:
+            cloudpickle.dump(custom_objects, f)
+
+        base_meta = model_blob_meta.ModelBlobMeta(
+            name=name,
+            model_type=cls.HANDLER_TYPE,
+            handler_version=cls.HANDLER_VERSION,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
+        )
+        model_meta.models[name] = base_meta
+        model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
+
+        dependencies = [
+            model_env.ModelDependency(requirement="keras>=3", pip_name="keras"),
+        ]
+        keras_backend = keras.backend.backend()
+        if keras_backend == "tensorflow":
+            dependencies.append(model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"))
+        elif keras_backend == "torch":
+            dependencies.append(model_env.ModelDependency(requirement="pytorch", pip_name="torch"))
+        elif keras_backend == "jax":
+            dependencies.append(model_env.ModelDependency(requirement="jax", pip_name="jax"))
+        else:
+            raise ValueError(f"Unsupported backend {keras_backend}")
+
+        model_meta.env.include_if_absent(
+            dependencies,
+            check_local_version=True,
+        )
+        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+
+    @classmethod
+    def load_model(
+        cls,
+        name: str,
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
+    ) -> "keras.Model":
+        import keras
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        model_blobs_metadata = model_meta.models
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+
+        custom_object_save_path = os.path.join(model_blob_path, cls.CUSTOM_OBJECT_SAVE_PATH)
+        with open(custom_object_save_path, "rb") as f:
+            custom_objects = cloudpickle.load(f)
+        load_path = os.path.join(model_blob_path, model_blob_filename)
+        m = keras.models.load_model(load_path, custom_objects=custom_objects, safe_mode=False)
+
+        return cast(keras.Model, m)
+
+    @classmethod
+    def convert_as_custom_model(
+        cls,
+        raw_model: "keras.Model",
+        model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
+    ) -> custom_model.CustomModel:
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "keras.Model",
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "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}
+
+                @custom_model.inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    res = getattr(raw_model, target_method)(X.astype(dtype_map), verbose=0)
+
+                    if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
+                        # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
+                        # return a list of ndarrays. We need to deal them separately
+                        df = numpy_handler.SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                return fn
+
+            type_method_dict = {}
+            for target_method_name, sig in model_meta.signatures.items():
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
+
+            _KerasModel = type(
+                "_KerasModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _KerasModel
+
+        _KerasModel = _create_custom_model(raw_model, model_meta)
+        keras_model = _KerasModel(custom_model.ModelContext())
+
+        return keras_model