added model files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,36 @@
+{
+    "architectures": [
+      "StepAudio2ForCausalLM"
+    ],
+    "auto_map": {
+      "AutoConfig": "configuration_step_audio_2.StepAudio2Config",
+      "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",
+    "audio_encoder_config": {
+      "n_mels": 128,
+      "n_audio_ctx": 1500,
+      "n_audio_state": 1280,
+      "n_audio_head": 20,
+      "n_audio_layer": 32,
+      "n_codebook_size": 4096,
+      "llm_dim": 3584,
+      "kernel_size": 3,
+      "adapter_stride": 2
+    }
+}

configuration_step_audio_2.py

@@ -0,0 +1,56 @@
+from typing import Optional
+
+from transformers import Qwen2Config
+from transformers.configuration_utils import PretrainedConfig
+
+
+class StepAudio2EncoderConfig(PretrainedConfig):
+    model_type = "step_audio_2_encoder"
+
+    def __init__(
+        self,
+        n_mels=128,
+        n_audio_ctx=1500,
+        n_audio_state=512,
+        n_audio_head=8,
+        n_audio_layer=6,
+        llm_dim=4096,
+        kernel_size=3,
+        adapter_stride=2,
+        **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.llm_dim     = llm_dim
+        self.kernel_size = kernel_size
+        self.adapter_stride = adapter_stride
+        super().__init__(**kwargs)
+
+class StepAudio2Config(PretrainedConfig):
+    model_type = "step_audio_2"
+    architectures = ["StepAudio2ForCausalLM"]
+
+    def __init__(
+        self,
+        audio_encoder_config=None,
+        use_sliding_window: bool = False,
+        sliding_window: Optional[int] = 2048,
+        max_window_layers: Optional[int] = None,
+        **kwargs
+    ):
+        kwargs.setdefault("use_sliding_window", use_sliding_window)
+        kwargs.setdefault("sliding_window", sliding_window)
+        if max_window_layers is None:
+            max_window_layers = kwargs.get("num_hidden_layers", None)
+        kwargs.setdefault("max_window_layers", max_window_layers)
+        super().__init__(**kwargs)
+
+        self.text_config = Qwen2Config(**kwargs)
+
+        if audio_encoder_config is None:
+            self.audio_encoder_config = StepAudio2EncoderConfig()
+        elif isinstance(audio_encoder_config, dict):
+            self.audio_encoder_config = StepAudio2EncoderConfig(**audio_encoder_config)

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b88e02b0b8c643412ec68cae009b3952dbd8e27642d61626065a2c420a8b73c
+size 4925370984

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d412c8d2fc17ca3351751f3171d48ff5b139af623aa05749062f132ac2585f1
+size 4932751008

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:135ae4a891350e8ebf9791ef073d310314e1f75192bece0971bfab7b86c5587c
+size 4988307424

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d35bf0ec42ff9ec160dfc6c5cb20a65247f0f8ba1c6edc620398c2ef49a66295
+size 1784019520

model.safetensors.index.json

@@ -0,0 +1,839 @@
+{
+  "metadata": {
+    "total_size": 16630358528
+  },
+  "weight_map": {
+    "adapter.conv.bias": "model-00004-of-00004.safetensors",
+    "adapter.conv.weight": "model-00004-of-00004.safetensors",
+    "adapter.linear1.bias": "model-00004-of-00004.safetensors",
+    "adapter.linear1.weight": "model-00004-of-00004.safetensors",
+    "adapter.linear2.bias": "model-00004-of-00004.safetensors",
+    "adapter.linear2.weight": "model-00004-of-00004.safetensors",
+    "encoder.after_norm.bias": "model-00004-of-00004.safetensors",
+    "encoder.after_norm.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.0.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.0.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.1.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.10.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.11.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.12.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.13.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.14.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.15.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.16.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.16.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.16.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.16.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.16.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.16.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.17.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.18.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.19.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.2.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.2.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.20.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.20.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.21.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.22.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.23.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.24.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.25.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.26.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.27.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.28.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.29.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.3.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.3.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.30.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.30.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.key.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.out.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.out.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.query.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.query.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.value.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn.value.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.attn_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp.0.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp.0.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp.2.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp.2.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp_ln.bias": "model-00004-of-00004.safetensors",
+    "encoder.blocks.31.mlp_ln.weight": "model-00004-of-00004.safetensors",
+    "encoder.blocks.4.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.4.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.5.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.6.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.7.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.8.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.key.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.out.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.out.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.query.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.query.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.value.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn.value.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.attn_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp.0.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp.0.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp.2.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp.2.weight": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp_ln.bias": "model-00003-of-00004.safetensors",
+    "encoder.blocks.9.mlp_ln.weight": "model-00003-of-00004.safetensors",
+    "encoder.conv1.bias": "model-00003-of-00004.safetensors",
+    "encoder.conv1.weight": "model-00003-of-00004.safetensors",
+    "encoder.conv2.bias": "model-00003-of-00004.safetensors",
+    "encoder.conv2.weight": "model-00003-of-00004.safetensors",
+    "encoder.positional_embedding.weight": "model-00003-of-00004.safetensors",
+    "lm_head.weight": "model-00004-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.norm.weight": "model-00003-of-00004.safetensors"
+  }
+}

modeling_step_audio_2.py

@@ -0,0 +1,425 @@
+from typing import Iterable, Optional, Tuple
+
+import librosa
+import torch
+import torch.nn.functional as F
+import torchaudio
+from torch import Tensor, nn
+from transformers import PreTrainedModel, Qwen2Model
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_step_audio_2 import StepAudio2Config
+
+
+def _mel_filters(n_mels: int) -> torch.Tensor:
+    """Load the mel filterbank matrix for projecting STFT into a Mel spectrogram."""
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    if n_mels == 128:
+        return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128))
+    else:
+        return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=80))
+
+
+def load_audio(file_path, target_rate=16000, max_length=None):
+    """
+    Open an audio file and read as mono waveform, resampling as necessary
+    If max_length is provided, truncate the audio to that length
+    """
+    waveform, sample_rate = torchaudio.load(file_path)
+    if sample_rate != target_rate:
+        waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_rate)(waveform)
+    audio = waveform[0]  # get the first channel
+
+    # Truncate audio if it exceeds max_length
+    if max_length is not None and audio.shape[0] > max_length:
+        audio = audio[:max_length]
+
+    return audio
+
+def log_mel_spectrogram(audio, n_mels=128, padding=479, device=None):
+    """
+    Compute the log-Mel spectrogram with specific padding for StepAudio
+    """
+    if not torch.is_tensor(audio):
+        if isinstance(audio, str):
+            audio = load_audio(audio)
+        audio = torch.from_numpy(audio)
+    if device is not None:
+        audio = audio.to(device)
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(400).to(audio.device)
+    stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+    filters = _mel_filters(n_mels)
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+def compute_token_num(max_feature_len):
+    # First, audio goes through encoder:
+    # 1. conv1: kernel=3, stride=1, padding=1 -> size unchanged
+    # 2. conv2: kernel=3, stride=2, padding=1 -> size/2
+    # 3. avg_pooler: kernel=2, stride=2 -> size/2
+    max_feature_len = max_feature_len - 2  # remove padding
+    encoder_output_dim = (max_feature_len + 1) // 2 // 2  # after conv2 and avg_pooler
+    
+    # Then through adaptor (parameters from config file):
+    padding = 1
+    kernel_size = 3  # from config: audio_encoder_config.kernel_size
+    stride = 2      # from config: audio_encoder_config.adapter_stride
+    adapter_output_dim = (encoder_output_dim + 2 * padding - kernel_size) // stride + 1
+    return adapter_output_dim
+
+def make_non_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Make mask tensor containing indices of non-padded part.
+
+    The sequences in a batch may have different lengths. To enable
+    batch computing, padding is need to make all sequence in same
+    size. To avoid the padding part pass value to context dependent
+    block such as attention or convolution , this padding part is
+    masked.
+
+    1 for non-padded part and 0 for padded part.
+
+    Parameters
+    ----------
+        lengths (torch.Tensor): Batch of lengths (B,).
+
+    Returns:
+    -------
+        torch.Tensor: Mask tensor containing indices of padded part (B, max_T).
+
+    Examples:
+        >>> import torch
+        >>> import s3tokenizer
+        >>> lengths = torch.tensor([5, 3, 2])
+        >>> masks = s3tokenizer.make_non_pad_mask(lengths)
+        masks = [[1, 1, 1, 1, 1],
+                 [1, 1, 1, 0, 0],
+                 [1, 1, 0, 0, 0]]
+    """
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0,
+                             max_len,
+                             dtype=torch.int64,
+                             device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return ~mask
+
+def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Convert bool-tensor to float-tensor for flash attention.
+
+    Parameters
+    ----------
+        lengths (torch.Tensor): Batch of lengths (B, ?).
+
+    Returns:
+    -------
+        torch.Tensor: Mask tensor containing indices of padded part (B, ?).
+
+    Examples:
+        >>> import torch
+        >>> import s3tokenizer
+        >>> lengths = torch.tensor([5, 3, 2])
+        >>> masks = s3tokenizer.make_non_pad_mask(lengths)
+        masks = [[1, 1, 1, 1, 1],
+                 [1, 1, 1, 0, 0],
+                 [1, 1, 0, 0, 0]]
+        >>> new_masks = s3tokenizer.mask_to_bias(masks, torch.float32)
+        new_masks = [[-0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00],
+                    [-0.0000e+00, -0.0000e+00, -0.0000e+00, -1.0000e+10, -1.0000e+10],
+                    [-0.0000e+00, -0.0000e+00, -1.0000e+10, -1.0000e+10, -1.0000e+10]]
+    """
+    assert mask.dtype == torch.bool
+    assert dtype in [torch.float32, torch.bfloat16, torch.float16]
+    mask = mask.to(dtype)
+    # attention mask bias
+    # NOTE(Mddct): torch.finfo jit issues
+    #     chunk_masks = (1.0 - chunk_masks) * torch.finfo(dtype).min
+    mask = (1.0 - mask) * -1.0e+10
+    return mask
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, input: Tensor) -> Tensor:
+        return super().forward(input).type(input.dtype)
+
+class Linear(nn.Linear):
+    def forward(self, input: Tensor) -> Tensor:
+        return F.linear(
+            input,
+            self.weight.to(input.dtype),
+            None if self.bias is None else self.bias.to(input.dtype),
+        )
+
+class Conv1d(nn.Conv1d):
+    def _conv_forward(
+        self, input: Tensor, weight: Tensor, bias: Optional[Tensor]
+    ) -> Tensor:
+        return super()._conv_forward(
+            input, weight.to(input.dtype), None if bias is None else bias.to(input.dtype)
+        )
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+        self.n_head = n_head
+        self.query = Linear(n_state, n_state)
+        self.key = Linear(n_state, n_state, bias=False)
+        self.value = Linear(n_state, n_state)
+        self.out = Linear(n_state, n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        mask: Optional[Tensor] = None,
+    ):
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+
+        wv, qk = self.qkv_attention(q, k, v, mask)
+        return self.out(wv), qk
+
+    def qkv_attention(
+        self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None
+    ):
+        _, T, D = q.shape
+        scale = (D // self.n_head) ** -0.25
+        q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
+        k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
+        v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
+
+        qk = q @ k  # (B, n_head, T, T)
+        if mask is not None:
+            qk = qk + mask
+        qk = qk.float()
+
+        w = F.softmax(qk, dim=-1).to(q.dtype)
+        return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+
+        self.attn = MultiHeadAttention(n_state, n_head)
+        self.attn_ln = LayerNorm(n_state)
+
+        n_mlp = n_state * 4
+        self.mlp = nn.Sequential(
+            Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)
+        )
+        self.mlp_ln = LayerNorm(n_state)
+
+    def forward(
+        self,
+        x: Tensor,
+        mask: Optional[Tensor] = None,
+    ):
+        x = x + self.attn(self.attn_ln(x.contiguous()), mask=mask)[0]
+        x = x + self.mlp(self.mlp_ln(x.contiguous()))
+        return x
+
+class AudioEncoder(nn.Module):
+    def __init__(
+        self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int
+    ):
+        super().__init__()
+        self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
+        self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
+        self.positional_embedding = nn.Embedding(n_ctx, n_state)
+        self.positional_embedding.requires_grad_(False)
+        self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
+            [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
+        )
+        self.avg_pooler = nn.AvgPool1d(2, stride=2)
+        self.after_norm = LayerNorm(n_state)
+        self.gradient_checkpointing = False
+
+    def forward(self, x: Tensor, x_len: Tensor) -> Tuple[Tensor, Tensor]:
+        T = x.size(-1)
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)  # (B, T // 2, n_state)
+        mask = make_non_pad_mask(x_len, T).unsqueeze(1)  # (B, 1, T)
+        mask = mask_to_bias(mask[:, :, (T + 1) % 2::2], x.dtype)  # (B, 1, T // 2)
+        x = (x + self.positional_embedding.weight[:x.shape[1], :]).to(x.dtype)
+        for block in self.blocks:
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(block, x, mask.unsqueeze(1))
+            else:
+                x = block(x, mask.unsqueeze(1))
+        x = x.permute(0, 2, 1)
+        x = self.avg_pooler(x)
+        x = x.permute(0, 2, 1)
+        x_len = (x_len + 1) // 2 // 2
+        x = self.after_norm(x.contiguous())
+        return x, x_len
+
+class Adaptor(nn.Module):
+    def __init__(
+        self,
+        n_state: int = 1280,
+        n_hidden: int = 3072,
+        kernel_size: int = 7,
+        stride: int = 4
+    ):
+        super().__init__()
+        self.stride = stride
+        if self.stride != -1:
+            # print("self.stride: {}".format(self.stride))
+            self.conv = Conv1d(n_state, n_state, kernel_size, stride, padding=1)
+        self.linear1 = nn.Linear(n_state, 2048)
+        self.relu = nn.ReLU()
+        self.linear2 = nn.Linear(2048, n_hidden)
+        self.gradient_checkpointing = False
+
+    def forward(self, x: Tensor) -> Tuple[Tensor]:
+        T = x.size(-1)
+        if self.stride != -1:
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(self.conv, x.permute(0, 2, 1))
+                x = x.permute(0, 2, 1)
+            else:
+                x = x.permute(0, 2, 1)
+                x = F.gelu(self.conv(x))
+                x = x.permute(0, 2, 1)
+        if self.gradient_checkpointing and self.training:
+            x = torch.utils.checkpoint.checkpoint(self.linear1, x)
+            x = torch.utils.checkpoint.checkpoint(self.relu, x)
+            x = torch.utils.checkpoint.checkpoint(self.linear2, x)
+        else:
+            x = self.linear1(x)
+            x = self.relu(x)
+            x = self.linear2(x)
+        return x
+
+class StepAudio2ForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = StepAudio2Config
+    main_input_name = "input_ids"
+    # Important: Add this attribute to make HF recognize it as a model with generation capability
+    # _keys_to_ignore_on_load_missing = ["lm_head.weight"]
+    supports_gradient_checkpointing = True  # 新增，声明支持gradient checkpointing
+
+    def __init__(self, config: StepAudio2Config):
+        super().__init__(config)
+        if isinstance(config.torch_dtype, str):
+            dtype = getattr(torch, config.torch_dtype)
+        else:
+            dtype = config.torch_dtype
+        self.model = Qwen2Model(config.text_config)
+        self.bf16 = dtype==torch.bfloat16
+        self.encoder = AudioEncoder(
+            config.audio_encoder_config.n_mels, config.audio_encoder_config.n_audio_ctx, config.audio_encoder_config.n_audio_state,
+            config.audio_encoder_config.n_audio_head, config.audio_encoder_config.n_audio_layer
+        )
+        self.adapter = Adaptor(
+            config.audio_encoder_config.n_audio_state, config.audio_encoder_config.llm_dim,
+            config.audio_encoder_config.kernel_size, config.audio_encoder_config.adapter_stride
+        )
+        if self.bf16:
+            self.encoder = self.encoder.bfloat16()
+            self.adapter = self.adapter.bfloat16()
+        self.lm_head = torch.nn.Linear(
+            config.hidden_size,
+            config.vocab_size,
+            bias=False,
+            dtype=dtype
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        wavs=None,
+        wav_lens=None,
+        attention_mask=None,
+        **kwargs
+    ):
+        hidden_states = self.model.embed_tokens(input_ids)
+        if wavs is not None:
+            if self.bf16:
+                wavs = wavs.bfloat16()
+            out, feat_lens = self.encoder(wavs, wav_lens)
+            out = self.adapter(out)
+            feat_lens = (feat_lens - 1) // 2 + 1
+            insert_location = torch.nonzero(input_ids == 151688)
+            insert_location[:,1] += 1
+            for idx in range(len(insert_location)):
+                i,s = insert_location[idx]
+                hidden_states[i][s : s+feat_lens[idx]] = out[idx][:feat_lens[idx]]
+
+        x = self.model(inputs_embeds=hidden_states, attention_mask=attention_mask)[0]
+        logits = self.lm_head(x)
+        return CausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None
+        )
+
+    def get_input_embeddings(self):
+        """Return the model's input embeddings - required for GenerationMixin"""
+        return self.model.embed_tokens
+
+    def get_output_embeddings(self):
+        """Return the model's output embeddings (LM head) - required for GenerationMixin"""
+        return self.lm_head
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        """Prepare inputs for generation - required for GenerationMixin"""
+        # Keep the wavs and wav_lens from the initial call
+        wavs = kwargs.get("wavs", None)
+        wav_lens = kwargs.get("wav_lens", None)
+
+        # For generation steps after the first, we don't need to process audio again
+        # because the audio tokens have already been replaced in the input sequence
+        if "past_key_values" in kwargs and kwargs["past_key_values"] is not None:
+            # We're in a generation step, no need to process audio again
+            return {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "past_key_values": kwargs.get("past_key_values")
+            }
+
+        # First generation step, include audio processing
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "wavs": wavs,
+            "wav_lens": wav_lens
+        }
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        """Reorder the cache for beam search - required for GenerationMixin if using beam search"""
+        # If you're not using past_key_values or beam search, this can be a simple pass-through
+        # Otherwise implement according to your model's cache structure
+        return past_key_values
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        # For Qwen2Model
+        if hasattr(self.model, 'gradient_checkpointing'):
+            self.model.gradient_checkpointing = value
+
+            # Add the missing _gradient_checkpointing_func method to Qwen2Model
+            # This is what Qwen2Model tries to use when gradient_checkpointing=True
+            if value and not hasattr(self.model, '_gradient_checkpointing_func'):
+                def _gradient_checkpointing_func(module_to_run, *args, **kwargs):
+                    # This function wraps torch.utils.checkpoint.checkpoint
+                    # and is used by Qwen2Model to perform checkpointing
+                    return torch.utils.checkpoint.checkpoint(module_to_run, *args, **kwargs)
+
+                self.model._gradient_checkpointing_func = _gradient_checkpointing_func
+
+        # For custom encoder and adapter
+        if hasattr(self.encoder, 'gradient_checkpointing'):
+            self.encoder.gradient_checkpointing = value
+        if hasattr(self.adapter, 'gradient_checkpointing'):
+            self.adapter.gradient_checkpointing = value

special_tokens_map.json

@@ -0,0 +1,43 @@
+{
+  "additional_special_tokens": [
+    "<|EOT|>",
+    "<|BOT|>",
+    "<|CALL_START|>",
+    "<|CALL_END|>",
+    "<|THINK_START|>",
+    "<|THINK_END|>",
+    "<|IMG_START|>",
+    "<|IMG_END|>",
+    "<|META_START|>",
+    "<|META_END|>",
+    "<im_patch>",
+    "<im_start>",
+    "<im_end>",
+    "<dream>",
+    "<dream_start>",
+    "<dream_end>",
+    "<|MASK_1e69f|>",
+    "<|UNMASK_1e69f|>",
+    "<video_start>",
+    "<video_end>",
+    "<patch_start>",
+    "<patch_end>",
+    "<patch_newline>",
+    "<audio_start>",
+    "<audio_end>",
+    "<audio_patch>",
+    "<audio_patch_pad>",
+    "<|SC|>",
+    "<tts_start>",
+    "<tts_end>",
+    "<tts_pad>"
+  ],
+  "eos_token": "<|EOT|>",
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

token2wav/campplus.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6ac6a63997761ae2997373e2ee1c47040854b4b759ea41ec48e4e42df0f4d73
+size 28303423

token2wav/flow.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15ccff24256ff61537c7f8b51e025116b83405f3fb017b54b008fc97da115446
+size 623466603

token2wav/flow.yaml

@@ -0,0 +1,34 @@
+flow: !new:cosyvoice2.flow.flow.CausalMaskedDiffWithXvec
+    input_size: 512
+    output_size: 80
+    spk_embed_dim: 192
+    output_type: 'mel'
+    vocab_size: 6561
+    encoder: !new:cosyvoice2.transformer.upsample_encoder_v2.UpsampleConformerEncoderV2
+        input_size: 512
+        output_size: 512
+        input_layer: 'linear'
+        pre_lookahead_len: 3
+        num_blocks: 6
+        num_up_blocks: 4
+        up_stride: 2
+        up_scale_factor: 2
+        attention_heads: 8
+        pos_enc_layer_type: 'rel_pos_espnet'
+        selfattention_layer_type: 'rel_selfattn'
+        key_bias: true
+        linear_units: 2048
+        dropout_rate: 0.1
+        positional_dropout_rate: 0.1
+        attention_dropout_rate: 0.1
+        normalize_before: True
+    decoder: !new:cosyvoice2.flow.flow_matching.CausalConditionalCFM
+        inference_cfg_rate: 0.7
+        estimator: !new:cosyvoice2.flow.decoder_dit.DiT
+            in_channels: 320
+            out_channels: 80
+            mlp_ratio: 4.0
+            depth: 16
+            num_heads: 8
+            head_dim: 64
+            hidden_size: 512

token2wav/hift.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3386cc880324d4e98e05987b99107f49e40ed925b8ecc87c1f4939432d429879
+size 83390254

token2wav/speech_tokenizer_v2_25hz.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d43342aa12163a80bf07bffb94c9de2e120a8df2f9917cd2f642e7f4219c6f71
+size 496082973

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:529f599b059f0e73c30a25205b906a84add3988c964de0b7db090f34d71e2e6e
+size 12684784