from transformers import PretrainedConfig


class OnnxParameterStoreConfig(PretrainedConfig):
    model_type = "onnx_parameter_store"

    def __init__(
        self,
        parameter_spec_file: str = "onnx_parameter_map.json",
        param_prefix: str = "weights",
        source_onnx_repo_id: str | None = None,
        source_onnx_file: str = "speech_tokenizer_v3.onnx",
        audio_sample_rate: int = 16000,
        n_mels: int = 128,
        n_fft: int = 512,
        win_length: int = 400,
        hop_length: int = 160,
        **kwargs,
    ):
        self.parameter_spec_file = parameter_spec_file
        self.param_prefix = param_prefix
        self.source_onnx_repo_id = source_onnx_repo_id
        self.source_onnx_file = source_onnx_file
        self.audio_sample_rate = int(audio_sample_rate)
        self.n_mels = int(n_mels)
        self.n_fft = int(n_fft)
        self.win_length = int(win_length)
        self.hop_length = int(hop_length)
        super().__init__(**kwargs)