from transformers import PretrainedConfig, AutoConfig

class BERTMultiAttentionConfig(PretrainedConfig):
    model_type = "bert_multi_attention"
    keys_to_ignore_at_inference = ["dropout"]

    def __init__(
        self,
        transformer="bert-base-uncased",
        hidden_size=768,
        num_heads=8,
        dropout=0.1,
        rnn_hidden_size=128,
        rnn_num_layers=2,
        rnn_bidirectional=True,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.transformer = transformer
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.dropout = dropout
        self.rnn_hidden_size = rnn_hidden_size
        self.rnn_num_layers = rnn_num_layers
        self.rnn_bidirectional = rnn_bidirectional


AutoConfig.register("bert_multi_attention", BERTMultiAttentionConfig)