updated config.json

config.json

@@ -7,21 +7,23 @@
       "AutoModelForCausalLM": "modeling_step_audio_2.StepAudio2ForCausalLM"
     },
     "model_type": "step_audio_2",
-    "hidden_size": 3584,
-    "intermediate_size": 18944,
-    "num_attention_heads": 28,
-    "num_attention_groups": 4,
-    "num_key_value_heads": 4,
-    "num_hidden_layers": 28,
-    "max_seq_len": 16384,
-    "vocab_size": 158720,
-    "rms_norm_eps": 1e-06,
-    "eos_token_id": 151643,
-    "pad_token_id": 151643,
-    "rope_theta": 1000000.0,
-    "max_position_embeddings": 16384,
-    "rope_scaling": null,
-    "torch_dtype": "bfloat16",
+    "text_config": {
+      "hidden_size": 3584,
+      "intermediate_size": 18944,
+      "num_attention_heads": 28,
+      "num_attention_groups": 4,
+      "num_key_value_heads": 4,
+      "num_hidden_layers": 28,
+      "max_seq_len": 16384,
+      "vocab_size": 158720,
+      "rms_norm_eps": 1e-06,
+      "eos_token_id": 151643,
+      "pad_token_id": 151643,
+      "rope_theta": 1000000.0,
+      "max_position_embeddings": 16384,
+      "rope_scaling": null,
+      "torch_dtype": "bfloat16"
+    },
     "audio_encoder_config": {
       "n_mels": 128,
       "n_audio_ctx": 1500,

configuration_step_audio_2.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 
 from transformers import Qwen2Config
 from transformers.configuration_utils import PretrainedConfig
@@ -29,13 +29,80 @@ class StepAudio2EncoderConfig(PretrainedConfig):
         self.adapter_stride = adapter_stride
         super().__init__(**kwargs)
 
+class StepAudio2TextConfig(PretrainedConfig):
+    model_type = "step_audio_2_text"
+
+    def __init__(
+        self,
+        vocab_size=64012,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=48,
+        num_attention_heads=32,
+        num_attention_groups=4,
+        num_key_value_heads=4,
+        hidden_act="silu",
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        rope_theta=1000000.0,
+        rope_scaling=None,
+        eos_token_id=None,
+        **kwargs
+    ):
+
+        if eos_token_id is not None:
+            if isinstance(eos_token_id, list):
+                eos_token_id = list(set([151643, 151645, 151665] + eos_token_id))
+            else:
+                eos_token_id = [151643, 151645, 151665, eos_token_id]
+        else:
+            eos_token_id = [151643, 151645, 151665]
+
+        super().__init__(
+            eos_token_id=eos_token_id,
+            **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_attention_groups = num_attention_groups
+        self.num_key_value_heads = num_key_value_heads
+        assert self.num_attention_groups == self.num_key_value_heads, "num_attention_groups must be equal to num_key_value_heads"
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.text_config = Qwen2Config(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            hidden_act=hidden_act,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            architectures=["Qwen2ForCausalLM"],
+            torch_dtype=getattr(self, "torch_dtype", "bfloat16"),
+        )
+
 class StepAudio2Config(PretrainedConfig):
     model_type = "step_audio_2"
     architectures = ["StepAudio2ForCausalLM"]
 
     def __init__(
         self,
-        audio_encoder_config=None,
+        audio_encoder_config :Optional[Union[dict, StepAudio2EncoderConfig]] = None,
+        text_config: Optional[Union[dict, StepAudio2TextConfig]] = None,
         use_sliding_window: bool = False,
         sliding_window: Optional[int] = 2048,
         max_window_layers: Optional[int] = None,
@@ -48,7 +115,12 @@ class StepAudio2Config(PretrainedConfig):
         kwargs.setdefault("max_window_layers", max_window_layers)
         super().__init__(**kwargs)
 
-        self.text_config = Qwen2Config(**kwargs)
+        if text_config is None:
+            text_config = StepAudio2TextConfig().text_config
+        elif isinstance(text_config, dict):
+            text_config = StepAudio2TextConfig(**text_config).text_config
+
+        self.text_config = text_config
 
         if audio_encoder_config is None:
             self.audio_encoder_config = StepAudio2EncoderConfig()

modeling_step_audio_2.py

@@ -328,8 +328,8 @@ class StepAudio2ForCausalLM(PreTrainedModel, GenerationMixin):
             self.encoder = self.encoder.bfloat16()
             self.adapter = self.adapter.bfloat16()
         self.lm_head = torch.nn.Linear(
-            config.hidden_size,
-            config.vocab_size,
+            config.text_config.hidden_size,
+            config.text_config.vocab_size,
             bias=False,
             dtype=dtype
         )