from transformers import PretrainedConfig

class BitNetConfig(PretrainedConfig):
    model_type = "bitnet"

    def __init__(self, weight_bits=1, input_bits=8, **kwargs):
        super().__init__(**kwargs)
        self.weight_bits = weight_bits
        self.input_bits = input_bits