utanvir/telco_churn
0
1# src/components/model_trainer.py2 3import os4import sys5from dataclasses import dataclass6import numpy as np7 8from xgboost import XGBClassifier9from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, AdaBoostClassifier10from sklearn.linear_model import LogisticRegression11from sklearn.neighbors import KNeighborsClassifier12from sklearn.tree import DecisionTreeClassifier13from sklearn.svm import SVC14from sklearn.naive_bayes import GaussianNB15from sklearn.neural_network import MLPClassifier16 17from src.exception import CustomException18from src.logger import logging19from src.utils import save_object, evaluate_models # your utils, course-style API20 21@dataclass22class ModelTrainerConfig:23 trained_model_file_path = os.path.join("artifact", "model.pkl")24 25class ModelTrainer:26 def __init__(self):27 self.model_trainer_config = ModelTrainerConfig()28 29 def initiate_model_trainer(self, train_array, test_array):30 """31 Course-style: separate `models` and `params`, hand both to evaluate_models(..., param=params).32 Hyperparameter tuning depth is intentionally small; you said you'll do full HPO later.33 """34 try:35 logging.info("Split training and test input data")36 X_train, y_train, X_test, y_test = (37 train_array[:, :-1],38 train_array[:, -1],39 test_array[:, :-1],40 test_array[:, -1],41 )42 43 # ---------------------------44 # Models (sane defaults)45 # ---------------------------46 models = {47 "LogisticRegression": LogisticRegression(48 solver="lbfgs",49 penalty="l2",50 class_weight="balanced",51 max_iter=2000,52 n_jobs=-1,53 ),54 "RandomForestClassifier": RandomForestClassifier(55 n_estimators=200,56 class_weight="balanced_subsample",57 n_jobs=-1,58 random_state=42,59 ),60 "GradientBoostingClassifier": GradientBoostingClassifier(61 random_state=4262 ),63 "AdaBoostClassifier": AdaBoostClassifier(64 random_state=4265 ),66 "SVC": SVC(67 kernel="rbf",68 C=1.0,69 gamma="scale",70 probability=False, # AUC via decision_function in your evaluator71 random_state=42,72 ),73 "KNN": KNeighborsClassifier(74 n_neighbors=5,75 n_jobs=-176 ),77 "DecisionTreeClassifier": DecisionTreeClassifier(78 class_weight="balanced",79 random_state=4280 ),81 "GaussianNB": GaussianNB(),82 "MLPClassifier": MLPClassifier(83 hidden_layer_sizes=(128,),84 max_iter=500,85 random_state=4286 ),87 "XGBClassifier": XGBClassifier(88 eval_metric="logloss",89 tree_method="hist",90 n_estimators=300,91 random_state=4292 ),93 }94 95 # ---------------------------96 # Param grids 97 # Keys must match the model names above.98 # ---------------------------99 params = {100 "LogisticRegression": {101 "C": [0.1, 1.0, 10.0]102 },103 "RandomForestClassifier": {104 "n_estimators": [100, 200, 400],105 "max_depth": [None, 10, 20]106 },107 "GradientBoostingClassifier": {108 "n_estimators": [100, 200],109 "learning_rate": [0.05, 0.1],110 "max_depth": [2, 3]111 },112 "AdaBoostClassifier": {113 "n_estimators": [100, 300],114 "learning_rate": [0.1, 1.0]115 },116 "SVC": {117 "C": [0.5, 1.0, 2.0],118 "gamma": ["scale", "auto"],119 "kernel": ["rbf"]120 },121 "KNN": {122 "n_neighbors": [3, 5, 7],123 "weights": ["uniform", "distance"]124 },125 "DecisionTreeClassifier": {126 "max_depth": [None, 10, 20],127 "min_samples_split": [2, 10]128 },129 "GaussianNB": {130 # tiny smoothing sweep; cheap131 "var_smoothing": [1e-9, 1e-8, 1e-7]132 },133 "MLPClassifier": {134 "hidden_layer_sizes": [(64,), (128,)],135 "alpha": [1e-4, 1e-3]136 },137 "XGBClassifier": {138 "learning_rate": [0.05, 0.1],139 "max_depth": [3, 5],140 "n_estimators": [200, 300]141 },142 }143 144 logging.info("Evaluating models with course-style evaluator")145 model_report, fitted = evaluate_models(146 X_train=X_train, y_train=y_train,147 X_test=X_test, y_test=y_test,148 models=models, 149 param=params150 )151 if not model_report:152 raise CustomException("evaluate_models returned empty report", sys)153 154 # Rank by AUC first, f1 as tiebreaker (handles Nan Safely)155 PRIMARY, FALLBACK = "auc_roc", "f1_score"156 157 def score_tuple(m:dict):158 auc = np.nan_to_num(m.get(PRIMARY, np.nan), nan=-1.0)159 f1 = np.nan_to_num(m.get(FALLBACK, np.nan), nan=-1.0)160 return (auc, f1)161 162 163 best_model_name = max(model_report.keys(), key=lambda k: score_tuple(model_report[k]))164 best_metric = model_report[best_model_name]165 166 # Use AUC if present else F1 for theshold/logging167 best_model_score = (168 best_metric[PRIMARY]169 if not np.isnan(best_metric.get(PRIMARY, np.nan))170 else best_metric[FALLBACK]171 )172 173 # Minimal sanity check; adjust to your tolerance174 if best_model_score < 0.6:175 raise CustomException(f"No best model found (best={best_model_name}, score={best_model_score:.4f})", sys)176 177 logging.info(f"Best model: {best_model_name} | score: {best_model_score:.4f}")178 179 180 best_model = fitted[best_model_name]181 save_object(182 file_path=self.model_trainer_config.trained_model_file_path,183 obj=best_model184 )185 logging.info(f"Saved best model to {self.model_trainer_config.trained_model_file_path} | score: {best_model_score:.4f}")186 187 return (188 best_model_score,189 best_model_name190 )191 192 193 except Exception as e:194 logging.error(f"initiate_model_trainer failed: {e}")195 raise CustomException(e, sys)196