h2o models no longer require creating a zip file in SAS Model Manager

Author: djm21

src/sasctl/pzmm/import_model.py

@@ -136,7 +136,7 @@ def import_model(
         model_prefix: str,
         project: Union[str, dict, RestObj],
         input_data: Optional[DataFrame] = None,
-        predict_method: [Callable[..., List], List[Any]] = None,
+        predict_method: Union[Callable[..., List], List[Any]] = None,
         score_metrics: Optional[List[str]] = None,
         pickle_type: str = "pickle",
         project_version: str = "latest",

src/sasctl/pzmm/pickle_model.py

@@ -7,6 +7,7 @@
 import shutil
 from pathlib import Path
 from typing import Optional, Union
+import h2o
 
 from ..utils.misc import check_if_jupyter
 
@@ -41,8 +42,7 @@ def pickle_trained_model(
         Parameters
         ---------------
         trained_model : model object, str, or Path
-            For non-H2O models, this argument contains the model variable. Otherwise,
-            this should be the file path of the MOJO file.
+            The trained model to be exported
         model_prefix : str or Path
             Variable name for the model to be displayed in SAS Open Model Manager
             (i.e. hmeqClassTree + [Score.py || .pickle]).
@@ -112,15 +112,7 @@ def pickle_trained_model(
                 binary_file.rename(binary_file.with_suffix(PICKLE))
             # For MOJO H2O models, gzip the model file and adjust the file extension
             elif is_h2o_model and pickle_path:
-                with open(Path(trained_model), "rb") as fileIn, gzip.open(
-                    Path(pickle_path) / (model_prefix + ".mojo"), "wb"
-                ) as fileOut:
-                    fileOut.writelines(fileIn)
-                if cls.notebook_output:
-                    print(
-                        f"MOJO model {model_prefix} was successfully gzipped and saved "
-                        f"to {Path(pickle_path) / (model_prefix + '.mojo')}."
-                    )
+                h2o.save_model(model=trained_model, force=True, path=pickle_path, filename=f'{model_prefix}.mojo')
             else:
                 raise ValueError(
                     "There is currently no support for file-less H2O.ai model handling."

src/sasctl/pzmm/write_score_code.py

@@ -24,7 +24,7 @@ def write_score_code(
         cls,
         model_prefix: str,
         input_data: Union[DataFrame, List[dict]],
-        predict_method: [Callable[..., List], List[Any]],
+        predict_method: Union[Callable[..., List], List[Any]],
         target_variable: Optional[str] = None,
         target_values: Optional[List] = None,
         score_metrics: Optional[List[str]] = None,
@@ -515,23 +515,8 @@ def _viya4_model_load(
         pickle_type = pickle_type if pickle_type else "pickle"
 
         if mojo_model:
-            cls.score_code += (
-                f"with gzip.open(Path(settings.pickle_path) / "
-                '"{model_file_name}", "r") as fileIn, '
-                "open(Path(settings.pickle_path) / "
-                f"\"{str(Path(model_file_name).with_suffix('.zip'))}\","
-                f" \"wb\") as fileOut:\n{'':4}shutil.copyfileobj(fileIn,"
-                " fileOut)\nos.chmod(Path(settings.pickle_path) / "
-                f"\"{str(Path(model_file_name).with_suffix('.zip'))}\""
-                ", 0o777)\nmodel = h2o.import_mojo("
-                "Path(settings.pickle_path) / "
-                f"\"{str(Path(model_file_name).with_suffix('.zip'))}\")"
-                "\n\n"
-            )
-            return (
-                f"{'':8}model = h2o.import_mojo(Path(settings.pickle_path) / "
-                f"\"{str(Path(model_file_name).with_suffix('.zip'))}\")\n\n"
-            )
+            cls.score_code += "model = h2o.import_mojo(Path(settings.pickle_path))\n\n"
+            return f"{'':8}model = h2o.import_mojo(Path(settings.pickle_path))\n\n"
         elif binary_h2o_model:
             cls.score_code += "model = h2o.load(Path(settings.pickle_path))\n\n"
             return f"{'':8}model = h2o.load(Path(settings.pickle_path))\n\n"