models/registry.py

from sklearn.linear_model import LinearRegression, LogisticRegression, Perceptron
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.naive_bayes import GaussianNB
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.svm import SVC, SVR
from sklearn.neural_network import MLPClassifier, MLPRegressor
from sklearn.ensemble import (
    RandomForestClassifier, RandomForestRegressor,
    ExtraTreesClassifier, ExtraTreesRegressor,
    AdaBoostClassifier, AdaBoostRegressor,
    GradientBoostingClassifier, GradientBoostingRegressor,
    StackingClassifier, StackingRegressor
)

REGRESSION_MODELS = {
    "Linear Regression": LinearRegression(),
    "KNN Regressor": KNeighborsRegressor(),
    "Decision Tree Regressor": DecisionTreeRegressor(),
    "SVR": SVR(),
    "MLP Regressor": MLPRegressor(max_iter=1000),
}

CLASSIFICATION_MODELS = {
    "Logistic Regression": LogisticRegression(max_iter=500),
    "KNN Classifier": KNeighborsClassifier(),
    "Naive Bayes": GaussianNB(),
    "Perceptron": Perceptron(),
    "Decision Tree Classifier": DecisionTreeClassifier(),
    "SVM Classifier": SVC(probability=True),
    "MLP Classifier": MLPClassifier(max_iter=1000),
}

MODEL_GROUPS = {
    "Basic": {
        "Regression": REGRESSION_MODELS,
        "Classification": CLASSIFICATION_MODELS,
    },
    "Bagging": {
        "Regression": {
            "Random Forest Regressor": RandomForestRegressor(),
            "Extra Trees Regressor": ExtraTreesRegressor(),
        },
        "Classification": {
            "Random Forest Classifier": RandomForestClassifier(),
            "Extra Trees Classifier": ExtraTreesClassifier(),
        },
    },
    "Boosting": {
        "Regression": {
            "AdaBoost Regressor": AdaBoostRegressor(),
            "Gradient Boosting Regressor": GradientBoostingRegressor(),
        },
        "Classification": {
            "AdaBoost Classifier": AdaBoostClassifier(),
            "Gradient Boosting Classifier": GradientBoostingClassifier(),
        },
    },
    "Stacking": {
        "Regression": {
            "Stacking Regressor": StackingRegressor(
                estimators=[("lr", LinearRegression())]
            ),
        },
        "Classification": {
            "Stacking Classifier": StackingClassifier(
                estimators=[("lr", LogisticRegression(max_iter=500))]
            ),
        },
    },
}

CLASSIFICATION_GRAPHS = [
    "Confusion Matrix",
    "ROC Curve",
    "Per-Class Metrics Table",
    "Precision-Recall Curve",
    "Probability Histogram",
]

REGRESSION_GRAPHS = [
    "Actual vs Predicted",
    "Residual Plot",
    "Residual Histogram",
    "Feature Importance",
    "Learning Curve",
]