import os import sys import joblib import argparse import collections import mlflow import numpy as np import lightgbm as lgbm from sklearn.svm import SVC from sklearn.decomposition import PCA from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler from sklearn.linear_model import LogisticRegression from sklearn.experimental import enable_iterative_imputer from sklearn.metrics import accuracy_score, f1_score, make_scorer from sklearn.impute import KNNImputer, SimpleImputer, IterativeImputer from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier from sklearn.model_selection import cross_validate, train_test_split, GridSearchCV, KFold from data_utils import read_csv_file, get_data_from_data_frame def load_ml_model(pkl_file_name): model_pipeline = mlflow.sklearn.load_model(pkl_file_name) return model_pipeline def get_imputer(imputer_type): # setup parameter search space for different imputers imputer, imputer_params = None, None if imputer_type == "simple": imputer = SimpleImputer() imputer_params = { "imputer__strategy": ["mean", "median", "most_frequent"], } elif imputer_type == "knn": imputer = KNNImputer() imputer_params = { "imputer__n_neighbors": [5, 7], "imputer__weights": ["uniform", "distance"], } elif imputer_type == "iterative": imputer = IterativeImputer() imputer_params = { "imputer__initial_strategy": ["mean", "median", "most_frequent"], "imputer__imputation_order": ["ascending", "descending"], } else: print(f"unidentified option for arg, imputer_type: {imputer_type}") sys.exit(0) return imputer, imputer_params def get_scaler(): scaler = StandardScaler() return scaler def get_pca(max_num_feats): pca = PCA() pca_params = { "pca__n_components": np.arange(2, max_num_feats+1), } return pca, pca_params def get_classifier(classifier_type): # setup parameter search space for different classifiers classifier, classifier_params = None, None if classifier_type == "ada_boost": classifier = AdaBoostClassifier() classifier_params = { "classifier__learning_rate": [0.5, 1, 1.5, 2, 2.5, 3], "classifier__n_estimators": [100, 200, 500], } elif classifier_type == "log_reg": classifier = LogisticRegression(max_iter=200, solver="saga") classifier_params = { "classifier__penalty": ["l1", "l2", "elasticnet"], "classifier__class_weight": [None, "balanced"], "classifier__C": [0.1, 0.5, 1, 2], "classifier__l1_ratio": np.arange(0.1, 1, 0.1), } elif classifier_type == "random_forest": classifier = RandomForestClassifier() classifier_params = { "classifier__n_estimators": [100, 250], "classifier__criterion": ["gini", "entropy"], "classifier__max_depth": [None, 10, 25, 50, 75], "classifier__min_samples_leaf": [1, 5, 10, 20], "classifier__min_samples_split": [2, 3, 4, 5], } elif classifier_type == "svc": classifier = SVC() classifier_params = { "classifier__C": [0.5, 1, 1.5, 2, 2.5], "classifier__kernel": ["linear", "poly", "rbf", "sigmoid"], "classifier__degree": [2, 3, 4], } elif classifier_type == "light_gbm": classifier = lgbm.LGBMClassifier( boosting_type="gbdt", objective="binary", metric="auc", verbosity=-1) classifier_params = { "classifier__num_leaves": [15, 31, 63, 127, 255], "classifier__learning_rate": [0.1, 0.5, 1, 2], "classifier__n_estimators": [100, 500, 1000], "classifier__reg_lambda": [0.1, 0.5, 1], "classifier__min_data_in_leaf": [10, 20, 30, 50], } else: print(f"unidentified option for arg, classifier_type: {classifier_type}") sys.exit(0) return classifier, classifier_params def get_pipeline_params(imputer_params, classifier_params): pipeline_params = {**imputer_params, **classifier_params} return pipeline_params def train_model(df_train, df_test, imputer_type, classifier_type): # get data arrays from the data frame for train and test sets X_train, Y_train = get_data_from_data_frame(df_train) X_test, Y_test = get_data_from_data_frame(df_test) # get imputer and its params imputer, imputer_params = get_imputer(imputer_type) # get classifier and its params classifier, classifier_params = get_classifier(classifier_type) # get the pipeline params pipeline_params = get_pipeline_params(imputer_params, classifier_params) print("\n" + "-"*100) # build the model pipeline if classifier_type == "svc" or classifier_type == "log_reg": scaler = get_scaler() pca, pca_params = get_pca(X_train.shape[1]) print(f"Started training the model with the imputer: {imputer_type}, preprocessing: std_scaler + pca, classifier: {classifier_type}") pipeline = Pipeline([("imputer", imputer), ("scaler", scaler), ("pca", pca), ("classifier", classifier)]) pipeline_params = get_pipeline_params(pipeline_params, pca_params) else: print(f"Started training the model with the imputer: {imputer_type}, classifier: {classifier_type}") pipeline = Pipeline([("imputer", imputer), ("classifier", classifier)]) print("Model pipeline params space: ") print(pipeline_params) print("-"*100) # setup grid search with k-fold cross validation k_fold_cv = KFold(n_splits=5, shuffle=True, random_state=4) grid_cv = GridSearchCV(pipeline, pipeline_params, scoring="f1", cv=k_fold_cv) grid_cv.fit(X_train, Y_train) # get the cross validation score and the params for the best estimator cv_best_estimator = grid_cv.best_estimator_ cv_best_f1 = grid_cv.best_score_ cv_best_params = grid_cv.best_params_ # predict and compute train set metrics Y_train_pred = cv_best_estimator.predict(X_train) train_f1 = f1_score(Y_train, Y_train_pred) train_acc = accuracy_score(Y_train, Y_train_pred) # predict and compute test set metrics Y_test_pred = cv_best_estimator.predict(X_test) test_f1 = f1_score(Y_test, Y_test_pred) test_acc = accuracy_score(Y_test, Y_test_pred) print("\n" + "-"*50) # begin mlflow logging for the best estimator mlflow.set_experiment("water_potability") experiment = mlflow.get_experiment_by_name("water_potability") print(f"Started mlflow logging for the best estimator") with mlflow.start_run(experiment_id=experiment.experiment_id): # log the model and the metrics mlflow.sklearn.log_model(cv_best_estimator, f"{imputer_type}_{classifier_type}") mlflow.sklearn.save_model(cv_best_estimator, f"{imputer_type}_{classifier_type}") mlflow.log_params(cv_best_params) mlflow.log_metric("cv_f1_score", cv_best_f1) mlflow.log_metric("train_f1_score", train_f1) mlflow.log_metric("train_acc_score", train_acc) mlflow.log_metric("test_f1_score", test_f1) mlflow.log_metric("test_acc_score", test_acc) # end mlflow logging mlflow.end_run() print(f"Completed mlflow logging for the best estimator") print("-"*50) return def init_and_train_model(ARGS): df_csv = read_csv_file(ARGS.file_csv) df_train, df_test = train_test_split(df_csv, test_size=0.1, random_state=4) num_samples_train = df_train.shape[0] num_samples_test = df_test.shape[0] print("\n" + "-"*40) print("Num samples after splitting the dataset") print("-"*40) print(f"train: {num_samples_train}, test: {num_samples_test}") print("\n" + "-"*40) print("A few samples from train data") print("-"*40) print(df_train.head()) if ARGS.is_train: train_model(df_train, df_test, ARGS.imputer_type, ARGS.classifier_type) return def main(): file_csv = "dataset/water_potability.csv" classifier_type = "ada_boost" imputer_type = "knn" is_train = 1 parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument("--file_csv", default=file_csv, type=str, help="full path to dataset csv file") parser.add_argument("--is_train", default=is_train, type=int, choices=[0, 1], help="to train or not") parser.add_argument("--classifier_type", default=classifier_type, type=str, choices=["ada_boost", "log_reg", "random_forest", "svc", "light_gbm"], help="classifier to be used in the training model pipeline") parser.add_argument("--imputer_type", default=imputer_type, type=str, choices=["simple", "knn", "iterative"], help="imputer to be used in the training model pipeline") ARGS, unparsed = parser.parse_known_args() init_and_train_model(ARGS) return if __name__ == "__main__": main()