mdlm-owt / configuration_mdlm.py
yairschiff's picture
Upload MDLM
25ce5c5 verified
raw
history blame
No virus
740 Bytes
"""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