| from transformers import PretrainedConfig | |
| class MLPConfig(PretrainedConfig): | |
| model_type = "mlp" | |
| def __init__( | |
| self, | |
| num_hidden_layers: int = 2, | |
| input_size: int = 64, | |
| hidden_size: list[int] = [256, 256], | |
| output_size: int = 2, | |
| hidden_act: str = "relu", | |
| initializer_range: float = 0.02, | |
| **kwargs | |
| ): | |
| if len(hidden_size) != num_hidden_layers: | |
| raise ValueError("num_hidden_layers should equal to len(hidden_size)") | |
| self.num_hidden_layers = num_hidden_layers | |
| self.input_size = input_size | |
| self.hidden_size = hidden_size | |
| self.output_size = output_size | |
| self.hidden_act = hidden_act | |
| self.initializer_range = initializer_range | |
| super().__init__(**kwargs) |