snowflake.ml.model.custom_model.inference_api¶

snowflake.ml.model.custom_model.inference_api(func: Callable[[Concatenate[CustomModelType, DataFrame, InferenceParams]], DataFrame]) → Callable[[Concatenate[CustomModelType, DataFrame, InferenceParams]], DataFrame]¶

Decorator to mark a method as an inference API in a CustomModel.

Methods decorated with @inference_api are exposed as callable inference endpoints when the model is logged to the registry and deployed.

The decorated method must accept a pandas DataFrame as its first argument (after self) and return a pandas DataFrame. It may also accept additional keyword-only parameters that can be passed at inference time.

Parameters:

func – The method to decorate.

Returns:

The decorated function with inference API metadata.

Example:

class MyModel(CustomModel):
    @inference_api
    def predict(self, input_df: pd.DataFrame) -> pd.DataFrame:
        return pd.DataFrame({"output": input_df["feature"] * 2})

    @inference_api
    def predict_with_params(
        self,
        input_df: pd.DataFrame,
        *,
        temperature: float = 1.0,
    ) -> pd.DataFrame:
        return pd.DataFrame({"output": input_df["feature"] * temperature})

Copy