Spaces:
Build error
Build error
| # src/models/ensemble.py | |
| import numpy as np | |
| import pandas as pd | |
| import lightgbm as lgb | |
| import xgboost as xgb | |
| from catboost import CatBoostClassifier | |
| from sklearn.linear_model import LogisticRegression | |
| import joblib | |
| import os | |
| class SuperLearnerEnsemble: | |
| """ | |
| Ensemble of XGBoost, LightGBM, CatBoost with a logistic regression meta-learner. | |
| Now stores feature names to prevent shape mismatches. | |
| """ | |
| def __init__(self): | |
| self.xgb_model = None | |
| self.lgb_model = None | |
| self.cat_model = None | |
| self.meta_learner = LogisticRegression(max_iter=1000) | |
| self.is_fitted = False | |
| self.feature_names = None # NEW: store feature names | |
| def fit(self, X_train, y_train, X_val=None, y_val=None, feature_names=None): | |
| """ | |
| Train base models and meta-learner. | |
| Args: | |
| X_train: Training features (numpy array or DataFrame) | |
| y_train: Training targets | |
| X_val: Validation features for meta-learner training | |
| y_val: Validation targets for meta-learner training | |
| feature_names: List of feature names (optional, extracted if X is DataFrame) | |
| """ | |
| # Extract feature names if available | |
| if feature_names is not None: | |
| self.feature_names = feature_names | |
| elif hasattr(X_train, 'columns'): | |
| self.feature_names = list(X_train.columns) | |
| else: | |
| self.feature_names = [f"feature_{i}" for i in range(X_train.shape[1])] | |
| # Convert to numpy if DataFrame | |
| if hasattr(X_train, 'values'): | |
| X_train = X_train.values | |
| if X_val is not None and hasattr(X_val, 'values'): | |
| X_val = X_val.values | |
| # Initialize base models | |
| self.xgb_model = xgb.XGBClassifier( | |
| n_estimators=200, max_depth=6, learning_rate=0.05, | |
| objective='binary:logistic', random_state=42, | |
| use_label_encoder=False, eval_metric='logloss' | |
| ) | |
| self.lgb_model = lgb.LGBMClassifier( | |
| n_estimators=200, max_depth=6, learning_rate=0.05, | |
| random_state=42, verbose=-1 | |
| ) | |
| self.cat_model = CatBoostClassifier( | |
| n_estimators=200, depth=6, learning_rate=0.05, | |
| random_state=42, verbose=0 | |
| ) | |
| # Train base models | |
| self.xgb_model.fit(X_train, y_train) | |
| self.lgb_model.fit(X_train, y_train) | |
| self.cat_model.fit(X_train, y_train) | |
| # Meta-learner training | |
| if X_val is not None and y_val is not None: | |
| meta_X = np.column_stack([ | |
| self.xgb_model.predict_proba(X_val)[:, 1], | |
| self.lgb_model.predict_proba(X_val)[:, 1], | |
| self.cat_model.predict_proba(X_val)[:, 1] | |
| ]) | |
| self.meta_learner.fit(meta_X, y_val) | |
| else: | |
| split = int(len(X_train) * 0.8) | |
| X_train_base, X_meta = X_train[:split], X_train[split:] | |
| y_train_base, y_meta = y_train[:split], y_train[split:] | |
| self.xgb_model.fit(X_train_base, y_train_base) | |
| self.lgb_model.fit(X_train_base, y_train_base) | |
| self.cat_model.fit(X_train_base, y_train_base) | |
| meta_X = np.column_stack([ | |
| self.xgb_model.predict_proba(X_meta)[:, 1], | |
| self.lgb_model.predict_proba(X_meta)[:, 1], | |
| self.cat_model.predict_proba(X_meta)[:, 1] | |
| ]) | |
| self.meta_learner.fit(meta_X, y_meta) | |
| self.is_fitted = True | |
| return self | |
| def predict_proba(self, X): | |
| if not self.is_fitted: | |
| raise ValueError("Model not fitted.") | |
| # Convert to numpy if DataFrame | |
| if hasattr(X, 'values'): | |
| X = X.values | |
| xgb_proba = self.xgb_model.predict_proba(X)[:, 1] | |
| lgb_proba = self.lgb_model.predict_proba(X)[:, 1] | |
| cat_proba = self.cat_model.predict_proba(X)[:, 1] | |
| meta_X = np.column_stack([xgb_proba, lgb_proba, cat_proba]) | |
| return self.meta_learner.predict_proba(meta_X)[:, 1] | |
| def predict(self, X): | |
| proba = self.predict_proba(X) | |
| return (proba >= 0.5).astype(int) | |
| def save(self, path_prefix: str): | |
| os.makedirs(os.path.dirname(path_prefix), exist_ok=True) | |
| joblib.dump(self.xgb_model, f"{path_prefix}_xgb.joblib") | |
| joblib.dump(self.lgb_model, f"{path_prefix}_lgb.joblib") | |
| joblib.dump(self.cat_model, f"{path_prefix}_cat.joblib") | |
| joblib.dump(self.meta_learner, f"{path_prefix}_meta.joblib") | |
| # Save feature names | |
| joblib.dump(self.feature_names, f"{path_prefix}_features.joblib") | |
| print(f"✅ Ensemble saved with prefix '{path_prefix}'") | |
| def load(self, path_prefix: str): | |
| self.xgb_model = joblib.load(f"{path_prefix}_xgb.joblib") | |
| self.lgb_model = joblib.load(f"{path_prefix}_lgb.joblib") | |
| self.cat_model = joblib.load(f"{path_prefix}_cat.joblib") | |
| self.meta_learner = joblib.load(f"{path_prefix}_meta.joblib") | |
| # Load feature names if available | |
| try: | |
| self.feature_names = joblib.load(f"{path_prefix}_features.joblib") | |
| except: | |
| self.feature_names = None | |
| print("⚠️ No feature names file found.") | |
| self.is_fitted = True | |
| print(f"✅ Ensemble loaded from '{path_prefix}_*.joblib'") | |
| return self |