from transformers import PretrainedConfig

class MusicNNConfig(PretrainedConfig):
    model_type = 'musicnn'

    def __init__(
        self,
        num_classes=50,
        mid_filt=64,
        backend_units=200,
        dataset='MTT',
        **kwargs
    ):
        self.num_classes = num_classes
        self.mid_filt = mid_filt
        self.backend_units = backend_units
        self.dataset = dataset
        super().__init__(**kwargs)