Upload folder using huggingface_hub

.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

.hfd/last_download_command

@@ -0,0 +1 @@
+REPO_ID=leungtianle/aaaaaaaasikao-two2one TOOL=aria2c INCLUDE_PATTERNS= EXCLUDE_PATTERNS= DATASET=0 HF_USERNAME= HF_TOKEN= HF_TOKEN=https://hf-mirror.com REVISION=main

.hfd/repo_metadata.json

@@ -0,0 +1 @@
+{"_id":"694119d009f9267fc5b168bd","id":"leungtianle/aaaaaaaasikao-two2one","private":false,"tags":["safetensors","step_audio_2","custom_code","region:us"],"downloads":0,"likes":0,"modelId":"leungtianle/aaaaaaaasikao-two2one","author":"leungtianle","sha":"210d275a5bc5846f8dce7bd8c6e54f95c7b979b2","lastModified":"2025-12-16T09:15:01.000Z","gated":false,"disabled":false,"config":{"architectures":["StepAudio2ForCausalLM"],"auto_map":{"AutoConfig":"configuration_step_audio_2.StepAudio2Config","AutoModelForCausalLM":"modeling_step_audio_2.StepAudio2ForCausalLM"},"model_type":"step_audio_2","tokenizer_config":{"bos_token":null,"chat_template":"{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n","eos_token":"<|EOT|>","pad_token":"<|endoftext|>","unk_token":null}},"siblings":[{"rfilename":".gitattributes"},{"rfilename":"added_tokens.json"},{"rfilename":"config.json"},{"rfilename":"configuration_step_audio_2.py"},{"rfilename":"generation_config.json"},{"rfilename":"merges.txt"},{"rfilename":"model-00001-of-00004.safetensors"},{"rfilename":"model-00002-of-00004.safetensors"},{"rfilename":"model-00003-of-00004.safetensors"},{"rfilename":"model-00004-of-00004.safetensors"},{"rfilename":"model.safetensors.index.json"},{"rfilename":"special_tokens_map.json"},{"rfilename":"tokenizer.json"},{"rfilename":"tokenizer_config.json"},{"rfilename":"trainer_state.json"},{"rfilename":"training_args.bin"},{"rfilename":"vocab.json"}],"spaces":[],"createdAt":"2025-12-16T08:35:28.000Z","safetensors":{"parameters":{"BF16":8315179264},"total":8315179264},"usedStorage":16643140584}

config.json

@@ -0,0 +1,45 @@
+{
+  "_name_or_path": "/root/autodl-tmp/Step-Audio2/models/Step-Audio-2-mini",
+  "architectures": [
+    "StepAudio2ForCausalLM"
+  ],
+  "audio_encoder_config": {
+    "llm_dim": 3584,
+    "model_type": "step_audio_2_encoder",
+    "n_audio_head": 20,
+    "n_audio_layer": 32,
+    "n_audio_state": 1280,
+    "n_codebook_size": 4096
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_step_audio_2.StepAudio2Config",
+    "AutoModelForCausalLM": "modeling_step_audio_2.StepAudio2ForCausalLM"
+  },
+  "eos_token_id": [
+    151665,
+    151694
+  ],
+  "max_window_layers": null,
+  "model_type": "step_audio_2",
+  "pad_token_id": 151643,
+  "sliding_window": 2048,
+  "text_config": {
+    "_attn_implementation_autoset": true,
+    "architectures": [
+      "Qwen2ForCausalLM"
+    ],
+    "hidden_size": 3584,
+    "intermediate_size": 18944,
+    "max_position_embeddings": 16384,
+    "model_type": "qwen2",
+    "num_attention_heads": 28,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 4,
+    "rope_theta": 1000000.0,
+    "torch_dtype": "bfloat16",
+    "vocab_size": 158720
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.49.0",
+  "use_sliding_window": false
+}

configuration_step_audio_2.py

@@ -0,0 +1,128 @@
+from typing import Optional, Union
+
+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 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 :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,
+        **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)
+
+        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()
+        elif isinstance(audio_encoder_config, dict):
+            self.audio_encoder_config = StepAudio2EncoderConfig(**audio_encoder_config)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.49.0"
+}

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4407e8e7d550cb1bee0d7c6fda64975e128a0ba30b89aec3660b1bcf60a7725e
+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.text_config.hidden_size,
+            config.text_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:c23796e0498b651e92b0d514d43636d0dfd556534f8dde7b72ed0e2ff1d07744
+size 12684616

trainer_state.json

@@ -0,0 +1,1783 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0571428571428572,
+  "eval_steps": 500,
+  "global_step": 2500,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0020408163265306124,
+      "grad_norm": 1453.32421875,
+      "learning_rate": 2.0408163265306121e-07,
+      "loss": 22.518,
+      "step": 10
+    },
+    {
+      "epoch": 0.004081632653061225,
+      "grad_norm": 192.08140563964844,
+      "learning_rate": 4.0816326530612243e-07,
+      "loss": 19.6125,
+      "step": 20
+    },
+    {
+      "epoch": 0.006122448979591836,
+      "grad_norm": 325.45782470703125,
+      "learning_rate": 6.122448979591837e-07,
+      "loss": 20.4047,
+      "step": 30
+    },
+    {
+      "epoch": 0.00816326530612245,
+      "grad_norm": 1713.63623046875,
+      "learning_rate": 8.163265306122449e-07,
+      "loss": 18.383,
+      "step": 40
+    },
+    {
+      "epoch": 0.01020408163265306,
+      "grad_norm": 70.65190887451172,
+      "learning_rate": 1.0204081632653063e-06,
+      "loss": 16.599,
+      "step": 50
+    },
+    {
+      "epoch": 0.012244897959183673,
+      "grad_norm": 81.9298324584961,
+      "learning_rate": 1.2244897959183673e-06,
+      "loss": 15.8105,
+      "step": 60
+    },
+    {
+      "epoch": 0.014285714285714285,
+      "grad_norm": 71.76364135742188,
+      "learning_rate": 1.4285714285714286e-06,
+      "loss": 14.5338,
+      "step": 70
+    },
+    {
+      "epoch": 0.0163265306122449,
+      "grad_norm": 71.15852355957031,
+      "learning_rate": 1.6326530612244897e-06,
+      "loss": 13.8039,
+      "step": 80
+    },
+    {
+      "epoch": 0.018367346938775512,
+      "grad_norm": 64.56929016113281,
+      "learning_rate": 1.8367346938775512e-06,
+      "loss": 12.8447,
+      "step": 90
+    },
+    {
+      "epoch": 0.02040816326530612,
+      "grad_norm": 75.6944808959961,
+      "learning_rate": 2.0408163265306125e-06,
+      "loss": 12.8626,
+      "step": 100
+    },
+    {
+      "epoch": 0.022448979591836733,
+      "grad_norm": 63.737022399902344,
+      "learning_rate": 2.244897959183674e-06,
+      "loss": 12.3816,
+      "step": 110
+    },
+    {
+      "epoch": 0.024489795918367346,
+      "grad_norm": 47.80781936645508,
+      "learning_rate": 2.4489795918367347e-06,
+      "loss": 12.4277,
+      "step": 120
+    },
+    {
+      "epoch": 0.026530612244897958,
+      "grad_norm": 67.43174743652344,
+      "learning_rate": 2.6530612244897964e-06,
+      "loss": 12.3672,
+      "step": 130
+    },
+    {
+      "epoch": 0.02857142857142857,
+      "grad_norm": 48.110782623291016,
+      "learning_rate": 2.8571428571428573e-06,
+      "loss": 12.4754,
+      "step": 140
+    },
+    {
+      "epoch": 0.030612244897959183,
+      "grad_norm": 51.14405059814453,
+      "learning_rate": 3.0612244897959185e-06,
+      "loss": 12.5354,
+      "step": 150
+    },
+    {
+      "epoch": 0.0326530612244898,
+      "grad_norm": 43.35947799682617,
+      "learning_rate": 3.2653061224489794e-06,
+      "loss": 12.326,
+      "step": 160
+    },
+    {
+      "epoch": 0.03469387755102041,
+      "grad_norm": 58.11931610107422,
+      "learning_rate": 3.469387755102041e-06,
+      "loss": 11.3223,
+      "step": 170
+    },
+    {
+      "epoch": 0.036734693877551024,
+      "grad_norm": 44.09523391723633,
+      "learning_rate": 3.6734693877551024e-06,
+      "loss": 11.5811,
+      "step": 180
+    },
+    {
+      "epoch": 0.03877551020408163,
+      "grad_norm": 50.09272766113281,
+      "learning_rate": 3.877551020408164e-06,
+      "loss": 11.4436,
+      "step": 190
+    },
+    {
+      "epoch": 0.04081632653061224,
+      "grad_norm": 42.839256286621094,
+      "learning_rate": 4.081632653061225e-06,
+      "loss": 10.7325,
+      "step": 200
+    },
+    {
+      "epoch": 0.04285714285714286,
+      "grad_norm": 37.5905647277832,
+      "learning_rate": 4.2857142857142855e-06,
+      "loss": 11.7708,
+      "step": 210
+    },
+    {
+      "epoch": 0.044897959183673466,
+      "grad_norm": 47.18428421020508,
+      "learning_rate": 4.489795918367348e-06,
+      "loss": 11.2905,
+      "step": 220
+    },
+    {
+      "epoch": 0.04693877551020408,
+      "grad_norm": 50.881309509277344,
+      "learning_rate": 4.693877551020409e-06,
+      "loss": 10.9084,
+      "step": 230
+    },
+    {
+      "epoch": 0.04897959183673469,
+      "grad_norm": 43.27337646484375,
+      "learning_rate": 4.897959183673469e-06,
+      "loss": 11.535,
+      "step": 240
+    },
+    {
+      "epoch": 0.05102040816326531,
+      "grad_norm": 36.66025924682617,
+      "learning_rate": 5.1020408163265315e-06,
+      "loss": 11.1937,
+      "step": 250
+    },
+    {
+      "epoch": 0.053061224489795916,
+      "grad_norm": 44.64875411987305,
+      "learning_rate": 5.306122448979593e-06,
+      "loss": 11.4403,
+      "step": 260
+    },
+    {
+      "epoch": 0.05510204081632653,
+      "grad_norm": 41.23341751098633,
+      "learning_rate": 5.510204081632653e-06,
+      "loss": 10.9498,
+      "step": 270
+    },
+    {
+      "epoch": 0.05714285714285714,
+      "grad_norm": 40.457698822021484,
+      "learning_rate": 5.7142857142857145e-06,
+      "loss": 11.5234,
+      "step": 280
+    },
+    {
+      "epoch": 0.05918367346938776,
+      "grad_norm": 46.501705169677734,
+      "learning_rate": 5.918367346938776e-06,
+      "loss": 11.5757,
+      "step": 290
+    },
+    {
+      "epoch": 0.061224489795918366,
+      "grad_norm": 47.70826721191406,
+      "learning_rate": 6.122448979591837e-06,
+      "loss": 11.5367,
+      "step": 300
+    },
+    {
+      "epoch": 0.06326530612244897,
+      "grad_norm": 42.181766510009766,
+      "learning_rate": 6.326530612244899e-06,
+      "loss": 11.4378,
+      "step": 310
+    },
+    {
+      "epoch": 0.0653061224489796,
+      "grad_norm": 38.154624938964844,
+      "learning_rate": 6.530612244897959e-06,
+      "loss": 11.6781,
+      "step": 320
+    },
+    {
+      "epoch": 0.0673469387755102,
+      "grad_norm": 44.771705627441406,
+      "learning_rate": 6.734693877551021e-06,
+      "loss": 10.5691,
+      "step": 330
+    },
+    {
+      "epoch": 0.06938775510204082,
+      "grad_norm": 42.923789978027344,
+      "learning_rate": 6.938775510204082e-06,
+      "loss": 11.2529,
+      "step": 340
+    },
+    {
+      "epoch": 0.07142857142857142,
+      "grad_norm": 39.77996063232422,
+      "learning_rate": 7.1428571428571436e-06,
+      "loss": 11.1836,
+      "step": 350
+    },
+    {
+      "epoch": 0.07346938775510205,
+      "grad_norm": 44.5689811706543,
+      "learning_rate": 7.346938775510205e-06,
+      "loss": 11.1589,
+      "step": 360
+    },
+    {
+      "epoch": 0.07551020408163266,
+      "grad_norm": 39.37873458862305,
+      "learning_rate": 7.551020408163265e-06,
+      "loss": 11.0871,
+      "step": 370
+    },
+    {
+      "epoch": 0.07755102040816327,
+      "grad_norm": 43.95620346069336,
+      "learning_rate": 7.755102040816327e-06,
+      "loss": 10.9601,
+      "step": 380
+    },
+    {
+      "epoch": 0.07959183673469387,
+      "grad_norm": 39.072547912597656,
+      "learning_rate": 7.959183673469388e-06,
+      "loss": 11.2358,
+      "step": 390
+    },
+    {
+      "epoch": 0.08163265306122448,
+      "grad_norm": 40.79281234741211,
+      "learning_rate": 8.16326530612245e-06,
+      "loss": 11.195,
+      "step": 400
+    },
+    {
+      "epoch": 0.0836734693877551,
+      "grad_norm": 39.67601776123047,
+      "learning_rate": 8.36734693877551e-06,
+      "loss": 11.5086,
+      "step": 410
+    },
+    {
+      "epoch": 0.08571428571428572,
+      "grad_norm": 47.2785530090332,
+      "learning_rate": 8.571428571428571e-06,
+      "loss": 11.7385,
+      "step": 420
+    },
+    {
+      "epoch": 0.08775510204081632,
+      "grad_norm": 38.76987075805664,
+      "learning_rate": 8.775510204081633e-06,
+      "loss": 10.8836,
+      "step": 430
+    },
+    {
+      "epoch": 0.08979591836734693,
+      "grad_norm": 42.80223846435547,
+      "learning_rate": 8.979591836734695e-06,
+      "loss": 12.0561,
+      "step": 440
+    },
+    {
+      "epoch": 0.09183673469387756,
+      "grad_norm": 4399.8046875,
+      "learning_rate": 9.183673469387756e-06,
+      "loss": 11.0337,
+      "step": 450
+    },
+    {
+      "epoch": 0.09387755102040816,
+      "grad_norm": 43.93349075317383,
+      "learning_rate": 9.387755102040818e-06,
+      "loss": 11.4056,
+      "step": 460
+    },
+    {
+      "epoch": 0.09591836734693877,
+      "grad_norm": 29.04706573486328,
+      "learning_rate": 9.591836734693878e-06,
+      "loss": 11.1528,
+      "step": 470
+    },
+    {
+      "epoch": 0.09795918367346938,
+      "grad_norm": 887.9359741210938,
+      "learning_rate": 9.795918367346939e-06,
+      "loss": 11.5702,
+      "step": 480
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 38.69070053100586,
+      "learning_rate": 1e-05,
+      "loss": 11.3116,
+      "step": 490
+    },
+    {
+      "epoch": 0.10204081632653061,
+      "grad_norm": 44.86007308959961,
+      "learning_rate": 9.999873129474573e-06,
+      "loss": 10.3874,
+      "step": 500
+    },
+    {
+      "epoch": 0.10408163265306122,
+      "grad_norm": 40.4832763671875,
+      "learning_rate": 9.999492524336743e-06,
+      "loss": 10.7138,
+      "step": 510
+    },
+    {
+      "epoch": 0.10612244897959183,
+      "grad_norm": 30.557941436767578,
+      "learning_rate": 9.99885820390154e-06,
+      "loss": 10.9528,
+      "step": 520
+    },
+    {
+      "epoch": 0.10816326530612246,
+      "grad_norm": 36.02729034423828,
+      "learning_rate": 9.997970200359592e-06,
+      "loss": 11.3563,
+      "step": 530
+    },
+    {
+      "epoch": 0.11020408163265306,
+      "grad_norm": 33.15877914428711,
+      "learning_rate": 9.996828558775486e-06,
+      "loss": 11.2402,
+      "step": 540
+    },
+    {
+      "epoch": 0.11224489795918367,
+      "grad_norm": 33.34375762939453,
+      "learning_rate": 9.995433337085492e-06,
+      "loss": 11.1627,
+      "step": 550
+    },
+    {
+      "epoch": 0.11428571428571428,
+      "grad_norm": 39.86962890625,
+      "learning_rate": 9.993784606094612e-06,
+      "loss": 11.1086,
+      "step": 560
+    },
+    {
+      "epoch": 0.11632653061224489,
+      "grad_norm": 39.242279052734375,
+      "learning_rate": 9.991882449472994e-06,
+      "loss": 11.4848,
+      "step": 570
+    },
+    {
+      "epoch": 0.11836734693877551,
+      "grad_norm": 38.2926025390625,
+      "learning_rate": 9.989726963751683e-06,
+      "loss": 11.2885,
+      "step": 580
+    },
+    {
+      "epoch": 0.12040816326530612,
+      "grad_norm": 47.0013313293457,
+      "learning_rate": 9.987318258317718e-06,
+      "loss": 11.1414,
+      "step": 590
+    },
+    {
+      "epoch": 0.12244897959183673,
+      "grad_norm": 30.91645622253418,
+      "learning_rate": 9.984656455408591e-06,
+      "loss": 11.146,
+      "step": 600
+    },
+    {
+      "epoch": 0.12448979591836734,
+      "grad_norm": 37.095882415771484,
+      "learning_rate": 9.981741690106035e-06,
+      "loss": 10.8618,
+      "step": 610
+    },
+    {
+      "epoch": 0.12653061224489795,
+      "grad_norm": 31.8751277923584,
+      "learning_rate": 9.978574110329174e-06,
+      "loss": 10.9426,
+      "step": 620
+    },
+    {
+      "epoch": 0.12857142857142856,
+      "grad_norm": 35.010799407958984,
+      "learning_rate": 9.975153876827008e-06,
+      "loss": 11.5526,
+      "step": 630
+    },
+    {
+      "epoch": 0.1306122448979592,
+      "grad_norm": 42.3626594543457,
+      "learning_rate": 9.97148116317027e-06,
+      "loss": 10.4616,
+      "step": 640
+    },
+    {
+      "epoch": 0.1326530612244898,
+      "grad_norm": 40.87386703491211,
+      "learning_rate": 9.9675561557426e-06,
+      "loss": 11.0368,
+      "step": 650
+    },
+    {
+      "epoch": 0.1346938775510204,
+      "grad_norm": 38.359169006347656,
+      "learning_rate": 9.963379053731104e-06,
+      "loss": 10.5361,
+      "step": 660
+    },
+    {
+      "epoch": 0.13673469387755102,
+      "grad_norm": 35.00498962402344,
+      "learning_rate": 9.95895006911623e-06,
+      "loss": 10.6864,
+      "step": 670
+    },
+    {
+      "epoch": 0.13877551020408163,
+      "grad_norm": 34.58518600463867,
+      "learning_rate": 9.954269426661023e-06,
+      "loss": 10.2779,
+      "step": 680
+    },
+    {
+      "epoch": 0.14081632653061224,
+      "grad_norm": 34.32389831542969,
+      "learning_rate": 9.949337363899709e-06,
+      "loss": 11.1447,
+      "step": 690
+    },
+    {
+      "epoch": 0.14285714285714285,
+      "grad_norm": 37.98169708251953,
+      "learning_rate": 9.944154131125643e-06,
+      "loss": 10.9091,
+      "step": 700
+    },
+    {
+      "epoch": 0.14489795918367346,
+      "grad_norm": 39.24800491333008,
+      "learning_rate": 9.938719991378614e-06,
+      "loss": 10.98,
+      "step": 710
+    },
+    {
+      "epoch": 0.1469387755102041,
+      "grad_norm": 37.687591552734375,
+      "learning_rate": 9.933035220431489e-06,
+      "loss": 10.6501,
+      "step": 720
+    },
+    {
+      "epoch": 0.1489795918367347,
+      "grad_norm": 34.72271728515625,
+      "learning_rate": 9.927100106776213e-06,
+      "loss": 10.8848,
+      "step": 730
+    },
+    {
+      "epoch": 0.1510204081632653,
+      "grad_norm": 33.22209548950195,
+      "learning_rate": 9.920914951609189e-06,
+      "loss": 11.3228,
+      "step": 740
+    },
+    {
+      "epoch": 0.15306122448979592,
+      "grad_norm": 34.376155853271484,
+      "learning_rate": 9.914480068815964e-06,
+      "loss": 11.8528,
+      "step": 750
+    },
+    {
+      "epoch": 0.15510204081632653,
+      "grad_norm": 36.82484817504883,
+      "learning_rate": 9.907795784955327e-06,
+      "loss": 11.0567,
+      "step": 760
+    },
+    {
+      "epoch": 0.15714285714285714,
+      "grad_norm": 35.90745162963867,
+      "learning_rate": 9.900862439242719e-06,
+      "loss": 10.6546,
+      "step": 770
+    },
+    {
+      "epoch": 0.15918367346938775,
+      "grad_norm": 34.09635543823242,
+      "learning_rate": 9.893680383533027e-06,
+      "loss": 10.6432,
+      "step": 780
+    },
+    {
+      "epoch": 0.16122448979591836,
+      "grad_norm": 39.50465774536133,
+      "learning_rate": 9.88624998230272e-06,
+      "loss": 10.471,
+      "step": 790
+    },
+    {
+      "epoch": 0.16326530612244897,
+      "grad_norm": 36.799888610839844,
+      "learning_rate": 9.878571612631364e-06,
+      "loss": 10.9606,
+      "step": 800
+    },
+    {
+      "epoch": 0.1653061224489796,
+      "grad_norm": 37.335994720458984,
+      "learning_rate": 9.870645664182478e-06,
+      "loss": 10.9557,
+      "step": 810
+    },
+    {
+      "epoch": 0.1673469387755102,
+      "grad_norm": 40.67239761352539,
+      "learning_rate": 9.862472539183757e-06,
+      "loss": 11.13,
+      "step": 820
+    },
+    {
+      "epoch": 0.16938775510204082,
+      "grad_norm": 38.582218170166016,
+      "learning_rate": 9.854052652406666e-06,
+      "loss": 11.0027,
+      "step": 830
+    },
+    {
+      "epoch": 0.17142857142857143,
+      "grad_norm": 32.440155029296875,
+      "learning_rate": 9.84538643114539e-06,
+      "loss": 10.5934,
+      "step": 840
+    },
+    {
+      "epoch": 0.17346938775510204,
+      "grad_norm": 35.402313232421875,
+      "learning_rate": 9.836474315195148e-06,
+      "loss": 10.5801,
+      "step": 850
+    },
+    {
+      "epoch": 0.17551020408163265,
+      "grad_norm": 31.911802291870117,
+      "learning_rate": 9.827316756829871e-06,
+      "loss": 10.5228,
+      "step": 860
+    },
+    {
+      "epoch": 0.17755102040816326,
+      "grad_norm": 34.77294921875,
+      "learning_rate": 9.817914220779258e-06,
+      "loss": 10.4315,
+      "step": 870
+    },
+    {
+      "epoch": 0.17959183673469387,
+      "grad_norm": 34.85014343261719,
+      "learning_rate": 9.808267184205182e-06,
+      "loss": 10.9996,
+      "step": 880
+    },
+    {
+      "epoch": 0.1816326530612245,
+      "grad_norm": 32.950138092041016,
+      "learning_rate": 9.798376136677486e-06,
+      "loss": 11.2809,
+      "step": 890
+    },
+    {
+      "epoch": 0.1836734693877551,
+      "grad_norm": 39.478111267089844,
+      "learning_rate": 9.788241580149123e-06,
+      "loss": 10.4687,
+      "step": 900
+    },
+    {
+      "epoch": 0.18571428571428572,
+      "grad_norm": 24.397432327270508,
+      "learning_rate": 9.777864028930705e-06,
+      "loss": 11.1725,
+      "step": 910
+    },
+    {
+      "epoch": 0.18775510204081633,
+      "grad_norm": 34.38156509399414,
+      "learning_rate": 9.767244009664376e-06,
+      "loss": 10.2822,
+      "step": 920
+    },
+    {
+      "epoch": 0.18979591836734694,
+      "grad_norm": 36.80168533325195,
+      "learning_rate": 9.75638206129711e-06,
+      "loss": 11.7147,
+      "step": 930
+    },
+    {
+      "epoch": 0.19183673469387755,
+      "grad_norm": 30.4884090423584,
+      "learning_rate": 9.745278735053345e-06,
+      "loss": 10.6125,
+      "step": 940
+    },
+    {
+      "epoch": 0.19387755102040816,
+      "grad_norm": 31.85347557067871,
+      "learning_rate": 9.733934594407012e-06,
+      "loss": 10.8146,
+      "step": 950
+    },
+    {
+      "epoch": 0.19591836734693877,
+      "grad_norm": 27.185823440551758,
+      "learning_rate": 9.722350215052946e-06,
+      "loss": 10.2984,
+      "step": 960
+    },
+    {
+      "epoch": 0.19795918367346937,
+      "grad_norm": 26.676513671875,
+      "learning_rate": 9.710526184877667e-06,
+      "loss": 10.459,
+      "step": 970
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 39.182044982910156,
+      "learning_rate": 9.698463103929542e-06,
+      "loss": 11.1628,
+      "step": 980
+    },
+    {
+      "epoch": 0.20204081632653062,
+      "grad_norm": 37.69694519042969,
+      "learning_rate": 9.68616158438834e-06,
+      "loss": 10.6955,
+      "step": 990
+    },
+    {
+      "epoch": 0.20408163265306123,
+      "grad_norm": 26.145389556884766,
+      "learning_rate": 9.673622250534155e-06,
+      "loss": 10.6714,
+      "step": 1000
+    },
+    {
+      "epoch": 0.20612244897959184,
+      "grad_norm": 25.507823944091797,
+      "learning_rate": 9.660845738715743e-06,
+      "loss": 10.6368,
+      "step": 1010
+    },
+    {
+      "epoch": 0.20816326530612245,
+      "grad_norm": 29.48382568359375,
+      "learning_rate": 9.647832697318207e-06,
+      "loss": 10.5836,
+      "step": 1020
+    },
+    {
+      "epoch": 0.21020408163265306,
+      "grad_norm": 34.13095474243164,
+      "learning_rate": 9.63458378673011e-06,
+      "loss": 10.7477,
+      "step": 1030
+    },
+    {
+      "epoch": 0.21224489795918366,
+      "grad_norm": 28.783233642578125,
+      "learning_rate": 9.621099679309948e-06,
+      "loss": 9.9524,
+      "step": 1040
+    },
+    {
+      "epoch": 0.21428571428571427,
+      "grad_norm": 28.7989501953125,
+      "learning_rate": 9.60738105935204e-06,
+      "loss": 10.4375,
+      "step": 1050
+    },
+    {
+      "epoch": 0.2163265306122449,
+      "grad_norm": 34.467567443847656,
+      "learning_rate": 9.593428623051793e-06,
+      "loss": 10.6146,
+      "step": 1060
+    },
+    {
+      "epoch": 0.21836734693877552,
+      "grad_norm": 31.702016830444336,
+      "learning_rate": 9.57924307847038e-06,
+      "loss": 10.4064,
+      "step": 1070
+    },
+    {
+      "epoch": 0.22040816326530613,
+      "grad_norm": 27.013303756713867,
+      "learning_rate": 9.564825145498795e-06,
+      "loss": 10.1444,
+      "step": 1080
+    },
+    {
+      "epoch": 0.22244897959183674,
+      "grad_norm": 30.175722122192383,
+      "learning_rate": 9.550175555821333e-06,
+      "loss": 10.8813,
+      "step": 1090
+    },
+    {
+      "epoch": 0.22448979591836735,
+      "grad_norm": 34.38163375854492,
+      "learning_rate": 9.53529505287845e-06,
+      "loss": 10.4208,
+      "step": 1100
+    },
+    {
+      "epoch": 0.22653061224489796,
+      "grad_norm": 29.349163055419922,
+      "learning_rate": 9.520184391829037e-06,
+      "loss": 9.8474,
+      "step": 1110
+    },
+    {
+      "epoch": 0.22857142857142856,
+      "grad_norm": 30.701316833496094,
+      "learning_rate": 9.504844339512096e-06,
+      "loss": 10.9479,
+      "step": 1120
+    },
+    {
+      "epoch": 0.23061224489795917,
+      "grad_norm": 30.510988235473633,
+      "learning_rate": 9.489275674407826e-06,
+      "loss": 10.7473,
+      "step": 1130
+    },
+    {
+      "epoch": 0.23265306122448978,
+      "grad_norm": 34.73878860473633,
+      "learning_rate": 9.473479186598115e-06,
+      "loss": 10.8098,
+      "step": 1140
+    },
+    {
+      "epoch": 0.23469387755102042,
+      "grad_norm": 30.67318344116211,
+      "learning_rate": 9.457455677726447e-06,
+      "loss": 10.3694,
+      "step": 1150
+    },
+    {
+      "epoch": 0.23673469387755103,
+      "grad_norm": 27.99265480041504,
+      "learning_rate": 9.441205960957221e-06,
+      "loss": 10.5622,
+      "step": 1160
+    },
+    {
+      "epoch": 0.23877551020408164,
+      "grad_norm": 34.095008850097656,
+      "learning_rate": 9.424730860934474e-06,
+      "loss": 10.4537,
+      "step": 1170
+    },
+    {
+      "epoch": 0.24081632653061225,
+      "grad_norm": 29.445648193359375,
+      "learning_rate": 9.408031213740045e-06,
+      "loss": 10.5178,
+      "step": 1180
+    },
+    {
+      "epoch": 0.24285714285714285,
+      "grad_norm": 32.19340896606445,
+      "learning_rate": 9.391107866851143e-06,
+      "loss": 10.3766,
+      "step": 1190
+    },
+    {
+      "epoch": 0.24489795918367346,
+      "grad_norm": 31.312353134155273,
+      "learning_rate": 9.37396167909733e-06,
+      "loss": 10.9467,
+      "step": 1200
+    },
+    {
+      "epoch": 0.24693877551020407,
+      "grad_norm": 36.48698806762695,
+      "learning_rate": 9.356593520616948e-06,
+      "loss": 10.7217,
+      "step": 1210
+    },
+    {
+      "epoch": 0.24897959183673468,
+      "grad_norm": 32.96704864501953,
+      "learning_rate": 9.33900427281295e-06,
+      "loss": 11.0355,
+      "step": 1220
+    },
+    {
+      "epoch": 0.2510204081632653,
+      "grad_norm": 32.133941650390625,
+      "learning_rate": 9.321194828308185e-06,
+      "loss": 10.7212,
+      "step": 1230
+    },
+    {
+      "epoch": 0.2530612244897959,
+      "grad_norm": 32.15143585205078,
+      "learning_rate": 9.303166090900082e-06,
+      "loss": 10.3076,
+      "step": 1240
+    },
+    {
+      "epoch": 0.25510204081632654,
+      "grad_norm": 31.86509895324707,
+      "learning_rate": 9.284918975514798e-06,
+      "loss": 10.1564,
+      "step": 1250
+    },
+    {
+      "epoch": 0.2571428571428571,
+      "grad_norm": 24.602365493774414,
+      "learning_rate": 9.266454408160779e-06,
+      "loss": 10.335,
+      "step": 1260
+    },
+    {
+      "epoch": 0.25918367346938775,
+      "grad_norm": 31.0269775390625,
+      "learning_rate": 9.24777332588177e-06,
+      "loss": 11.5042,
+      "step": 1270
+    },
+    {
+      "epoch": 0.2612244897959184,
+      "grad_norm": 24.591663360595703,
+      "learning_rate": 9.22887667670926e-06,
+      "loss": 11.0083,
+      "step": 1280
+    },
+    {
+      "epoch": 0.26326530612244897,
+      "grad_norm": 34.8664436340332,
+      "learning_rate": 9.209765419614375e-06,
+      "loss": 11.2456,
+      "step": 1290
+    },
+    {
+      "epoch": 0.2653061224489796,
+      "grad_norm": 28.38134765625,
+      "learning_rate": 9.190440524459203e-06,
+      "loss": 10.9234,
+      "step": 1300
+    },
+    {
+      "epoch": 0.2673469387755102,
+      "grad_norm": 30.83916473388672,
+      "learning_rate": 9.170902971947589e-06,
+      "loss": 10.4725,
+      "step": 1310
+    },
+    {
+      "epoch": 0.2693877551020408,
+      "grad_norm": 33.667266845703125,
+      "learning_rate": 9.151153753575351e-06,
+      "loss": 11.2772,
+      "step": 1320
+    },
+    {
+      "epoch": 0.2714285714285714,
+      "grad_norm": 82.6866226196289,
+      "learning_rate": 9.131193871579975e-06,
+      "loss": 10.3814,
+      "step": 1330
+    },
+    {
+      "epoch": 0.27346938775510204,
+      "grad_norm": 38.70499801635742,
+      "learning_rate": 9.111024338889748e-06,
+      "loss": 10.8925,
+      "step": 1340
+    },
+    {
+      "epoch": 0.2755102040816326,
+      "grad_norm": 27.011062622070312,
+      "learning_rate": 9.090646179072352e-06,
+      "loss": 10.5997,
+      "step": 1350
+    },
+    {
+      "epoch": 0.27755102040816326,
+      "grad_norm": 27.594419479370117,
+      "learning_rate": 9.070060426282924e-06,
+      "loss": 10.6429,
+      "step": 1360
+    },
+    {
+      "epoch": 0.2795918367346939,
+      "grad_norm": 31.248558044433594,
+      "learning_rate": 9.049268125211577e-06,
+      "loss": 10.6206,
+      "step": 1370
+    },
+    {
+      "epoch": 0.2816326530612245,
+      "grad_norm": 29.369421005249023,
+      "learning_rate": 9.028270331030373e-06,
+      "loss": 10.7207,
+      "step": 1380
+    },
+    {
+      "epoch": 0.2836734693877551,
+      "grad_norm": 29.121797561645508,
+      "learning_rate": 9.007068109339783e-06,
+      "loss": 10.9274,
+      "step": 1390
+    },
+    {
+      "epoch": 0.2857142857142857,
+      "grad_norm": 27.654598236083984,
+      "learning_rate": 8.985662536114614e-06,
+      "loss": 10.0533,
+      "step": 1400
+    },
+    {
+      "epoch": 0.28775510204081634,
+      "grad_norm": 35.461509704589844,
+      "learning_rate": 8.964054697649389e-06,
+      "loss": 11.1923,
+      "step": 1410
+    },
+    {
+      "epoch": 0.2897959183673469,
+      "grad_norm": 34.02452850341797,
+      "learning_rate": 8.94224569050324e-06,
+      "loss": 10.993,
+      "step": 1420
+    },
+    {
+      "epoch": 0.29183673469387755,
+      "grad_norm": 29.19350814819336,
+      "learning_rate": 8.920236621444243e-06,
+      "loss": 10.6823,
+      "step": 1430
+    },
+    {
+      "epoch": 0.2938775510204082,
+      "grad_norm": 31.622766494750977,
+      "learning_rate": 8.89802860739326e-06,
+      "loss": 10.8623,
+      "step": 1440
+    },
+    {
+      "epoch": 0.29591836734693877,
+      "grad_norm": 27.097755432128906,
+      "learning_rate": 8.87562277536726e-06,
+      "loss": 10.6118,
+      "step": 1450
+    },
+    {
+      "epoch": 0.2979591836734694,
+      "grad_norm": 27.950016021728516,
+      "learning_rate": 8.853020262422111e-06,
+      "loss": 10.3304,
+      "step": 1460
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 32.812591552734375,
+      "learning_rate": 8.83022221559489e-06,
+      "loss": 10.1491,
+      "step": 1470
+    },
+    {
+      "epoch": 0.3020408163265306,
+      "grad_norm": 29.947307586669922,
+      "learning_rate": 8.807229791845673e-06,
+      "loss": 10.3768,
+      "step": 1480
+    },
+    {
+      "epoch": 0.3040816326530612,
+      "grad_norm": 29.019681930541992,
+      "learning_rate": 8.78404415799881e-06,
+      "loss": 10.736,
+      "step": 1490
+    },
+    {
+      "epoch": 0.30612244897959184,
+      "grad_norm": 27.67340087890625,
+      "learning_rate": 8.76066649068372e-06,
+      "loss": 10.3324,
+      "step": 1500
+    },
+    {
+      "epoch": 0.3081632653061224,
+      "grad_norm": 27.788095474243164,
+      "learning_rate": 8.737097976275177e-06,
+      "loss": 10.6988,
+      "step": 1510
+    },
+    {
+      "epoch": 0.31020408163265306,
+      "grad_norm": 33.040504455566406,
+      "learning_rate": 8.713339810833105e-06,
+      "loss": 10.7684,
+      "step": 1520
+    },
+    {
+      "epoch": 0.3122448979591837,
+      "grad_norm": 30.00772476196289,
+      "learning_rate": 8.689393200041878e-06,
+      "loss": 10.7234,
+      "step": 1530
+    },
+    {
+      "epoch": 0.3142857142857143,
+      "grad_norm": 31.760778427124023,
+      "learning_rate": 8.665259359149132e-06,
+      "loss": 10.6876,
+      "step": 1540
+    },
+    {
+      "epoch": 0.3163265306122449,
+      "grad_norm": 32.49633026123047,
+      "learning_rate": 8.640939512904097e-06,
+      "loss": 10.3881,
+      "step": 1550
+    },
+    {
+      "epoch": 0.3183673469387755,
+      "grad_norm": 32.45978546142578,
+      "learning_rate": 8.61643489549544e-06,
+      "loss": 9.7477,
+      "step": 1560
+    },
+    {
+      "epoch": 0.32040816326530613,
+      "grad_norm": 28.2198429107666,
+      "learning_rate": 8.591746750488639e-06,
+      "loss": 10.0062,
+      "step": 1570
+    },
+    {
+      "epoch": 0.3224489795918367,
+      "grad_norm": 32.73541259765625,
+      "learning_rate": 8.566876330762861e-06,
+      "loss": 10.8813,
+      "step": 1580
+    },
+    {
+      "epoch": 0.32448979591836735,
+      "grad_norm": 30.3160400390625,
+      "learning_rate": 8.541824898447399e-06,
+      "loss": 11.127,
+      "step": 1590
+    },
+    {
+      "epoch": 0.32653061224489793,
+      "grad_norm": 21.944679260253906,
+      "learning_rate": 8.516593724857598e-06,
+      "loss": 10.8866,
+      "step": 1600
+    },
+    {
+      "epoch": 0.32857142857142857,
+      "grad_norm": 29.54482078552246,
+      "learning_rate": 8.491184090430365e-06,
+      "loss": 10.3385,
+      "step": 1610
+    },
+    {
+      "epoch": 0.3306122448979592,
+      "grad_norm": 32.63984680175781,
+      "learning_rate": 8.465597284659163e-06,
+      "loss": 10.7293,
+      "step": 1620
+    },
+    {
+      "epoch": 0.3326530612244898,
+      "grad_norm": 30.096092224121094,
+      "learning_rate": 8.439834606028594e-06,
+      "loss": 10.8387,
+      "step": 1630
+    },
+    {
+      "epoch": 0.3346938775510204,
+      "grad_norm": 26.435718536376953,
+      "learning_rate": 8.413897361948484e-06,
+      "loss": 10.7297,
+      "step": 1640
+    },
+    {
+      "epoch": 0.336734693877551,
+      "grad_norm": 29.123828887939453,
+      "learning_rate": 8.387786868687549e-06,
+      "loss": 10.2166,
+      "step": 1650
+    },
+    {
+      "epoch": 0.33877551020408164,
+      "grad_norm": 28.276975631713867,
+      "learning_rate": 8.361504451306585e-06,
+      "loss": 10.2483,
+      "step": 1660
+    },
+    {
+      "epoch": 0.3408163265306122,
+      "grad_norm": 27.21367073059082,
+      "learning_rate": 8.335051443591236e-06,
+      "loss": 10.6175,
+      "step": 1670
+    },
+    {
+      "epoch": 0.34285714285714286,
+      "grad_norm": 27.0053768157959,
+      "learning_rate": 8.308429187984298e-06,
+      "loss": 11.1318,
+      "step": 1680
+    },
+    {
+      "epoch": 0.3448979591836735,
+      "grad_norm": 32.046974182128906,
+      "learning_rate": 8.281639035517591e-06,
+      "loss": 10.2026,
+      "step": 1690
+    },
+    {
+      "epoch": 0.3469387755102041,
+      "grad_norm": 20.883655548095703,
+      "learning_rate": 8.254682345743406e-06,
+      "loss": 10.4051,
+      "step": 1700
+    },
+    {
+      "epoch": 0.3489795918367347,
+      "grad_norm": 26.248065948486328,
+      "learning_rate": 8.227560486665498e-06,
+      "loss": 10.0888,
+      "step": 1710
+    },
+    {
+      "epoch": 0.3510204081632653,
+      "grad_norm": 27.3173770904541,
+      "learning_rate": 8.200274834669675e-06,
+      "loss": 10.6525,
+      "step": 1720
+    },
+    {
+      "epoch": 0.35306122448979593,
+      "grad_norm": 29.715709686279297,
+      "learning_rate": 8.172826774453937e-06,
+      "loss": 9.9913,
+      "step": 1730
+    },
+    {
+      "epoch": 0.3551020408163265,
+      "grad_norm": 34.22426223754883,
+      "learning_rate": 8.145217698958213e-06,
+      "loss": 10.4346,
+      "step": 1740
+    },
+    {
+      "epoch": 0.35714285714285715,
+      "grad_norm": 33.10298156738281,
+      "learning_rate": 8.117449009293668e-06,
+      "loss": 10.263,
+      "step": 1750
+    },
+    {
+      "epoch": 0.35918367346938773,
+      "grad_norm": 36.919559478759766,
+      "learning_rate": 8.089522114671603e-06,
+      "loss": 10.1763,
+      "step": 1760
+    },
+    {
+      "epoch": 0.36122448979591837,
+      "grad_norm": 29.963897705078125,
+      "learning_rate": 8.061438432331935e-06,
+      "loss": 10.1127,
+      "step": 1770
+    },
+    {
+      "epoch": 0.363265306122449,
+      "grad_norm": 26.98243522644043,
+      "learning_rate": 8.033199387471278e-06,
+      "loss": 10.5181,
+      "step": 1780
+    },
+    {
+      "epoch": 0.3653061224489796,
+      "grad_norm": 34.38954162597656,
+      "learning_rate": 8.004806413170613e-06,
+      "loss": 11.2506,
+      "step": 1790
+    },
+    {
+      "epoch": 0.3673469387755102,
+      "grad_norm": 27.817028045654297,
+      "learning_rate": 7.976260950322572e-06,
+      "loss": 9.9205,
+      "step": 1800
+    },
+    {
+      "epoch": 0.3693877551020408,
+      "grad_norm": 29.21490478515625,
+      "learning_rate": 7.9475644475583e-06,
+      "loss": 9.7749,
+      "step": 1810
+    },
+    {
+      "epoch": 0.37142857142857144,
+      "grad_norm": 36.39332962036133,
+      "learning_rate": 7.918718361173951e-06,
+      "loss": 10.1115,
+      "step": 1820
+    },
+    {
+      "epoch": 0.373469387755102,
+      "grad_norm": 29.05034637451172,
+      "learning_rate": 7.889724155056776e-06,
+      "loss": 9.8026,
+      "step": 1830
+    },
+    {
+      "epoch": 0.37551020408163266,
+      "grad_norm": 31.34714126586914,
+      "learning_rate": 7.860583300610849e-06,
+      "loss": 9.9526,
+      "step": 1840
+    },
+    {
+      "epoch": 0.37755102040816324,
+      "grad_norm": 30.960601806640625,
+      "learning_rate": 7.83129727668237e-06,
+      "loss": 10.3724,
+      "step": 1850
+    },
+    {
+      "epoch": 0.3795918367346939,
+      "grad_norm": 27.005430221557617,
+      "learning_rate": 7.801867569484635e-06,
+      "loss": 10.0523,
+      "step": 1860
+    },
+    {
+      "epoch": 0.3816326530612245,
+      "grad_norm": 32.369354248046875,
+      "learning_rate": 7.772295672522615e-06,
+      "loss": 10.5823,
+      "step": 1870
+    },
+    {
+      "epoch": 0.3836734693877551,
+      "grad_norm": 27.78652572631836,
+      "learning_rate": 7.742583086517151e-06,
+      "loss": 9.75,
+      "step": 1880
+    },
+    {
+      "epoch": 0.38571428571428573,
+      "grad_norm": 50.487850189208984,
+      "learning_rate": 7.712731319328798e-06,
+      "loss": 10.7037,
+      "step": 1890
+    },
+    {
+      "epoch": 0.3877551020408163,
+      "grad_norm": 25.972570419311523,
+      "learning_rate": 7.682741885881314e-06,
+      "loss": 10.418,
+      "step": 1900
+    },
+    {
+      "epoch": 0.38979591836734695,
+      "grad_norm": 32.993350982666016,
+      "learning_rate": 7.652616308084774e-06,
+      "loss": 9.834,
+      "step": 1910
+    },
+    {
+      "epoch": 0.39183673469387753,
+      "grad_norm": 30.098678588867188,
+      "learning_rate": 7.622356114758328e-06,
+      "loss": 10.5912,
+      "step": 1920
+    },
+    {
+      "epoch": 0.39387755102040817,
+      "grad_norm": 33.64377212524414,
+      "learning_rate": 7.591962841552627e-06,
+      "loss": 10.5525,
+      "step": 1930
+    },
+    {
+      "epoch": 0.39591836734693875,
+      "grad_norm": 40.15250015258789,
+      "learning_rate": 7.561438030871886e-06,
+      "loss": 10.4681,
+      "step": 1940
+    },
+    {
+      "epoch": 0.3979591836734694,
+      "grad_norm": 19.717267990112305,
+      "learning_rate": 7.530783231795615e-06,
+      "loss": 9.9533,
+      "step": 1950
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 22.303646087646484,
+      "learning_rate": 7.500000000000001e-06,
+      "loss": 9.7424,
+      "step": 1960
+    },
+    {
+      "epoch": 0.4020408163265306,
+      "grad_norm": 25.239171981811523,
+      "learning_rate": 7.469089897678958e-06,
+      "loss": 10.306,
+      "step": 1970
+    },
+    {
+      "epoch": 0.40408163265306124,
+      "grad_norm": 33.91167449951172,
+      "learning_rate": 7.438054493464859e-06,
+      "loss": 10.0006,
+      "step": 1980
+    },
+    {
+      "epoch": 0.4061224489795918,
+      "grad_norm": 25.457975387573242,
+      "learning_rate": 7.406895362348916e-06,
+      "loss": 10.2729,
+      "step": 1990
+    },
+    {
+      "epoch": 0.40816326530612246,
+      "grad_norm": 30.050312042236328,
+      "learning_rate": 7.375614085601265e-06,
+      "loss": 10.4351,
+      "step": 2000
+    },
+    {
+      "epoch": 0.41020408163265304,
+      "grad_norm": 27.26880645751953,
+      "learning_rate": 7.344212250690712e-06,
+      "loss": 9.7646,
+      "step": 2010
+    },
+    {
+      "epoch": 0.4122448979591837,
+      "grad_norm": 26.21525764465332,
+      "learning_rate": 7.312691451204178e-06,
+      "loss": 10.0521,
+      "step": 2020
+    },
+    {
+      "epoch": 0.4142857142857143,
+      "grad_norm": 30.313282012939453,
+      "learning_rate": 7.281053286765816e-06,
+      "loss": 10.2,
+      "step": 2030
+    },
+    {
+      "epoch": 0.4163265306122449,
+      "grad_norm": 24.833383560180664,
+      "learning_rate": 7.249299362955846e-06,
+      "loss": 9.8594,
+      "step": 2040
+    },
+    {
+      "epoch": 0.41836734693877553,
+      "grad_norm": 32.73284149169922,
+      "learning_rate": 7.217431291229068e-06,
+      "loss": 10.1855,
+      "step": 2050
+    },
+    {
+      "epoch": 0.4204081632653061,
+      "grad_norm": 39.904109954833984,
+      "learning_rate": 7.185450688833083e-06,
+      "loss": 9.1649,
+      "step": 2060
+    },
+    {
+      "epoch": 0.42244897959183675,
+      "grad_norm": 29.54826545715332,
+      "learning_rate": 7.153359178726222e-06,
+      "loss": 10.1546,
+      "step": 2070
+    },
+    {
+      "epoch": 0.42448979591836733,
+      "grad_norm": 21.587451934814453,
+      "learning_rate": 7.121158389495187e-06,
+      "loss": 9.8794,
+      "step": 2080
+    },
+    {
+      "epoch": 0.42653061224489797,
+      "grad_norm": 34.06581497192383,
+      "learning_rate": 7.088849955272396e-06,
+      "loss": 9.1326,
+      "step": 2090
+    },
+    {
+      "epoch": 0.42857142857142855,
+      "grad_norm": 25.830421447753906,
+      "learning_rate": 7.056435515653059e-06,
+      "loss": 10.9292,
+      "step": 2100
+    },
+    {
+      "epoch": 0.4306122448979592,
+      "grad_norm": 26.324182510375977,
+      "learning_rate": 7.023916715611969e-06,
+      "loss": 10.341,
+      "step": 2110
+    },
+    {
+      "epoch": 0.4326530612244898,
+      "grad_norm": 28.438213348388672,
+      "learning_rate": 6.991295205420028e-06,
+      "loss": 10.7136,
+      "step": 2120
+    },
+    {
+      "epoch": 0.4346938775510204,
+      "grad_norm": 27.563526153564453,
+      "learning_rate": 6.9585726405604915e-06,
+      "loss": 10.1619,
+      "step": 2130
+    },
+    {
+      "epoch": 0.43673469387755104,
+      "grad_norm": 22.530603408813477,
+      "learning_rate": 6.925750681644954e-06,
+      "loss": 10.9317,
+      "step": 2140
+    },
+    {
+      "epoch": 0.4387755102040816,
+      "grad_norm": 31.136659622192383,
+      "learning_rate": 6.892830994329089e-06,
+      "loss": 10.3304,
+      "step": 2150
+    },
+    {
+      "epoch": 0.44081632653061226,
+      "grad_norm": 29.56623649597168,
+      "learning_rate": 6.859815249228106e-06,
+      "loss": 9.6344,
+      "step": 2160
+    },
+    {
+      "epoch": 0.44285714285714284,
+      "grad_norm": 24.045392990112305,
+      "learning_rate": 6.8267051218319766e-06,
+      "loss": 9.3313,
+      "step": 2170
+    },
+    {
+      "epoch": 0.4448979591836735,
+      "grad_norm": 25.503190994262695,
+      "learning_rate": 6.793502292420402e-06,
+      "loss": 10.3088,
+      "step": 2180
+    },
+    {
+      "epoch": 0.44693877551020406,
+      "grad_norm": 30.17300033569336,
+      "learning_rate": 6.760208445977551e-06,
+      "loss": 10.0106,
+      "step": 2190
+    },
+    {
+      "epoch": 0.4489795918367347,
+      "grad_norm": 26.85445213317871,
+      "learning_rate": 6.726825272106539e-06,
+      "loss": 9.552,
+      "step": 2200
+    },
+    {
+      "epoch": 0.45102040816326533,
+      "grad_norm": 32.817073822021484,
+      "learning_rate": 6.693354464943689e-06,
+      "loss": 10.7139,
+      "step": 2210
+    },
+    {
+      "epoch": 0.4530612244897959,
+      "grad_norm": 25.95104217529297,
+      "learning_rate": 6.659797723072558e-06,
+      "loss": 10.8172,
+      "step": 2220
+    },
+    {
+      "epoch": 1.0020408163265306,
+      "grad_norm": 24.709354400634766,
+      "learning_rate": 6.626156749437736e-06,
+      "loss": 9.865,
+      "step": 2230
+    },
+    {
+      "epoch": 1.0040816326530613,
+      "grad_norm": 30.658660888671875,
+      "learning_rate": 6.592433251258423e-06,
+      "loss": 9.4195,
+      "step": 2240
+    },
+    {
+      "epoch": 1.006122448979592,
+      "grad_norm": 26.727264404296875,
+      "learning_rate": 6.558628939941792e-06,
+      "loss": 10.2027,
+      "step": 2250
+    },
+    {
+      "epoch": 1.0081632653061225,
+      "grad_norm": 30.268173217773438,
+      "learning_rate": 6.524745530996137e-06,
+      "loss": 10.0028,
+      "step": 2260
+    },
+    {
+      "epoch": 1.010204081632653,
+      "grad_norm": 30.007192611694336,
+      "learning_rate": 6.490784743943819e-06,
+      "loss": 10.251,
+      "step": 2270
+    },
+    {
+      "epoch": 1.0122448979591836,
+      "grad_norm": 32.65926742553711,
+      "learning_rate": 6.456748302233995e-06,
+      "loss": 10.119,
+      "step": 2280
+    },
+    {
+      "epoch": 1.0142857142857142,
+      "grad_norm": 28.551115036010742,
+      "learning_rate": 6.4226379331551625e-06,
+      "loss": 9.8939,
+      "step": 2290
+    },
+    {
+      "epoch": 1.0163265306122449,
+      "grad_norm": 25.979280471801758,
+      "learning_rate": 6.388455367747503e-06,
+      "loss": 9.8234,
+      "step": 2300
+    },
+    {
+      "epoch": 1.0183673469387755,
+      "grad_norm": 30.16221046447754,
+      "learning_rate": 6.354202340715027e-06,
+      "loss": 9.1357,
+      "step": 2310
+    },
+    {
+      "epoch": 1.0204081632653061,
+      "grad_norm": 25.45052146911621,
+      "learning_rate": 6.319880590337549e-06,
+      "loss": 9.2762,
+      "step": 2320
+    },
+    {
+      "epoch": 1.0224489795918368,
+      "grad_norm": 31.70077133178711,
+      "learning_rate": 6.2854918583824745e-06,
+      "loss": 9.2601,
+      "step": 2330
+    },
+    {
+      "epoch": 1.0244897959183674,
+      "grad_norm": 27.903934478759766,
+      "learning_rate": 6.251037890016396e-06,
+      "loss": 9.3681,
+      "step": 2340
+    },
+    {
+      "epoch": 1.026530612244898,
+      "grad_norm": 35.41598892211914,
+      "learning_rate": 6.216520433716544e-06,
+      "loss": 9.5135,
+      "step": 2350
+    },
+    {
+      "epoch": 1.0285714285714285,
+      "grad_norm": 32.13457489013672,
+      "learning_rate": 6.181941241182044e-06,
+      "loss": 9.5838,
+      "step": 2360
+    },
+    {
+      "epoch": 1.030612244897959,
+      "grad_norm": 28.334228515625,
+      "learning_rate": 6.1473020672450275e-06,
+      "loss": 10.0139,
+      "step": 2370
+    },
+    {
+      "epoch": 1.0326530612244897,
+      "grad_norm": 25.681198120117188,
+      "learning_rate": 6.112604669781572e-06,
+      "loss": 9.8882,
+      "step": 2380
+    },
+    {
+      "epoch": 1.0346938775510204,
+      "grad_norm": 24.19563865661621,
+      "learning_rate": 6.0778508096224985e-06,
+      "loss": 8.789,
+      "step": 2390
+    },
+    {
+      "epoch": 1.036734693877551,
+      "grad_norm": 28.59593391418457,
+      "learning_rate": 6.043042250464005e-06,
+      "loss": 8.9804,
+      "step": 2400
+    },
+    {
+      "epoch": 1.0387755102040817,
+      "grad_norm": 32.59727478027344,
+      "learning_rate": 6.008180758778167e-06,
+      "loss": 8.9337,
+      "step": 2410
+    },
+    {
+      "epoch": 1.0408163265306123,
+      "grad_norm": 28.587871551513672,
+      "learning_rate": 5.973268103723293e-06,
+      "loss": 7.8557,
+      "step": 2420
+    },
+    {
+      "epoch": 1.042857142857143,
+      "grad_norm": 23.32282257080078,
+      "learning_rate": 5.938306057054139e-06,
+      "loss": 9.3373,
+      "step": 2430
+    },
+    {
+      "epoch": 1.0448979591836736,
+      "grad_norm": 28.35025978088379,
+      "learning_rate": 5.903296393031996e-06,
+      "loss": 8.694,
+      "step": 2440
+    },
+    {
+      "epoch": 1.046938775510204,
+      "grad_norm": 31.087772369384766,
+      "learning_rate": 5.8682408883346535e-06,
+      "loss": 8.2314,
+      "step": 2450
+    },
+    {
+      "epoch": 1.0489795918367346,
+      "grad_norm": 28.746549606323242,
+      "learning_rate": 5.8331413219662295e-06,
+      "loss": 8.8584,
+      "step": 2460
+    },
+    {
+      "epoch": 1.0510204081632653,
+      "grad_norm": 21.93814468383789,
+      "learning_rate": 5.797999475166897e-06,
+      "loss": 8.654,
+      "step": 2470
+    },
+    {
+      "epoch": 1.0530612244897959,
+      "grad_norm": 27.251869201660156,
+      "learning_rate": 5.762817131322482e-06,
+      "loss": 9.0667,
+      "step": 2480
+    },
+    {
+      "epoch": 1.0551020408163265,
+      "grad_norm": 28.02555274963379,
+      "learning_rate": 5.7275960758739655e-06,
+      "loss": 8.2027,
+      "step": 2490
+    },
+    {
+      "epoch": 1.0571428571428572,
+      "grad_norm": 27.519668579101562,
+      "learning_rate": 5.69233809622687e-06,
+      "loss": 9.0829,
+      "step": 2500
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 4900,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 607929349701632.0,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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