# coding=utf-8
# Copyright 2024 NAVER Cloud Corp. and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""CosyVoice2 audio encoder configuration"""

from transformers import AutoConfig
from transformers.configuration_utils import PretrainedConfig


class CosyVoice2Config(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a CosyVoice2 audio encoder model.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.

    Args:
        n_mels (`int`, *optional*, defaults to 128):
            Number of mel frequency bins for the input spectrogram.
        n_audio_ctx (`int`, *optional*, defaults to 1500):
            Maximum audio context length.
        n_audio_state (`int`, *optional*, defaults to 1280):
            Dimensionality of the encoder hidden states.
        n_audio_head (`int`, *optional*, defaults to 20):
            Number of attention heads in the encoder.
        n_audio_layer (`int`, *optional*, defaults to 6):
            Number of encoder layers.
        n_codebook_size (`int`, *optional*, defaults to 6561):
            Size of the codebook (3^8).
        use_sdpa (`bool`, *optional*, defaults to `True`):
            Whether to use Scaled Dot-Product Attention.
    """

    model_type = "cosyvoice2"

    def __init__(
        self,
        n_mels: int = 128,
        n_audio_ctx: int = 1500,
        n_audio_state: int = 1280,
        n_audio_head: int = 20,
        n_audio_layer: int = 6,
        n_codebook_size: int = 6561,  # 3**8
        use_sdpa: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.n_mels = n_mels
        self.n_audio_ctx = n_audio_ctx
        self.n_audio_state = n_audio_state
        self.n_audio_head = n_audio_head
        self.n_audio_layer = n_audio_layer
        self.n_codebook_size = n_codebook_size
        self.use_sdpa = use_sdpa


AutoConfig.register("cosyvoice2", CosyVoice2Config)

__all__ = ["CosyVoice2Config"]