mboukabous/train_unsupervised

1 2"""3train_clustering_model.py4 5A script to train clustering models (K-Means, DBSCAN, Gaussian Mixture, etc.).6It can optionally perform hyperparameter tuning using silhouette score,7trains the model, saves it, and visualizes clusters if requested.8"""9 10import os11import sys12import argparse13import importlib14import pandas as pd15import numpy as np16import joblib17 18from sklearn import datasets19from sklearn.metrics import silhouette_score20from sklearn.preprocessing import LabelEncoder21import matplotlib.pyplot as plt22import seaborn as sns23from timeit import default_timer as timer24 25def main(args):26    # Change to the project root if needed27    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))28    os.chdir(project_root)29    sys.path.insert(0, project_root)30 31    # Optional: import the unsupervised hyperparameter tuning function32    from utils.unsupervised_hyperparameter_tuning import clustering_hyperparameter_tuning33 34    # Dynamically import the chosen clustering model module35    model_module_path = f"models.unsupervised.clustering.{args.model_module}"36    model_module = importlib.import_module(model_module_path)37 38    # Retrieve the estimator and param grid from the model file39    estimator = model_module.estimator40    param_grid = getattr(model_module, 'param_grid', {})41    default_scoring = getattr(model_module, 'default_scoring', 'silhouette')  # fallback42 43    # Prepare results directory44    if args.results_path is None:45        # e.g., 'results/KMeans_Clustering'46        args.results_path = os.path.join('results', f"{estimator.__class__.__name__}_Clustering")47    os.makedirs(args.results_path, exist_ok=True)48    49    # Prepare model directory50    if args.model_path is None:51        # e.g., 'saved_model/KMeans_Clustering'52        args.model_path = os.path.join('saved_models', f"{estimator.__class__.__name__}_Clustering")53    os.makedirs(args.model_path, exist_ok=True)54 55    # Load data from CSV56    df = pd.read_csv(args.data_path)57    print(f"Data loaded from {args.data_path}, initial shape: {df.shape}")58 59    # Drop empty columns60    df = df.dropna(axis='columns', how='all')61    print("After dropping empty columns:", df.shape)62 63    # Drop specified columns if any64    if args.drop_columns:65        drop_cols = [col.strip() for col in args.drop_columns.split(',') if col.strip()]66        df = df.drop(columns=drop_cols, errors='ignore')67        print(f"Dropped columns: {drop_cols} | New shape: {df.shape}")68 69    # Select specified columns if any70    if args.select_columns:71        keep_cols = [col.strip() for col in args.select_columns.split(',') if col.strip()]72        # Keep only these columns (intersection with what's in df)73        df = df[keep_cols]74        print(f"Selected columns: {keep_cols} | New shape: {df.shape}")75 76    # For each non-numeric column, apply label encoding77    for col in df.columns:78        if df[col].dtype == 'object':79            le = LabelEncoder()80            df[col] = le.fit_transform(df[col])81 82    # Convert DataFrame to NumPy array for clustering83    X = df.values84    print(f"Final shape after dropping/selecting columns and encoding: {X.shape}")85 86    # If user wants hyperparam tuning87    if args.tune:88        print("Performing hyperparameter tuning...")89        best_model, best_params = clustering_hyperparameter_tuning(90            X, estimator, param_grid, scoring=default_scoring, cv=args.cv_folds91        )92        estimator = best_model  # the fitted best model93        print("Best Params:", best_params)94    else:95        # Just fit the model directly96        print("No hyperparameter tuning; fitting model with default parameters...")97        start_time = timer()98        estimator.fit(X)99        end_time = timer()100        print(f"Training time (no tuning): {end_time - start_time:.2f}s")101 102    # Ensure the model is fitted at this point103    model_output_path = os.path.join(args.model_path, "best_model.pkl")104    joblib.dump(estimator, model_output_path)105    print(f"Model saved to {model_output_path}")106 107    # Evaluate using silhouette if possible108    # Some clusterers use .labels_, others require .predict(X)109    if hasattr(estimator, 'labels_'):110        labels = estimator.labels_111    else:112        labels = estimator.predict(X)  # e.g. KMeans, GaussianMixture113 114    unique_labels = set(labels)115    if len(unique_labels) > 1:116        sil = silhouette_score(X, labels)117        print(f"Silhouette Score: {sil:.4f}")118        pd.DataFrame({"Silhouette": [sil]}).to_csv(119            os.path.join(args.results_path, "metrics.csv"), index=False120        )121    else:122        print("Only one cluster found; silhouette score not meaningful.")123 124    # Visualization125    if args.visualize:126        print("Creating cluster visualization...")127 128        # If X has more than 2 dims, do PCA => 2D129        if X.shape[1] > 2:130            from sklearn.decomposition import PCA131            pca = PCA(n_components=2)132            X_2d = pca.fit_transform(X)133            var_ratio = pca.explained_variance_ratio_134            pc1_var = var_ratio[0] * 100135            pc2_var = var_ratio[1] * 100136            x_label = f"PC1 ({pc1_var:.2f}% var)"137            y_label = f"PC2 ({pc2_var:.2f}% var)"138        elif X.shape[1] == 2:139            # If we know 'df' and shape matches, label with col names140            if df.shape[1] == 2:141                x_label = df.columns[0]142                y_label = df.columns[1]143            else:144                x_label = "Feature 1"145                y_label = "Feature 2"146            X_2d = X147        else:148            # 1D or 0D => skip149            if X.shape[1] == 1:150                print("Only 1 feature available; cannot create a 2D scatter plot.")151            else:152                print("No features available for plotting.")153            return154 155        plt.figure(figsize=(6, 5))156        plt.scatter(X_2d[:, 0], X_2d[:, 1], c=labels, cmap='viridis', s=30)157        plt.title(f"{estimator.__class__.__name__} Clusters")158        plt.xlabel(x_label)159        plt.ylabel(y_label)160 161        # Save the figure162        plot_path = os.path.join(args.results_path, "clusters.png")163        plt.savefig(plot_path)164        plt.show()165        print(f"Cluster plot saved to {plot_path}")166 167if __name__ == "__main__":168    parser = argparse.ArgumentParser(description="Train a clustering model.")169    parser.add_argument('--model_module', type=str, required=True,170                        help='Name of the clustering model module (e.g. kmeans, dbscan, etc.).')171    parser.add_argument('--data_path', type=str, required=True,172                        help='Path to the CSV dataset.')173    parser.add_argument('--model_path', type=str, default=None,174                        help='Path to save the trained model.')175    parser.add_argument('--results_path', type=str, default=None,176                        help='Directory to save results (metrics, plots).')177    parser.add_argument('--cv_folds', type=int, default=5,178                        help='Number of folds for hyperparam tuning.')179    parser.add_argument('--tune', action='store_true',180                        help='Perform hyperparameter tuning with silhouette score.')181    parser.add_argument('--visualize', action='store_true',182                        help='Generate a 2D visualization of the clusters.')183    parser.add_argument('--drop_columns', type=str, default='',184                        help='Comma-separated column names to drop from the dataset.')185    parser.add_argument('--select_columns', type=str, default='',186                        help='Comma-separated column names to keep (ignore all others).')187    args = parser.parse_args()188    main(args)189