diff --git a/examples/ml_blood_sugar_level/ml_blood_sugar_level.py b/examples/ml_blood_sugar_level/ml_blood_sugar_level.py index 2e5365d..60df493 100644 --- a/examples/ml_blood_sugar_level/ml_blood_sugar_level.py +++ b/examples/ml_blood_sugar_level/ml_blood_sugar_level.py @@ -1,4 +1,3 @@ -import json import textwrap from abc import ABC from pathlib import Path @@ -8,8 +7,9 @@ import pandas as pd import skops.io as sio from cosy.maestro import Maestro +from sklearn.base import RegressorMixin from sklearn.datasets import load_diabetes -from sklearn.linear_model import LassoLars, LinearRegression +from sklearn.linear_model import LassoLars from sklearn.metrics import root_mean_squared_error from sklearn.model_selection import train_test_split from sklearn.preprocessing import MinMaxScaler, RobustScaler @@ -57,82 +57,67 @@ def run(self): class FitTransformScaler(CoSyLuigiTask, ABC): splitted_data = CoSyLuigiTaskParameter(TrainTestSplit) + scaler_name: str + scaler: MinMaxScaler | RobustScaler - -class FitTransformMinMaxScaler(FitTransformScaler): def output(self): return { - "scaled_x_train": luigi.LocalTarget("minmax_scaled_x_train.json"), - "scaled_x_test": luigi.LocalTarget("minmax_scaled_x_test.json"), - "scaler": luigi.LocalTarget("minmax_scaler.skops"), + "scaled_x_train": luigi.LocalTarget(f"{self.scaler_name}_scaled_x_train.json"), + "scaled_x_test": luigi.LocalTarget(f"{self.scaler_name}_scaled_x_test.json"), + "scaler": luigi.LocalTarget(f"{self.scaler_name}_scaler.skops"), } + def scale(self, data_identifier: str): + x = pd.read_json(self.input()["splitted_data"][data_identifier].path) + self.scaler.fit(x) + x_train = pd.DataFrame(self.scaler.transform(x), columns=self.scaler.feature_names_in_, index=x.index) + x_train.to_json(self.output()[f"scaled_{data_identifier}"].path) + def run(self): - print(self.input()) - x_train = pd.read_json(self.input()["splitted_data"]["x_train"].path) - scaler = MinMaxScaler() - scaler.fit(x_train) - scaled_x_train = pd.DataFrame(scaler.transform(x_train), columns=scaler.feature_names_in_, index=x_train.index) - scaled_x_train.to_json(self.output()["scaled_x_train"].path) - - x_test = pd.read_json(self.input()["splitted_data"]["x_test"].path) - scaler.transform(x_test) - scaled_x_test = pd.DataFrame(scaler.transform(x_test), columns=scaler.feature_names_in_, index=x_test.index) - scaled_x_test.to_json(self.output()["scaled_x_test"].path) + self.scale("x_train") + self.scale("x_test") with open(self.output()["scaler"].path, "wb") as outfile: - sio.dump(scaler, outfile) + sio.dump(self.scaler, outfile) -class FitTransformRobustScaler(FitTransformScaler): - def output(self): - return { - "scaled_x_train": luigi.LocalTarget("robust_scaled_x_train.json"), - "scaled_x_test": luigi.LocalTarget("robust_scaled_x_test.json"), - "scaler": luigi.LocalTarget("robust_scaler.json"), - } - - def run(self): - x_train = pd.read_json(self.input()["splitted_data"]["x_train"].path) - scaler = RobustScaler() - scaler.fit(x_train) - scaled_x_train = pd.DataFrame(scaler.transform(x_train), columns=scaler.feature_names_in_, index=x_train.index) - scaled_x_train.to_json(self.output()["scaled_x_train"].path) +class FitTransformMinMaxScaler(FitTransformScaler): + scaler_name = "minmax" + scaler = MinMaxScaler() - x_test = pd.read_json(self.input()["splitted_data"]["x_test"].path) - scaler.transform(x_test) - scaled_x_test = pd.DataFrame(scaler.transform(x_test), columns=scaler.feature_names_in_, index=x_test.index) - scaled_x_test.to_json(self.output()["scaled_x_test"].path) - with open(self.output()["scaler"].path, "wb") as outfile: - json.dump(scaler, outfile) +class FitTransformRobustScaler(FitTransformScaler): + scaler_name = "robust" + scaler = RobustScaler() class TrainRegressionModel(CoSyLuigiTask, ABC): scaled_feats = CoSyLuigiTaskParameter(FitTransformScaler) splitted_data = CoSyLuigiTaskParameter(TrainTestSplit) + model_name: str + model: RegressorMixin def _get_variant_label(self): - return Path(self.input()["scaled_feats"]["scaled_x_train"].path).stem - + return f"{self.model_name}-{Path(self.input()['scaled_feats']['scaled_x_train'].path).stem}" -class TrainLinearRegressionModel(TrainRegressionModel): def output(self): - return {"model": luigi.LocalTarget("linear_reg" + "-" + self._get_variant_label() + ".skops")} + return {"model": luigi.LocalTarget(self._get_variant_label() + ".skops")} def run(self): x_train = pd.read_json(self.input()["scaled_feats"]["scaled_x_train"].path) y_train = pd.read_json(self.input()["splitted_data"]["y_train"].path) - reg = LinearRegression() - reg.fit(x_train, y_train) + self.model.fit(x_train, y_train) + + sio.dump(self.model, self.output()["model"].path) - sio.dump(reg, self.output()["model"].path) + +class TrainLinearRegressionModel(TrainRegressionModel): + model_name = "linear_reg" class TrainLassoLarsModel(TrainRegressionModel): - def output(self): - return {"model": luigi.LocalTarget("lasso_lars" + "-" + self._get_variant_label() + ".skops")} + model_name = "lasso_lars" def run(self): x_train = pd.read_json(self.input()["scaled_feats"]["scaled_x_train"].path) diff --git a/tests/test_unique_task_parameter.py b/tests/test_unique_task_parameter.py index 8cc556f..ef39c9f 100644 --- a/tests/test_unique_task_parameter.py +++ b/tests/test_unique_task_parameter.py @@ -62,28 +62,28 @@ class EvaluatePipeline(CoSyLuigiTask): @pytest.fixture -def repo_without_constraints(): +def repo_without_constraints() -> CoSyLuigiRepo: return CoSyLuigiRepo(TrainModel, ScaleDataABC, EvaluatePipeline) @pytest.fixture -def repo_with_constraints(): +def repo_with_constraints() -> CoSyLuigiRepo: return CoSyLuigiRepo(TrainModel, ScaleDataABC, EvaluatePipelineWithUniqueScaler) @pytest.fixture -def repo_with_manual_constraints(): +def repo_with_manual_constraints() -> CoSyLuigiRepo: return CoSyLuigiRepo(TrainModel, ScaleDataABC, EvaluatePipelineWithConstraintUniqueScaler) @pytest.fixture -def repo_with_non_abstract_super(): +def repo_with_non_abstract_super() -> CoSyLuigiRepo: return CoSyLuigiRepo( TrainModel, ScaleData, ScaleDataVariantA, ScaleDataVariantB, EvaluatePipelineWithUniqueScalerAndNonAbstractSuper ) -def test_implementation_is_not_unique_across_prior_tasks(repo_without_constraints): +def test_implementation_is_not_unique_across_prior_tasks(repo_without_constraints: CoSyLuigiRepo): maestro = Maestro( repo_without_constraints.cls_repo, repo_without_constraints.taxonomy, @@ -92,7 +92,7 @@ def test_implementation_is_not_unique_across_prior_tasks(repo_without_constraint assert len(results) == 18 -def test_implementation_is_unique_across_prior_tasks(repo_with_constraints): +def test_implementation_is_unique_across_prior_tasks(repo_with_constraints: CoSyLuigiRepo): maestro = Maestro( repo_with_constraints.cls_repo, repo_with_constraints.taxonomy, @@ -103,7 +103,9 @@ def test_implementation_is_unique_across_prior_tasks(repo_with_constraints): assert result.scaled_data == result.train_model.scaled_data -def test_implementation_is_unique_across_prior_tasks_with_manual_constraint(repo_with_manual_constraints): +def test_implementation_is_unique_across_prior_tasks_with_manual_constraint( + repo_with_manual_constraints: CoSyLuigiRepo, +): maestro = Maestro( repo_with_manual_constraints.cls_repo, repo_with_manual_constraints.taxonomy, @@ -116,7 +118,9 @@ def test_implementation_is_unique_across_prior_tasks_with_manual_constraint(repo assert result.scaled_data == result.train_model.scaled_data -def test_implementation_is_unique_across_prior_tasks_with_non_abstract_super(repo_with_non_abstract_super): +def test_implementation_is_unique_across_prior_tasks_with_non_abstract_super( + repo_with_non_abstract_super: CoSyLuigiRepo, +): maestro = Maestro( repo_with_non_abstract_super.cls_repo, repo_with_non_abstract_super.taxonomy,