diff --git a/src/fklearn/training/pipeline.py b/src/fklearn/training/pipeline.py index 2199f752..30efc512 100644 --- a/src/fklearn/training/pipeline.py +++ b/src/fklearn/training/pipeline.py @@ -82,8 +82,8 @@ def pipeline(data: pd.DataFrame) -> LearnerReturnType: current_data = new_data model_objects = {} - if learner_log.get("obj"): - model_objects["obj"] = learner_log.pop("obj") + if learner_log.get("object"): + model_objects["object"] = learner_log.pop("object") serialisation[learner_name].append({"fn": learner_fn, "log": learner_log, **model_objects}) logs.append(learner_log) diff --git a/tests/training/test_pipeline.py b/tests/training/test_pipeline.py index 22c3f9fb..0e97b595 100644 --- a/tests/training/test_pipeline.py +++ b/tests/training/test_pipeline.py @@ -152,7 +152,7 @@ def dummy_learner_2(df, fn, call): @fp.curry def dummy_learner_3(df, fn, call): - return fn, df, {f"dummy_learner_{call}": {}, "obj": "a"} + return fn, df, {f"dummy_learner_{call}": {}, "object": "a"} train_fn = build_pipeline( dummy_learner(fn=fn, call=1), @@ -166,10 +166,10 @@ def dummy_learner_3(df, fn, call): "features": ['id', 'x1', 'y'], "learners": {"dummy_learner": {"fn": fn, "log": {"dummy_learner_1": {}}}, "dummy_learner_2": {"fn": fn, "log": {"dummy_learner_2": {}}}, - "dummy_learner_3": {"fn": fn, "log": {"dummy_learner_3": {}}, "obj": "a"}}} + "dummy_learner_3": {"fn": fn, "log": {"dummy_learner_3": {}}, "object": "a"}}} assert log["__fkml__"] == fkml - assert "obj" not in log.keys() + assert "object" not in log.keys() @pytest.mark.parametrize("has_repeated_learners", [False, True]) @@ -247,4 +247,4 @@ def dummy_learner_2(df, fn, call): "dummy_learner_2": [{"fn": fn, "log": {"dummy_learner_2": {}}}]}} assert log["__fkml__"] == fkml - assert "obj" not in log.keys() + assert "object" not in log.keys()