File size: 740 Bytes
25ce5c5 |
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 29 30 31 32 |
"""MDLM config for Hugging Face.
"""
import transformers
class MDLMConfig(transformers.PretrainedConfig):
"""Hugging Face configuration class for MDLM."""
model_type = "mdlm"
def __init__(
self,
vocab_size: int = 50258,
model_length: int = 1024,
hidden_dim: int = 768,
cond_dim: int = 129,
n_blocks: int = 12,
n_heads: int = 12,
dropout: float = 0.1,
time_conditioning: bool = False,
** kwargs):
super().__init__(**kwargs)
self.vocab_size = vocab_size
self.model_length = model_length
self.hidden_dim = hidden_dim
self.cond_dim = cond_dim
self.n_blocks = n_blocks
self.n_heads = n_heads
self.dropout = dropout
self.time_conditioning = time_conditioning
|