| from transformers import PretrainedConfig | |
| class MoLMConfig(PretrainedConfig): | |
| model_type = "MoLM" | |
| def __init__( | |
| self, | |
| vocab_size=50304, | |
| n_embd=768, | |
| n_layer=12, | |
| n_head=12, | |
| sequence_length=1024, | |
| mlp_dim_exp_factor=1.0, | |
| dropout=0.0, | |
| bias=False, | |
| num_experts=6, | |
| expert_configs=None, | |
| use_router=False, | |
| top_k_experts=6, | |
| architectures=["MoLM"], | |
| auto_map={ | |
| "AutoConfig": "configuration.MoLMConfig", | |
| "AutoModelForCausalLM": "modeling.MoLM", | |
| "AutoTokenizer": "GPT2TokenizerFast" | |
| }, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.vocab_size = vocab_size | |
| self.n_embd = n_embd | |
| self.n_layer = n_layer | |
| self.n_head = n_head | |
| self.sequence_length = sequence_length | |
| self.mlp_dim_exp_factor = mlp_dim_exp_factor | |
| self.dropout = dropout | |
| self.bias = bias | |
| self.num_experts = num_experts | |
| self.expert_configs = expert_configs | |
| self.use_router = use_router | |
| self.architectures = architectures | |
| self.auto_map = auto_map | |
| self.top_k_experts = top_k_experts | |
| def to_dict(self): | |
| config_dict = super().to_dict() | |
| config_dict.update({ | |
| "vocab_size": self.vocab_size, | |
| "n_embd": self.n_embd, | |
| "n_layer": self.n_layer, | |
| "n_head": self.n_head, | |
| "sequence_length": self.sequence_length, | |
| "mlp_dim_exp_factor": self.mlp_dim_exp_factor, | |
| "dropout": self.dropout, | |
| "bias": self.bias, | |
| "num_experts": self.num_experts, | |
| "expert_configs": [ | |
| expert_config.to_dict() if not isinstance(expert_config, dict) else expert_config | |
| for expert_config in self.expert_configs | |
| ] if self.expert_configs else None, | |
| "use_router": self.use_router, | |
| "top_k_experts": self.top_k_experts, | |
| }) | |
| return config_dict |