Skip to content

Commit

Permalink
updates to score code + score code tests (EDMMMX-11741)
Browse files Browse the repository at this point in the history
  • Loading branch information
djm21 committed Dec 13, 2024
1 parent d945855 commit 204fd85
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 5 deletions.
39 changes: 38 additions & 1 deletion src/sasctl/pzmm/write_score_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,7 +497,7 @@ def _write_imports(
import codecs
binary_string = "<binary string>"
model = pickle.load(codecs.decode(binary_string.encode(), "base64"))
model = pickle.loads(codecs.decode(binary_string.encode(), "base64"))
"""

def _viya35_model_load(
Expand Down Expand Up @@ -562,6 +562,26 @@ def _viya35_model_load(
f'{model_id}/{model_file_name}")))'
)
else:
if pickle_type.lower() == 'pickle':
self.score_code += (
f'model_path = Path("/models/resources/viya/{model_id}'
f'")\nwith open(model_path / "{model_file_name}", '
f"\"rb\") as pickle_model:\n{'':4}model = pd.read_pickle"
"(pickle_model)\n\n"
)
"""
model_path = Path("/models/resources/viya/<UUID>")
with open(model_path / "model.pickle", "rb") as pickle_model:
model = pd.read_pickle(pickle_model)
"""
return (
f"{'':8}model_path = Path(\"/models/resources/viya/{model_id}"
f"\")\n{'':8}with open(model_path / \"{model_file_name}\", "
f"\"rb\") as pickle_model:\n{'':12}model = pd.read_pickle"
"(pickle_model)"
)

self.score_code += (
f'model_path = Path("/models/resources/viya/{model_id}'
f'")\nwith open(model_path / "{model_file_name}", '
Expand Down Expand Up @@ -658,6 +678,23 @@ def _viya4_model_load(
f"safe_mode=True)\n"
)
else:
if pickle_type.lower() == "pickle":
self.score_code += (
f"with open(Path(settings.pickle_path) / "
f'"{model_file_name}", "rb") as pickle_model:\n'
f"{'':4}model = pd.read_pickle(pickle_model)\n\n"
)
"""
with open(Path(settings.pickle_path) / "model.pickle", "rb") as pickle_model:
model = pd.read_pickle(pickle_model)
"""
return (
f"{'':8}with open(Path(settings.pickle_path) / "
f'"{model_file_name}", "rb") as pickle_model:\n'
f"{'':12}model = pd.read_pickle(pickle_model)\n\n"
)

self.score_code += (
f"with open(Path(settings.pickle_path) / "
f'"{model_file_name}", "rb") as pickle_model:\n'
Expand Down
18 changes: 14 additions & 4 deletions tests/unit/test_write_score_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,8 +118,13 @@ def test_viya35_model_load():
"""
sc = ScoreCode()
load_text = sc._viya35_model_load("1234", "normal")
assert "pickle.load(pickle_model)" in sc.score_code
assert "pickle.load(pickle_model)" in load_text
assert "pd.read_pickle(pickle_model)" in sc.score_code
assert "pd.read_pickle(pickle_model)" in load_text

sc = ScoreCode()
load_text = sc._viya35_model_load("1234", "normal", pickle_type="dill")
assert "dill.load(pickle_model)" in sc.score_code
assert "dill.load(pickle_model)" in load_text

sc = ScoreCode()
mojo_text = sc._viya35_model_load("2345", "mojo", mojo_model=True)
Expand All @@ -142,8 +147,13 @@ def test_viya4_model_load():
"""
sc = ScoreCode()
load_text = sc._viya4_model_load("normal")
assert "pickle.load(pickle_model)" in sc.score_code
assert "pickle.load(pickle_model)" in load_text
assert "pd.read_pickle(pickle_model)" in sc.score_code
assert "pd.read_pickle(pickle_model)" in load_text

sc = ScoreCode()
load_text = sc._viya35_model_load("1234", "normal", pickle_type="dill")
assert "dill.load(pickle_model)" in sc.score_code
assert "dill.load(pickle_model)" in load_text

sc = ScoreCode()
mojo_text = sc._viya4_model_load("mojo", mojo_model=True)
Expand Down

0 comments on commit 204fd85

Please sign in to comment.