PSLMssGerd/spaceapps25

1import pandas as pd2import joblib3import numpy as np4from sklearn.model_selection import train_test_split5from sklearn.preprocessing import StandardScaler6from sklearn.metrics import accuracy_score, classification_report7 8# Import the ensemble and pipeline tools9from sklearn.ensemble import VotingClassifier10from imblearn.pipeline import Pipeline11from imblearn.over_sampling import SMOTE12 13# Import the three models14from lightgbm import LGBMClassifier15from xgboost import XGBClassifier16from catboost import CatBoostClassifier17 18def train_ensemble_model():19    """20    Trains the ensemble model and SAVES performance metrics for the Streamlit app.21    """22    print("Starting ensemble model training on the merged dataset...")23 24    try:25        df = pd.read_csv('exoplanet_data_merged_for_ensemble.csv')26    except FileNotFoundError:27        print("Error: 'exoplanet_data_merged_for_ensemble.csv' not found.")28        print("Please run 'data_preparation.py' first.")29        return30 31    X = df.drop('disposition', axis=1)32    y = df['disposition']33    feature_columns = list(X.columns)34    35    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)36    print(f"Data split into {len(X_train)} training samples and {len(X_test)} testing samples.")37 38    # --- Create the Ensemble Model Pipeline ---39    clf1 = LGBMClassifier(random_state=42)40    clf2 = XGBClassifier(random_state=42, use_label_encoder=False, eval_metric='mlogloss')41    clf3 = CatBoostClassifier(random_state=42, verbose=0, loss_function='MultiClass')42 43    eclf1 = VotingClassifier(44        estimators=[('lgbm', clf1), ('xgb', clf2), ('catboost', clf3)],45        voting='soft'46    )47    48    pipeline = Pipeline([49        ('scaler', StandardScaler()),50        ('smote', SMOTE(random_state=42)),51        ('ensemble', eclf1)52    ])53 54    print("Training the full pipeline (Scaler -> SMOTE -> Ensemble)...")55    pipeline.fit(X_train, y_train)56 57    # --- Evaluate the Model and Prepare Metrics for Saving ---58    print("\nEvaluating model performance on the test set...")59    y_pred = pipeline.predict(X_test)60    accuracy = accuracy_score(y_test, y_pred)61    report = classification_report(y_test, y_pred, target_names=['False Positive', 'Candidate', 'Confirmed'], output_dict=True)62    63    print(f"Ensemble Model Accuracy: {accuracy:.4f}")64    print("\nClassification Report (Text):")65    print(classification_report(y_test, y_pred, target_names=['False Positive', 'Candidate', 'Confirmed']))66 67    # --- NEW: Get Feature Importances ---68    # We access the trained models inside the pipeline and average their feature importances69    base_models = pipeline.named_steps['ensemble'].estimators_70    importances = [model.feature_importances_ for model in base_models]71    avg_importances = np.mean(importances, axis=0)72    73    feature_importance_data = pd.DataFrame({'feature': feature_columns, 'importance': avg_importances})74    feature_importance_data = feature_importance_data.sort_values(by='importance', ascending=False)75    76    # --- NEW: Save all metrics in a dictionary ---77    metrics = {78        'accuracy': accuracy,79        'classification_report': report,80        'y_test': y_test.to_numpy(),81        'y_pred': y_pred,82        'feature_importances': feature_importance_data83    }84    85    # --- Save all Assets ---86    joblib.dump(pipeline, 'exoplanet_ensemble_model.joblib')87    joblib.dump(feature_columns, 'feature_columns.joblib')88    joblib.dump(metrics, 'training_metrics.joblib') # Save the new metrics file89    90    print("\nEnsemble pipeline, feature columns, and training metrics have been saved.")91 92if __name__ == "__main__":93    train_ensemble_model()