AlainDeLong/End-To-End-Machine-Learning-Project

1import os2import sys3from dataclasses import dataclass4 5from sklearn.metrics import r2_score6from sklearn.linear_model import LinearRegression7from sklearn.neighbors import KNeighborsRegressor8from sklearn.tree import DecisionTreeRegressor9from sklearn.ensemble import (10    RandomForestRegressor,11    AdaBoostRegressor,12    GradientBoostingRegressor,13)14from xgboost import XGBRegressor15from catboost import CatBoostRegressor16 17from src.logger import logging18from src.exception import CustomException19from src.utils import save_object, evaluate_models20 21 22@dataclass23class ModelTrainerConfig:24    trained_model_file_path = os.path.join("artifacts", "model.pkl")25 26 27class ModelTrainer:28    def __init__(self) -> None:29        self.model_trainer_config = ModelTrainerConfig()30 31    def initiate_model_trainer(self, train_array, test_array):32        try:33            logging.info("Split training and testing input data")34            X_train, y_train, X_test, y_test = (35                train_array[:, :-1],36                train_array[:, -1],37                test_array[:, :-1],38                test_array[:, -1],39            )40 41            models = {42                "Linear Regression": LinearRegression(),43                "K-Neighbors Regressor": KNeighborsRegressor(),44                "Decision Tree Regressor": DecisionTreeRegressor(),45                "Random Forest Regressor": RandomForestRegressor(),46                "AdaBoost Regressor": AdaBoostRegressor(),47                "Gradient Boosting Regressor": GradientBoostingRegressor(),48                "XGBRegressor": XGBRegressor(),49                "CatBoosting Regressor": CatBoostRegressor(verbose=False),50            }51 52            params_grid = {53                "Linear Regression": {},54                "K-Neighbors Regressor": {},55                "Decision Tree Regressor": {56                    "criterion": [57                        "squared_error",58                        "friedman_mse",59                        "absolute_error",60                        "poisson",61                    ],62                    # 'splitter':['best','random'],63                    # 'max_features':['sqrt','log2'],64                },65                "Random Forest Regressor": {66                    # 'criterion':['squared_error', 'friedman_mse', 'absolute_error', 'poisson'],67                    # 'max_features':['sqrt','log2',None],68                    "n_estimators": [8, 16, 32, 64, 128, 256]69                },70                "AdaBoost Regressor": {71                    "learning_rate": [0.1, 0.01, 0.5, 0.001],72                    # 'loss':['linear','square','exponential'],73                    "n_estimators": [8, 16, 32, 64, 128, 256],74                },75                "Gradient Boosting Regressor": {76                    # 'loss':['squared_error', 'huber', 'absolute_error', 'quantile'],77                    "learning_rate": [0.1, 0.01, 0.05, 0.001],78                    "subsample": [0.6, 0.7, 0.75, 0.8, 0.85, 0.9],79                    # 'criterion':['squared_error', 'friedman_mse'],80                    # 'max_features':['auto','sqrt','log2'],81                    "n_estimators": [8, 16, 32, 64, 128, 256],82                },83                "XGBRegressor": {84                    "learning_rate": [0.1, 0.01, 0.05, 0.001],85                    "n_estimators": [8, 16, 32, 64, 128, 256],86                },87                "CatBoosting Regressor": {88                    "depth": [6, 8, 10],89                    "learning_rate": [0.01, 0.05, 0.1],90                    "iterations": [30, 50, 100],91                },92            }93 94            model_report: dict = evaluate_models(95                X_train=X_train,96                y_train=y_train,97                X_test=X_test,98                y_test=y_test,99                models=models,100                params_grid=params_grid,101            )102 103            # To get best model score from dict104            best_model_score = max(sorted(model_report.values()))105 106            # To get best model name from dict107            best_model_name = list(model_report.keys())[108                list(model_report.values()).index(best_model_score)109            ]110 111            best_model = models[best_model_name]112            if best_model_score < 0.6:113                raise CustomException("No best model found", sys)114            logging.info(f"Best found model on both training and testing dataset")115 116            save_object(117                file_path=self.model_trainer_config.trained_model_file_path,118                obj=best_model,119            )120 121            print(best_model_name)122            predicted = best_model.predict(X_test)123            r2_square = r2_score(y_test, predicted)124            return r2_square125 126        except Exception as e:127            raise CustomException(e, sys)128