iBrokeTheCode/Multimodal_Product_Classification

1import warnings2from itertools import cycle3 4import matplotlib5 6# 💬 NOTE: Handle plots issues when running tests or displaying in notebooks7try:8    get_ipython  # Only exists in Jupyter9    matplotlib.use("module://matplotlib_inline.backend_inline")10except Exception:11    matplotlib.use("Agg")  # Fix error with tests12 13import matplotlib.pyplot as plt14import pandas as pd15import plotly.express as px16import seaborn as sns17from sklearn.decomposition import PCA18from sklearn.ensemble import RandomForestClassifier19from sklearn.linear_model import LogisticRegression20from sklearn.manifold import TSNE21from sklearn.metrics import (22    accuracy_score,23    auc,24    classification_report,25    confusion_matrix,26    f1_score,27    precision_score,28    recall_score,29    roc_curve,30)31 32warnings.filterwarnings("ignore")33 34 35def visualize_embeddings(36    X_train, X_test, y_train, y_test, plot_type="2D", method="PCA"37):38    """39    Visualizes high-dimensional embeddings (e.g., text or image embeddings) using dimensionality reduction techniques (PCA or t-SNE)40    and plots the results in 2D or 3D using Plotly for interactive visualizations.41 42    Args:43        X_train (np.ndarray): Training data embeddings of shape (n_samples, n_features).44        X_test (np.ndarray): Test data embeddings of shape (n_samples, n_features).45        y_train (np.ndarray): True labels for the training data.46        y_test (np.ndarray): True labels for the test data.47        plot_type (str, optional): Type of plot to generate, either '2D' or '3D'. Default is '2D'.48        method (str, optional): Dimensionality reduction method to use, either 'PCA' or 't-SNE'. Default is 'PCA'.49 50    Returns:51        None52 53    Side Effects:54        - Displays an interactive 2D or 3D scatter plot of the reduced embeddings, with points colored by their class labels.55 56    Notes:57        - PCA is a linear dimensionality reduction method, while t-SNE is non-linear and captures more complex relationships.58        - Perplexity is set to 10 for t-SNE. It can be tuned if necessary for better visualization of data clusters.59        - The function raises a `ValueError` if an invalid method is specified.60        - The function uses Plotly to display interactive plots.61 62    Example:63        visualize_embeddings(X_train, X_test, y_train, y_test, plot_type='3D', method='t-SNE')64 65    Visualization Details:66        - For 3D visualization, the reduced embeddings are plotted in a 3D scatter plot, with axes labeled as 'col1', 'col2', and 'col3'.67        - For 2D visualization, the embeddings are plotted in a 2D scatter plot, with axes labeled as 'col1' and 'col2'.68        - Class labels are represented by different colors in the scatter plots.69    """70    perplexity = 1071 72    if plot_type == "3D":73        if method == "PCA":74            # Create an instance of PCA for 3D visualization and fit it on the training data75            red = PCA(n_components=3)76            red.fit(X_train)77 78            # Use the trained model to transform the test data79            reduced_embeddings = red.transform(X_test)80        elif method == "t-SNE":81            # Implement t-SNE for 3D visualization82            red = TSNE(83                n_components=3, perplexity=perplexity, random_state=42, init="pca"84            )85 86            # Use the model to train and transform the test data87            reduced_embeddings = red.fit_transform(X_test)88        else:89            raise ValueError("Invalid method. Please choose either 'PCA' or 't-SNE'.")90 91        df_reduced = pd.DataFrame(reduced_embeddings, columns=["col1", "col2", "col3"])92        df_reduced["Class"] = y_test93 94        # 3D scatter plot95        fig = px.scatter_3d(96            df_reduced, x="col1", y="col2", z="col3", color="Class", title="3D"97        )98 99    else:  # 2D100        if method == "PCA":101            # Create an instance of PCA for 2D visualization and fit it on the training data102            red = PCA(n_components=2)103            red.fit(X_train)104 105            # Use the trained model to transform the test data106            reduced_embeddings = red.transform(X_test)107        elif method == "t-SNE":108            # Implement t-SNE for 2D visualization109            red = TSNE(110                n_components=2, perplexity=perplexity, random_state=42, init="pca"111            )112 113            # Use the model to train and transform the test data114            reduced_embeddings = red.fit_transform(X_test)115        else:116            raise ValueError("Invalid method. Please choose either 'PCA' or 't-SNE'.")117 118        df_reduced = pd.DataFrame(reduced_embeddings, columns=["col1", "col2"])119        df_reduced["Class"] = y_test120 121        # 2D scatter plot122        fig = px.scatter(df_reduced, x="col1", y="col2", color="Class", title="2D")123 124    fig.update_layout(125        title=f"Embeddings - {method} {plot_type} Visualization", scene=dict()126    )127 128    fig.show()129 130    return red131 132 133def test_model(X_test, y_test, model):134    """135    Evaluates a trained model on a test set by computing key performance metrics and visualizing the results.136 137    The function generates a confusion matrix, plots ROC curves (for binary or multi-class classification),138    and prints the classification report. It also computes overall accuracy, weighted precision, weighted recall,139    and weighted F1-score for the test data.140 141    Args:142        X_test (np.ndarray): Test set feature data.143        y_test (np.ndarray): True labels for the test set.144        model (sklearn-like model): A trained machine learning model with `predict` and `predict_proba` methods.145 146    Returns:147        tuple:148            - accuracy (float): Overall accuracy of the model on the test set.149            - precision (float): Weighted precision score across all classes.150            - recall (float): Weighted recall score across all classes.151            - f1 (float): Weighted F1-score across all classes.152 153    Side Effects:154        - Displays a confusion matrix as a heatmap.155        - Plots ROC curves for binary or multi-class classification.156        - Prints the classification report with precision, recall, F1-score, and support for each class.157 158    Example:159        accuracy, precision, recall, f1 = test_model(X_test, y_test, trained_model)160 161    Notes:162        - If `y_test` is multi-dimensional (e.g., one-hot encoded), it will be squeezed to 1D.163        - For binary classification, a single ROC curve is plotted. For multi-class classification,164          an ROC curve is plotted for each class with a unique color.165        - Weighted precision, recall, and F1-score are computed to handle class imbalance in multi-class classification.166 167    """168    y_pred = model.predict(X_test)169    y_pred_proba = model.predict_proba(X_test)170    y_test = y_test.squeeze() if y_test.ndim > 1 else y_test171 172    # Confusion matrix173    cm = confusion_matrix(y_test, y_pred)174    plt.figure(figsize=(10, 5))175    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues")176    plt.xlabel("Predicted")177    plt.ylabel("True")178    plt.title("Confusion Matrix")179    plt.show()180 181    # ROC curve182    fig, ax = plt.subplots(figsize=(6, 6))183 184    # Binary classification185    if y_pred_proba.shape[1] == 2:186        fpr, tpr, _ = roc_curve(y_test, y_pred_proba[:, 1])187        ax.plot(188            fpr,189            tpr,190            color="aqua",191            lw=2,192            label=f"ROC curve (area = {auc(fpr, tpr):.2f})",193        )194        ax.plot([0, 1], [0, 1], "k--", label="Chance level (AUC = 0.5)")195    # Multiclass classification196    else:197        y_onehot_test = pd.get_dummies(y_test).values198        colors = cycle(199            [200                "aqua",201                "darkorange",202                "cornflowerblue",203                "red",204                "green",205                "yellow",206                "purple",207                "pink",208                "brown",209                "black",210            ]211        )212 213        for class_id, color in zip(range(y_onehot_test.shape[1]), colors):214            fpr, tpr, _ = roc_curve(215                y_onehot_test[:, class_id], y_pred_proba[:, class_id]216            )217            ax.plot(218                fpr,219                tpr,220                color=color,221                lw=2,222                label=f"ROC curve for class {class_id} (area = {auc(fpr, tpr):.2f})",223            )224 225    ax.plot([0, 1], [0, 1], "k--", label="Chance level (AUC = 0.5)")226    ax.set_axisbelow(True)227    ax.set_xlabel("False Positive Rate")228    ax.set_ylabel("True Positive Rate")229    ax.set_title("ROC Curve")230    ax.legend(loc="lower right")231    plt.show()232 233    cr = classification_report(y_test, y_pred)234    print(cr)235 236    accuracy = accuracy_score(y_test, y_pred)237    precision = precision_score(y_test, y_pred, average="weighted")238    recall = recall_score(y_test, y_pred, average="weighted")239    f1 = f1_score(y_test, y_pred, average="weighted")240 241    return accuracy, precision, recall, f1242 243 244def train_and_evaluate_model(X_train, X_test, y_train, y_test, models=None, test=True):245    """246    Trains and evaluates multiple machine learning models on a given dataset, then visualizes the data embeddings247    using PCA before training. This function trains each model on the training data, evaluates them on the test data,248    and computes performance metrics (accuracy, precision, recall, and F1-score).249 250    Args:251        X_train (np.ndarray): Feature matrix for the training data.252        X_test (np.ndarray): Feature matrix for the test data.253        y_train (np.ndarray): True labels for the training data.254        y_test (np.ndarray): True labels for the test data.255        models (list of tuples, optional): A list of tuples, where each tuple contains the model name as a string and256                                           the corresponding scikit-learn model instance.257                                           If None, default models include Random Forest, Decision Tree, and Logistic Regression.258 259    Returns:260        list: A list of trained model tuples, where each tuple contains the model name and the trained model instance.261 262    Side Effects:263        - Displays a PCA 2D visualization of the embeddings using the `visualize_embeddings` function.264        - Trains each model on the training set.265        - Prints evaluation metrics (accuracy, precision, recall, F1-score) for each model on the test set.266        - Displays confusion matrix and ROC curve for each model using the `test_model` function.267 268    Example:269        models = train_and_evaluate_model(X_train, X_test, y_train, y_test)270 271    Notes:272        - The `models` argument can be customized to include any classification models from scikit-learn.273        - The function uses PCA for the embedding visualization. You can modify the `visualize_embeddings` function call for other visualization methods or dimensionality reduction techniques.274        - Default models include Random Forest, Decision Tree, and Logistic Regression.275    """276 277    visualize_embeddings(X_train, X_test, y_train, y_test, plot_type="2D", method="PCA")278 279    if not (models):280        # Implement the ML models281        models = [282            (283                "Random Forest",284                RandomForestClassifier(n_estimators=100, random_state=42),285            ),286            ("Logistic Regression", LogisticRegression(max_iter=1000, random_state=42)),287        ]288 289    for name, model in models:290        print("#" * 20, f" {name} ", "#" * 20)291        # Train the model on the training292        model.fit(X_train, y_train)293 294        # Evaluate the model on the test set using the test_model function295        if test:296            accuracy, precision, recall, f1 = test_model(X_test, y_test, model)297 298    return models299