# Copyright 2022 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0

from transformers import BertConfig as TransformersBertConfig


class BertConfig(TransformersBertConfig):

    auto_map = {
        "AutoConfig": "configuration_bert.BertConfig",
        "AutoModel": "bert_layers.BertModel",
        "AutoModelForMaskedLM": "bert_layers.BertForMaskedLM",
        "AutoModelForSequenceClassification": "bert_layers.BertForSequenceClassification",
    }

    def __init__(
        self,
        alibi_starting_size: int = 1024,
        attention_probs_dropout_prob: float = 0.0,
        **kwargs,
    ):
        super().__init__(
            attention_probs_dropout_prob=attention_probs_dropout_prob,
            **kwargs,
        )
        self.alibi_starting_size = alibi_starting_size