from ml_lib.utils.config import Config
from models.conv_classifier import ConvClassifier
from models.bandwise_conv_classifier import BandwiseConvClassifier
from models.bandwise_conv_multihead_classifier import BandwiseConvMultiheadClassifier
from models.ensemble import Ensemble


class MConfig(Config):
    # TODO: There should be a way to automate this.

    @property
    def _model_map(self):
        return dict(ConvClassifier=ConvClassifier,
                    BandwiseConvClassifier=BandwiseConvClassifier,
                    BandwiseConvMultiheadClassifier=BandwiseConvMultiheadClassifier,
                    Ensemble=Ensemble,
                    )