Chemfuser / configuration_diffusion.py
kelu01's picture
Update configuration_diffusion.py
897c3c0 verified
Raw History Blame Contribute Delete
637 Bytes
from transformers import PretrainedConfig
class DiffusionConfig(PretrainedConfig):
model_type = "gpt2"
def __init__(
self,
vocab_size=425,
d_model=512,
n_heads=16,
n_layers=6,
dropout=0.1,
max_len=2250,
num_steps=1000,
pad_token=0,
**kwargs
):
super().__init__(**kwargs)
self.vocab_size = vocab_size
self.d_model = d_model
self.n_heads = n_heads
self.n_layers = n_layers
self.dropout = dropout
self.max_len = max_len
self.num_steps = num_steps
self.pad_token = pad_token