File size: 690 Bytes
257d747 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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) |