from transformers import FNetConfig


class SpliNetConfig(FNetConfig):
    model_type = "splinet"

    def __init__(
        self,
        splinet_num_heads=12,
        splinet_radius=16,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.splinet_num_heads = int(
            splinet_num_heads
        )

        self.splinet_radius = int(
            splinet_radius
        )

        self.splinet_order = 2
        self.splinet_sidedness = "single"