File size: 4,554 Bytes
6aa989f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | 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"]
|