| from transformers import PretrainedConfig, AutoModel | |
| # class PragFormerConfig(PretrainedConfig): | |
| # model_type = "pragformer" | |
| # def __init__( | |
| # self, | |
| # dropout=0.1, | |
| # layers=[768, 512], | |
| # **kwargs, | |
| # ): | |
| # self.dropout = dropout | |
| # self.layers = layers | |
| # super().__init__(**kwargs) | |
| class PragFormerConfig(PretrainedConfig): | |
| model_type = "pragformer" | |
| def __init__(self, bert=None, dropout=0.2, fc1=512, fc2=2, softmax_dim=1, **kwargs): | |
| self.bert = bert | |
| self.dropout = dropout | |
| self.fc1 = fc1 | |
| self.fc2 = fc2 | |
| self.softmax_dim = softmax_dim | |
| super().__init__(**kwargs) |