wehub-resource-sync
26 lines
958 B
Python
26 lines
958 B
Python
import mlflow
|
|
from mlflow.pyfunc import PythonModel, load_model, log_model
|
|
|
|
|
|
def test_unwrap_python_model_from_pyfunc_class():
|
|
class MyModel(PythonModel):
|
|
def __init__(self, param_1: str, param_2: int):
|
|
self.param_1 = param_1
|
|
self.param_2 = param_2
|
|
|
|
def predict(self, context, model_input, params=None):
|
|
return model_input + self.param_2
|
|
|
|
def upper_param_1(self):
|
|
return self.param_1.upper()
|
|
|
|
with mlflow.start_run():
|
|
model = MyModel("this is test message", 2)
|
|
model_uri = log_model("mlruns", python_model=model).model_uri
|
|
loaded_model = load_model(model_uri).unwrap_python_model()
|
|
assert isinstance(loaded_model, MyModel)
|
|
assert loaded_model.param_1 == "this is test message"
|
|
assert loaded_model.param_2 == 2
|
|
assert loaded_model.predict(None, 1) == 3
|
|
assert loaded_model.upper_param_1() == "THIS IS TEST MESSAGE"
|