from transformers import PretrainedConfig


class BardConfig(PretrainedConfig):

    model_type = "bard"

    def __init__(
        self,
        vocab_size=16384,
        d_model=512,
        n_heads=8,
        n_layers=36,
        window=256,
        global_every=4,
        rope="local",
        parallel=True,
        context=4096,
        vars_weight=4.0,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.window = window
        self.global_every = global_every
        self.rope = rope
        self.parallel = parallel
        self.context = context
        self.vars_weight = vars_weight
        super().__init__(**kwargs)