brain_gcn/models/__init__.py

from .brain_gcn import (
    BrainGCNClassifier,
    ConnectivityMLPClassifier,
    GraphOnlyClassifier,
    TemporalGRUClassifier,
    build_model,
)
from .advanced_models import (
    GATClassifier,
    TransformerClassifier,
    CNN3DClassifier,
    GraphSAGEClassifier,
)
from .registry import ModelRegistry, ModelConfig, add_model_choice_argument

__all__ = [
    # Original models
    "BrainGCNClassifier",
    "ConnectivityMLPClassifier",
    "GraphOnlyClassifier",
    "TemporalGRUClassifier",
    # Advanced models
    "GATClassifier",
    "TransformerClassifier",
    "CNN3DClassifier",
    "GraphSAGEClassifier",
    # Utilities
    "build_model",
    "ModelRegistry",
    "ModelConfig",
    "add_model_choice_argument",
]