File size: 348 Bytes
deb51b4
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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