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)