from transformers.models.mistral.configuration_mistral import MistralConfig


class QConfig(MistralConfig):
    model_type = "q"

    def __init__(
        self,
        qk_norm=True,
        attention_scalar_gate=True,
        mlp_scalar_gate=True,
        gate_multiplier=2.0,
        nope_every_n=4,
        nope_layers=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.qk_norm = qk_norm
        self.attention_scalar_gate = attention_scalar_gate
        self.mlp_scalar_gate = mlp_scalar_gate
        self.gate_multiplier = gate_multiplier
        self.nope_every_n = nope_every_n
        if nope_layers is not None:
            self.nope_layers = list(nope_layers)
        elif nope_every_n:
            self.nope_layers = [
                i for i in range(self.num_hidden_layers) if (i + 1) % nope_every_n == 0
            ]
        else:
            self.nope_layers = []