def test_load_model_without_mlflow_version(): with TempDir(chdr=True) as tmp: model = Model(artifact_path="some/path", run_id="1234", mlflow_version=None) path = tmp.path("model") with open(path, "w") as out: model.to_yaml(out) loaded_model = Model.load(path) assert loaded_model.mlflow_version is None
def test_model_save_load(): m = Model(artifact_path="some/path", run_id="123", flavors={ "flavor1": { "a": 1, "b": 2 }, "flavor2": { "x": 1, "y": 2 }, }, signature=ModelSignature( inputs=Schema( [ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")])), saved_input_example_info={ "x": 1, "y": 2 }) assert m.get_input_schema() == m.signature.inputs assert m.get_output_schema() == m.signature.outputs x = Model(artifact_path="some/other/path", run_id="1234") assert x.get_input_schema() is None assert x.get_output_schema() is None n = Model(artifact_path="some/path", run_id="123", flavors={ "flavor1": { "a": 1, "b": 2 }, "flavor2": { "x": 1, "y": 2 }, }, signature=ModelSignature( inputs=Schema( [ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")])), saved_input_example_info={ "x": 1, "y": 2 }) n.utc_time_created = m.utc_time_created assert m == n n.signature = None assert m != n with TempDir() as tmp: m.save(tmp.path("model")) o = Model.load(tmp.path("model")) assert m == o assert m.to_json() == o.to_json() assert m.to_yaml() == o.to_yaml()