from transformers import PretrainedConfig | |
from typing import List | |
class TestConfig(PretrainedConfig): | |
model_type = "my_test_model" | |
def __init__( | |
self, | |
input_dim: int = 20, | |
output_dim: int = 10, | |
**kwargs, | |
): | |
self.input_dim = input_dim | |
self.output_dim = output_dim | |
super().__init__(**kwargs) | |