from __future__ import annotations from transformers import PretrainedConfig class ArkttsConfig(PretrainedConfig): model_type = "arktts" def __init__( self, vocab_size: int = 155776, dim: int = 896, n_layer: int = 24, n_head: int = 14, n_local_heads: int = 2, head_dim: int = 64, intermediate_size: int = 4864, max_seq_len: int = 2048, rope_base: float = 1_000_000, norm_eps: float = 1e-6, dropout: float = 0.0, attention_qkv_bias: bool = True, attention_qk_norm: bool = False, attention_o_bias: bool = False, tie_word_embeddings: bool = True, codebook_size: int = 4096, num_codebooks: int = 10, semantic_begin_id: int = 151678, semantic_end_id: int = 155773, n_fast_layer: int = 4, fast_dim: int = 896, fast_n_head: int = 14, fast_n_local_heads: int = 2, fast_head_dim: int = 64, fast_intermediate_size: int = 4864, fast_attention_qkv_bias: bool = False, fast_attention_qk_norm: bool = False, fast_attention_o_bias: bool = False, norm_fastlayer_input: bool = True, initializer_range: float = 0.02, use_gradient_checkpointing: bool = False, codec_filename: str = "codec.pth", codec_sample_rate: int = 44100, codec_frame_size: int = 2048, codec_post_n_layer: int = 8, codec_post_n_head: int = 16, codec_post_n_local_heads: int = 8, codec_post_intermediate_size: int = 1216, ras_window_size: int = 10, ras_temperature: float = 1.0, ras_top_p: float = 0.9, eos_token_id: int = 151645, pad_token_id: int = 151643, **kwargs, ): kwargs.pop("model_type", None) kwargs.pop("audio_embed_dim", None) kwargs.pop("is_reward_model", None) super().__init__( eos_token_id=eos_token_id, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) self.vocab_size = int(vocab_size) self.dim = int(dim) self.hidden_size = self.dim self.n_layer = int(n_layer) self.num_hidden_layers = self.n_layer self.n_head = int(n_head) self.num_attention_heads = self.n_head self.n_local_heads = int(n_local_heads) self.num_key_value_heads = self.n_local_heads self.head_dim = int(head_dim) self.intermediate_size = int(intermediate_size) self.max_seq_len = int(max_seq_len) self.max_position_embeddings = self.max_seq_len self.rope_base = float(rope_base) self.rope_theta = self.rope_base self.norm_eps = float(norm_eps) self.rms_norm_eps = self.norm_eps self.dropout = float(dropout) self.attention_qkv_bias = bool(attention_qkv_bias) self.attention_qk_norm = bool(attention_qk_norm) self.attention_o_bias = bool(attention_o_bias) self.codebook_size = int(codebook_size) self.num_codebooks = int(num_codebooks) self.semantic_begin_id = int(semantic_begin_id) self.semantic_end_id = int(semantic_end_id) self.n_fast_layer = int(n_fast_layer) self.fast_dim = int(fast_dim) self.fast_n_head = int(fast_n_head) self.fast_n_local_heads = int(fast_n_local_heads) self.fast_head_dim = int(fast_head_dim) self.fast_intermediate_size = int(fast_intermediate_size) self.fast_attention_qkv_bias = bool(fast_attention_qkv_bias) self.fast_attention_qk_norm = bool(fast_attention_qk_norm) self.fast_attention_o_bias = bool(fast_attention_o_bias) self.norm_fastlayer_input = bool(norm_fastlayer_input) self.initializer_range = float(initializer_range) self.use_gradient_checkpointing = bool(use_gradient_checkpointing) self.codec_filename = str(codec_filename) self.codec_sample_rate = int(codec_sample_rate) self.codec_frame_size = int(codec_frame_size) self.codec_post_n_layer = int(codec_post_n_layer) self.codec_post_n_head = int(codec_post_n_head) self.codec_post_n_local_heads = int(codec_post_n_local_heads) self.codec_post_intermediate_size = int(codec_post_intermediate_size) self.ras_window_size = int(ras_window_size) self.ras_temperature = float(ras_temperature) self.ras_top_p = float(ras_top_p) __all__ = ["ArkttsConfig"]