File size: 6,434 Bytes
a4a265d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.pipeline import Pipeline
from lightgbm import LGBMClassifier, LGBMRegressor
from typing import Tuple, Any, Dict
from src.monitoring.logger import get_logger

logger = get_logger(__name__)

class ModelSelectionEngine:
    """Dynamically selects the optimal model and its base parameters based on data sparsity, size and task."""

    @staticmethod
    def get_model_candidate(n_samples: int, n_features: int, is_sparse: bool, is_regression: bool = False, aggressive: bool = False) -> Any:
        """Selects a base model candidate for the dataset traits."""
        from sklearn.ensemble import StackingClassifier, StackingRegressor
        from sklearn.linear_model import LogisticRegression, Ridge
        
        # Aggressive Mode: Use XGBoost and CatBoost in the stack
        if aggressive:
            from xgboost import XGBClassifier, XGBRegressor
            from catboost import CatBoostClassifier, CatBoostRegressor
            
            if is_regression:
                logger.info("🔥 AGGRESSIVE MODE: Building Elite Regression Stack (RF + LGBM + XGB + CatBoost)")
                estimators = [
                    ('rf', RandomForestRegressor(n_estimators=100, n_jobs=-1, random_state=42)),
                    ('lgbm', LGBMRegressor(n_estimators=100, n_jobs=-1, random_state=42, verbose=-1)),
                    ('xgb', XGBRegressor(n_estimators=100, n_jobs=-1, random_state=42, verbosity=0)),
                    ('cat', CatBoostRegressor(n_estimators=100, random_state=42, verbose=0))
                ]
                return StackingRegressor(
                    estimators=estimators,
                    final_estimator=Ridge(),
                    cv=3,
                    n_jobs=-1
                )
            else:
                logger.info("🔥 AGGRESSIVE MODE: Building Elite Classification Stack (RF + LGBM + XGB + CatBoost)")
                estimators = [
                    ('rf', RandomForestClassifier(n_estimators=100, n_jobs=-1, random_state=42, class_weight='balanced')),
                    ('lgbm', LGBMClassifier(n_estimators=100, n_jobs=-1, random_state=42, verbose=-1, class_weight='balanced')),
                    ('xgb', XGBClassifier(n_estimators=100, n_jobs=-1, random_state=42, verbosity=0)),
                    ('cat', CatBoostClassifier(n_estimators=100, random_state=42, verbose=0))
                ]
                return StackingClassifier(
                    estimators=estimators,
                    final_estimator=LogisticRegression(),
                    cv=3,
                    n_jobs=-1
                )

        if is_regression:
            if n_samples < 1000:
                logger.info("Selecting RandomForestRegressor for small regression.")
                return RandomForestRegressor(n_jobs=-1, random_state=42)
            else:
                logger.info("Selecting StackingRegressor (RF + LGBM + Ridge) for peak accuracy.")
                estimators = [
                    ('rf', RandomForestRegressor(n_estimators=100, n_jobs=-1, random_state=42)),
                    ('lgbm', LGBMRegressor(n_estimators=100, n_jobs=-1, random_state=42, verbose=-1)),
                    ('ridge', Ridge(alpha=1.0))
                ]
                return StackingRegressor(
                    estimators=estimators,
                    final_estimator=Ridge(),
                    cv=3,
                    n_jobs=-1
                )
        else:
            if is_sparse and n_samples < 5000:
                logger.info("Selecting Logistic Regression for sparse data.")
                return LogisticRegression(solver='liblinear', max_iter=1000, class_weight='balanced')
            elif n_samples < 1000:
                logger.info("Selecting Random Forest for robust small-scale learning.")
                return RandomForestClassifier(n_jobs=-1, random_state=42, class_weight='balanced')
            else:
                logger.info("Selecting StackingClassifier (RF + LGBM + LogReg) for peak accuracy.")
                estimators = [
                    ('rf', RandomForestClassifier(n_estimators=100, n_jobs=-1, random_state=42, class_weight='balanced')),
                    ('lgbm', LGBMClassifier(n_estimators=100, n_jobs=-1, random_state=42, verbose=-1, class_weight='balanced')),
                    ('lr', LogisticRegression(max_iter=1000, class_weight='balanced'))
                ]
                return StackingClassifier(
                    estimators=estimators,
                    final_estimator=LogisticRegression(),
                    cv=3,
                    n_jobs=-1
                )

    @staticmethod
    def create_unified_pipeline(preprocessor: Any, model: Any) -> Pipeline:
        """Wraps preprocessor and model into a single versioned pipeline."""
        return Pipeline([
            ('preprocessor', preprocessor),
            ('model', model)
        ])

    @staticmethod
    def get_param_grid(model: Any) -> Dict[str, Any]:
        """Returns search space for the selected model."""
        from sklearn.ensemble import VotingClassifier, VotingRegressor, StackingClassifier, StackingRegressor
        
        if isinstance(model, (VotingClassifier, VotingRegressor, StackingClassifier, StackingRegressor)):
            grid = {}
            # Standard ensembles have (name, estimator) pairs in .estimators
            est_names = [name for name, _ in model.estimators]
            
            if 'rf' in est_names:
                grid['model__rf__n_estimators'] = [50, 100]
            if 'lgbm' in est_names:
                grid['model__lgbm__learning_rate'] = [0.05, 0.1]
            if 'xgb' in est_names:
                grid['model__xgb__max_depth'] = [3, 6]
            if 'cat' in est_names:
                grid['model__cat__depth'] = [4, 6]
            return grid
        
        if isinstance(model, LogisticRegression):
            return {'model__C': [0.1, 1.0, 10.0]}
        if isinstance(model, (RandomForestClassifier, RandomForestRegressor)):
            return {'model__n_estimators': [50, 100]}
        if isinstance(model, (LGBMClassifier, LGBMRegressor)):
            return {'model__learning_rate': [0.05, 0.1]}
        if isinstance(model, Ridge):
            return {'model__alpha': [0.1, 1.0, 10.0]}
        return {}