A2PRIVCOMP-Synthetic-Data-Generation-Models-Validation/train_sns.py at main · Vicomtech/A2PRIVCOMP-Synthetic-Data-Generation-Models-Validation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
# Stdlib imports
import os

# Third-party app imports
import shutil
import traceback
import pandas as pd
from tqdm import tqdm
from dotenv import load_dotenv
import shutil

# Imports from your apps
from src.LogClass import LogClass

from src.sdg_models.with_privacy_guarantes.smartnoise.AIMModel import (
    AIMModel,
    AIMModelParams
)
from src.sdg_models.with_privacy_guarantes.smartnoise.MSTModel import (
    MSTModelParams,
    MSTModel
)

from src.sdg_models.AbstractModelMLFlow import MLFlowExperimentInfo
from src.sd_evaluation.utility.UtilityEvaluator import UtilityEvaluator

logger = LogClass()

# def train_model(model: AbstractModel, data: pd.DataFrame):
#     model.initialize_model(input_data=data)
#     model.train_model(input_data=data)
#     return model


def train_evaluate_models(dataset: str):
    # Load environment variables
    real_data_path = os.environ.get("REAL_DATA_PATH")
    train_data_path = os.environ.get("TRAIN_DATA_PATH")
    test_data_path = os.environ.get("TEST_DATA_PATH")

    sdg_models = {
        # "MST-no-DP": MSTModel(
        #     model_params=MSTModelParams(epsilon=0),
        #     experiment_info=MLFlowExperimentInfo(),
        # ),
        # "DP-MST": MSTModel(
        #     model_params=MSTModelParams(),
        #     experiment_info=MLFlowExperimentInfo(),
        # ),
        "AIM-no-DP": AIMModel(
            model_params=AIMModelParams(epsilon=0,preprocessor_eps=0),
            experiment_info=MLFlowExperimentInfo(),
        ),
        # "DP-AIM": AIMModel(
        #     model_params=AIMModelParams(),
        #     experiment_info=MLFlowExperimentInfo(),
        # )
    }

    # Iterate over al csv files in train data path
    for file_name in tqdm(os.listdir(real_data_path)):

        # Verify if file is a csv file
        if file_name.endswith(".csv") and dataset in file_name:

            train_file_path = os.path.join(train_data_path, file_name)
            test_file_path = os.path.join(test_data_path, file_name)
            dataset_name = file_name.split(".")[0]
            logger.log_info(f" Reading {file_name} dataset...")

            # Read the datafile
            try:
                data_train = pd.read_csv(train_file_path, index_col=False)
                data_test = pd.read_csv(test_file_path, index_col=False)
            except Exception as e:
                logger.log_error(f" Failed to read {file_name}: {e}")

            logger.log_info(
                f" Performing TRTR for {file_name} dataset..."
            )
            utility_evaluator = UtilityEvaluator()
            utility_evaluator.perform_train_real_test_real(
                train_data=data_train,
                test_data=data_test,
            )

            # Iterate over all models
            for model_name in tqdm(sdg_models.keys()):

                # Create the experiment info data class
                experiment_info = MLFlowExperimentInfo(
                    name=dataset_name,
                    tags={"model": model_name},
                    run_name=model_name,
                    model_name=model_name,
                    tracking_uri="http://localhost:5000",
                )
                sdg_models[model_name].experiment_info = experiment_info

                # Train the SDG model
                try:
                    logger.log_info(
                        f" Training {model_name} model for {file_name} dataset..."
                    )
                    sdg_models[model_name].train_evaluate_model(
                        train_data=data_train, test_data=data_test, utility_evaluator=utility_evaluator,
                    )
                    logger.log_info(
                        f" {model_name} model trained and evaluated for {file_name} dataset!"
                    )

                except Exception as e:
                    logger.log_error(
                        f"Failed to train and evaluate {model_name} model for {file_name} dataset...: {e}"
                    )
                    traceback.print_exc()


if __name__ == "__main__":
    load_dotenv()
    # without hpc
    train_evaluate_models(dataset="aml_data")
    # train_evaluate_models(dataset="breast_cancer_wisonsin")
    # train_evaluate_models(dataset="heart_disease")
    # train_evaluate_models(dataset="pima_indians_diabetes")
    # train_evaluate_models(dataset="survey lung cancer")
    # train_evaluate_models(dataset="uci_heart_disease")

    # with hpc
    # train_evaluate_models(dataset="adult")
    # train_evaluate_models(dataset="Uber")
    # train_evaluate_models(dataset="Steel_industry_data_no_time")
    # train_evaluate_models(dataset="ADAS Mobility Data_no_time")
    # train_evaluate_models(dataset="bank")