from transformers import PretrainedConfig


class CageExtractorConfig(PretrainedConfig):
    model_type = "cage_extractor"

    def __init__(
        self,
        watermark_bits: int = 100,
        in_channels: int = 1,
        fine_kernel: int = 3,
        mid_kernel: int = 5,
        coarse_kernel: int = 7,
        sample_rate: int = 24000,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.watermark_bits = watermark_bits
        self.in_channels = in_channels
        self.fine_kernel = fine_kernel
        self.mid_kernel = mid_kernel
        self.coarse_kernel = coarse_kernel
        self.sample_rate = sample_rate