from transformers import PretrainedConfig
from transformers.utils import logging

logger = logging.get_logger(__name__)


class MSPFusionConfig(PretrainedConfig):
    model_type = "msp_fusion"

    def __init__(
        self,
        audio_hidden_size: int = 1024,
        visual_hidden_size: int = 1024,
        fusion_hidden_size: int = 1024,
        num_attention_heads: int = 16,
        attention_dropout: float = 0.1,
        activation_fn: str = "gelu",
        dropout: float = 0.1,
        layer_norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.audio_hidden_size = audio_hidden_size
        self.visual_hidden_size = visual_hidden_size
        self.fusion_hidden_size = fusion_hidden_size
        self.num_attention_heads = num_attention_heads
        self.attention_dropout = attention_dropout
        self.activation_fn = activation_fn
        self.dropout = dropout
        self.layer_norm_eps = layer_norm_eps
        self.initializer_range = initializer_range