models/__init__.py

# Lazy imports — modules are loaded on first access, not at package load time.
# This prevents a hard crash when xgboost's native library is momentarily absent.

__all__ = [
    "TemporalModel",
    "TGNWrapper",
    "TGATWrapper",
    "DyRepWrapper",
    "JODIEWrapper",
    "OracleMotifWrapper",
    "SequenceGRUWrapper",
    "StaticGNNWrapper",
    "XGBoostWrapper",
]


def __getattr__(name):
    if name == "TemporalModel":
        from models.base import TemporalModel
        return TemporalModel
    if name == "TGNWrapper":
        from models.tgn_wrapper import TGNWrapper
        return TGNWrapper
    if name == "TGATWrapper":
        from models.tgat import TGATWrapper
        return TGATWrapper
    if name == "DyRepWrapper":
        from models.dyrep import DyRepWrapper
        return DyRepWrapper
    if name == "JODIEWrapper":
        from models.jodie import JODIEWrapper
        return JODIEWrapper
    if name == "OracleMotifWrapper":
        from models.oracle_motif import OracleMotifWrapper
        return OracleMotifWrapper
    if name == "SequenceGRUWrapper":
        from models.sequence_gru import SequenceGRUWrapper
        return SequenceGRUWrapper
    if name == "StaticGNNWrapper":
        from models.static_gnn import StaticGNNWrapper
        return StaticGNNWrapper
    if name == "XGBoostWrapper":
        from models.xgboost_model import XGBoostWrapper
        return XGBoostWrapper
    raise AttributeError(f"module 'models' has no attribute {name!r}")