from transformers import PretrainedConfig | |
class MedTSConfig(PretrainedConfig): | |
model_type = "MedTS" | |
def __init__( | |
self, | |
vocab_size = 4, | |
n_embd = 128, | |
block_size=32, | |
n_layer=2, | |
n_head=2, | |
dropout=0.1, | |
**kwargs, | |
): | |
self.vocab_size = vocab_size | |
self.n_embd = n_embd | |
self.block_size = block_size | |
self.n_layer = n_layer | |
self.n_head = n_head | |
self.dropout = dropout | |
super().__init__(**kwargs) | |