File size: 801 Bytes
966807a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | # 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
|