Hum-Works/lodestone-base-4096-v1

1# Copyright 2022 MosaicML Examples authors2# SPDX-License-Identifier: Apache-2.03 4from transformers import BertConfig as TransformersBertConfig5 6 7class BertConfig(TransformersBertConfig):8 9    def __init__(10        self,11        alibi_starting_size: int = 512,12        attention_probs_dropout_prob: float = 0.0,13        **kwargs,14    ):15        """Configuration class for MosaicBert.16 17        Args:18            alibi_starting_size (int): Use `alibi_starting_size` to determine how large of an alibi tensor to19                create when initializing the model. You should be able to ignore this parameter in most cases.20                Defaults to 512.21            attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT 22                (otherwise, Flash Attention will be off by default). Defaults to 0.0.23        """24        super().__init__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)25        self.alibi_starting_size = alibi_starting_size  26