module.py - master - opendatas2017/P_Churn_prediction

module.py
276 строк13 KB
Igor
02 окт 2025, 16:52
02 окт 2025, 16:52481b477
import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import f1_score, precision_recall_curve, classification_report
from catboost import CatBoostClassifier
from xgboost import XGBClassifier
import optuna
from typing import List, Union, Dict, Any, Tuple


class ModelTrainer:
    def __init__(self, X_train_catboost: pd.DataFrame, y_train_catboost: pd.Series, 
                 X_train_xgboost: pd.DataFrame, y_train_xgboost: pd.Series, 
                 categorical_features: List[str] = None, n_trials: int = 50, cv_folds: int = 5):
        """
        Инициализация класса для обучения моделей с последующим подбором порога для F1.
        """
        self.X_train_catboost = X_train_catboost
        self.y_train_catboost = y_train_catboost
        self.X_train_xgboost = X_train_xgboost
        self.y_train_xgboost = y_train_xgboost
        if categorical_features is None:
            self.categorical_features = []
        elif hasattr(categorical_features, 'empty') and categorical_features.empty:
            self.categorical_features = []
        elif isinstance(categorical_features, (list, pd.Index)):
            self.categorical_features = list(categorical_features)
        else:
            self.categorical_features = []
        self.n_trials = n_trials
        self.cv_folds = cv_folds
        self.best_catboost_model = None
        self.best_xgboost_model = None
        self.best_catboost_params = None
        self.best_xgboost_params = None
        self.catboost_threshold = 0.5  # по умолчанию
        self.xgboost_threshold = 0.5   # по умолчанию

    def _find_best_threshold(self, y_true: np.ndarray, y_proba: np.ndarray) -> float:
        """
        Находит порог, максимизирующий F1, с использованием precision_recall_curve.
        """
        precision, recall, thresholds = precision_recall_curve(y_true, y_proba)
        # Убираем последнюю точку (где precision/recall не соответствуют порогу)
        f1_scores = 2 * (precision[:-1] * recall[:-1]) / (precision[:-1] + recall[:-1] + 1e-10)
        best_idx = np.argmax(f1_scores)
        return thresholds[best_idx]

    def objective_catboost(self, trial) -> float:
        """Функция для подбора гиперпараметров CatBoost"""
        params = {
            'iterations': trial.suggest_int('iterations', 100, 1000),
            'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3),
            'depth': trial.suggest_int('depth', 6, 8),
            'l2_leaf_reg': trial.suggest_float('l2_leaf_reg', 1, 10),
            'random_strength': trial.suggest_float('random_strength', 0.1, 10),
            'bagging_temperature': trial.suggest_float('bagging_temperature', 0.0, 1.0),
            'border_count': trial.suggest_int('border_count', 8, 128),
            'verbose': False,
            'random_state': 42
        }
        
        cv_scores = []
        skf = StratifiedKFold(n_splits=self.cv_folds, shuffle=True, random_state=42)
        
        X_array = self.X_train_catboost.values if hasattr(self.X_train_catboost, 'values') else self.X_train_catboost
        y_array = self.y_train_catboost.values if hasattr(self.y_train_catboost, 'values') else self.y_train_catboost
        
        for train_idx, val_idx in skf.split(X_array, y_array):
            if isinstance(self.X_train_catboost, pd.DataFrame):
                X_train_fold = self.X_train_catboost.iloc[train_idx]
                X_val_fold = self.X_train_catboost.iloc[val_idx]
                y_train_fold = self.y_train_catboost.iloc[train_idx]
                y_val_fold = self.y_train_catboost.iloc[val_idx]
            else:
                X_train_fold = self.X_train_catboost[train_idx]
                X_val_fold = self.X_train_catboost[val_idx]
                y_train_fold = self.y_train_catboost[train_idx]
                y_val_fold = self.y_train_catboost[val_idx]
            
            model = CatBoostClassifier(**params)
            model.fit(
                X_train_fold, y_train_fold,
                cat_features=self.categorical_features,
                eval_set=(X_val_fold, y_val_fold),
                early_stopping_rounds=50,
                verbose=False
            )
            
            y_pred = model.predict(X_val_fold)
            score = f1_score(y_val_fold, y_pred)
            cv_scores.append(score)
        
        return np.mean(cv_scores)
    
    def objective_xgboost(self, trial) -> float:
        """Функция для подбора гиперпараметров XGBoost"""
        params = {
            'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
            'max_depth': trial.suggest_int('max_depth', 6, 8),
            'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3),
            'subsample': trial.suggest_float('subsample', 0.6, 1.0),
            'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0),
            'reg_alpha': trial.suggest_float('reg_alpha', 0, 10),
            'reg_lambda': trial.suggest_float('reg_lambda', 0, 10),
            'min_child_weight': trial.suggest_int('min_child_weight', 1, 10),
            'random_state': 42,
            'eval_metric': 'logloss'
        }
        
        cv_scores = []
        skf = StratifiedKFold(n_splits=self.cv_folds, shuffle=True, random_state=42)
        
        X_array = self.X_train_xgboost.values if hasattr(self.X_train_xgboost, 'values') else self.X_train_xgboost
        y_array = self.y_train_xgboost.values if hasattr(self.y_train_xgboost, 'values') else self.y_train_xgboost
        
        for train_idx, val_idx in skf.split(X_array, y_array):
            if isinstance(self.X_train_xgboost, pd.DataFrame):
                X_train_fold = self.X_train_xgboost.iloc[train_idx]
                X_val_fold = self.X_train_xgboost.iloc[val_idx]
                y_train_fold = self.y_train_xgboost.iloc[train_idx]
                y_val_fold = self.y_train_xgboost.iloc[val_idx]
            else:
                X_train_fold = self.X_train_xgboost[train_idx]
                X_val_fold = self.X_train_xgboost[val_idx]
                y_train_fold = self.y_train_xgboost[train_idx]
                y_val_fold = self.y_train_xgboost[val_idx]
            
            model = XGBClassifier(**params)
            model.fit(
                X_train_fold, y_train_fold,
            )
            
            y_pred = model.predict(X_val_fold)
            score = f1_score(y_val_fold, y_pred)
            cv_scores.append(score)
        
        return np.mean(cv_scores)
    
    def optimize_hyperparameters(self) -> Tuple[Dict, Dict]:
        """Подбор гиперпараметров для обеих моделей"""
        print("Начинаем подбор гиперпараметров для CatBoost...")
        study_catboost = optuna.create_study(direction='maximize')
        study_catboost.optimize(self.objective_catboost, n_trials=self.n_trials)
        
        print("Начинаем подбор гиперпараметров для XGBoost...")
        study_xgboost = optuna.create_study(direction='maximize')
        study_xgboost.optimize(self.objective_xgboost, n_trials=self.n_trials)
        
        self.best_catboost_params = study_catboost.best_params
        self.best_xgboost_params = study_xgboost.best_params
        
        print(f"Лучшие параметры CatBoost: {self.best_catboost_params}")
        print(f"Лучшие параметры XGBoost: {self.best_xgboost_params}")
        
        return self.best_catboost_params, self.best_xgboost_params
    
    def find_best_thresholds(self) -> Tuple[float, float]:
        """
        Обучает модели с лучшими параметрами через CV и находит оптимальные пороги для F1
        на out-of-fold (OOF) предсказаниях.
        """
        print("Поиск оптимальных порогов для F1 на OOF-предсказаниях...")
        
        skf = StratifiedKFold(n_splits=self.cv_folds, shuffle=True, random_state=42)
        
        # --- CatBoost ---
        oof_proba_cat = np.zeros(len(self.y_train_catboost))
        for train_idx, val_idx in skf.split(self.X_train_catboost, self.y_train_catboost):
            X_train_fold = self.X_train_catboost.iloc[train_idx]
            X_val_fold = self.X_train_catboost.iloc[val_idx]
            y_train_fold = self.y_train_catboost.iloc[train_idx]
            
            model = CatBoostClassifier(**self.best_catboost_params, verbose=False, random_state=42)
            model.fit(
                X_train_fold, y_train_fold,
                cat_features=self.categorical_features,
                verbose=False
            )
            oof_proba_cat[val_idx] = model.predict_proba(X_val_fold)[:, 1]
        
        self.catboost_threshold = self._find_best_threshold(
            self.y_train_catboost.values, oof_proba_cat
        )
        
        # --- XGBoost ---
        oof_proba_xgb = np.zeros(len(self.y_train_xgboost))
        for train_idx, val_idx in skf.split(self.X_train_xgboost, self.y_train_xgboost):
            X_train_fold = self.X_train_xgboost.iloc[train_idx]
            X_val_fold = self.X_train_xgboost.iloc[val_idx]
            y_train_fold = self.y_train_xgboost.iloc[train_idx]
            
            model = XGBClassifier(**self.best_xgboost_params, random_state=42, eval_metric='logloss')
            model.fit(X_train_fold, y_train_fold, verbose=False)
            oof_proba_xgb[val_idx] = model.predict_proba(X_val_fold)[:, 1]
        
        self.xgboost_threshold = self._find_best_threshold(
            self.y_train_xgboost.values, oof_proba_xgb
        )
        
        # Оценка F1 с найденными порогами
        f1_cat = f1_score(
            self.y_train_catboost,
            (oof_proba_cat >= self.catboost_threshold).astype(int)
        )
        f1_xgb = f1_score(
            self.y_train_xgboost,
            (oof_proba_xgb >= self.xgboost_threshold).astype(int)
        )
        
        print(f"CatBoost — лучший порог: {self.catboost_threshold:.4f}, F1 на OOF: {f1_cat:.4f}")
        print(f"XGBoost — лучший порог: {self.xgboost_threshold:.4f}, F1 на OOF: {f1_xgb:.4f}")
        
        return self.catboost_threshold, self.xgboost_threshold

    def train_final_models(self) -> Tuple[CatBoostClassifier, XGBClassifier]:
        """Обучение финальных моделей на лучших параметрах"""
        # CatBoost
        self.best_catboost_model = CatBoostClassifier(
            **self.best_catboost_params,
            verbose=False,
            random_state=42
        )
        self.best_catboost_model.fit(
            self.X_train_catboost, 
            self.y_train_catboost,
            cat_features=self.categorical_features
        )
        
        # XGBoost
        self.best_xgboost_model = XGBClassifier(
            **self.best_xgboost_params,
            random_state=42,
            eval_metric='logloss'
        )
        self.best_xgboost_model.fit(
            self.X_train_xgboost, 
            self.y_train_xgboost
        )
        
        return self.best_catboost_model, self.best_xgboost_model

    def predict_with_threshold(self, X_catboost: pd.DataFrame = None, X_xgboost: pd.DataFrame = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        Предсказывает бинарные метки с использованием найденных оптимальных порогов.
        
        Возвращает:
            (catboost_preds, xgboost_preds)
        """
        cat_preds = None
        xgb_preds = None
        
        if X_catboost is not None and self.best_catboost_model is not None:
            proba = self.best_catboost_model.predict_proba(X_catboost)[:, 1]
            cat_preds = (proba >= self.catboost_threshold).astype(int)
        
        if X_xgboost is not None and self.best_xgboost_model is not None:
            proba = self.best_xgboost_model.predict_proba(X_xgboost)[:, 1]
            xgb_preds = (proba >= self.xgboost_threshold).astype(int)
        
        return cat_preds, xgb_preds

    def get_proba(self, X_catboost: pd.DataFrame = None, X_xgboost: pd.DataFrame = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        Возвращает вероятности (для класса 1).
        """
        cat_proba = None
        xgb_proba = None
        
        if X_catboost is not None and self.best_catboost_model is not None:
            cat_proba = self.best_catboost_model.predict_proba(X_catboost)[:, 1]
        
        if X_xgboost is not None and self.best_xgboost_model is not None:
            xgb_proba = self.best_xgboost_model.predict_proba(X_xgboost)[:, 1]
        
        return cat_proba, xgb_proba