from transformers import PretrainedConfig

class BERTClassifierConfig(PretrainedConfig):
    model_type = "bert_classifier"

    def __init__(
        self,
        drop_rate: float = 0.2,
        num_labels: int = 2,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.drop_rate = drop_rate
        self.num_labels = num_labels