diff --git a/src/sasctl/pzmm/import_model.py b/src/sasctl/pzmm/import_model.py index bfa9c6f7..3b2b9218 100644 --- a/src/sasctl/pzmm/import_model.py +++ b/src/sasctl/pzmm/import_model.py @@ -213,7 +213,7 @@ def import_model( target_values: Optional[List[str]] = None, overwrite_project_properties: Optional[bool] = False, target_index: Optional[int] = None, - preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None, + preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None, **kwargs, ) -> Tuple[RestObj, Union[dict, str, Path]]: """ diff --git a/src/sasctl/pzmm/write_score_code.py b/src/sasctl/pzmm/write_score_code.py index 9a6e80f8..e990d461 100644 --- a/src/sasctl/pzmm/write_score_code.py +++ b/src/sasctl/pzmm/write_score_code.py @@ -36,7 +36,7 @@ def write_score_code( score_cas: Optional[bool] = True, score_code_path: Union[Path, str, None] = None, target_index: Optional[int] = None, - preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None, + preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None, **kwargs, ) -> Union[dict, None]: """ @@ -763,7 +763,7 @@ def _predict_method( missing_values: Optional[Any] = None, statsmodels_model: Optional[bool] = False, tf_model: Optional[bool] = False, - preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None, + preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None, ) -> None: """ Write the model prediction section of the score code. @@ -2260,7 +2260,7 @@ def _viya35_score_code_import( mr.update_model(model) return mas_code, cas_code - def _add_preprocess_code(self, preprocess_function: Callable[DataFrame, DataFrame]): + def _add_preprocess_code(self, preprocess_function: Callable[[DataFrame], DataFrame]): """ Places the given preprocess function, which must both take a DataFrame as an argument and return a DataFrame, into the score code. If the preprocess function does not