update

.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

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|audio_end|>": 151653,
+  "<|audio_pad|>": 151654,
+  "<|audio_start|>": 151652,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656
+}

chat_template.jinja

@@ -0,0 +1,4 @@
+{% for message in messages %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}{% else %}{% for content in message['content'] %}{% if content.get('type') == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}<|im_end|>
+{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

config.json

@@ -0,0 +1,89 @@
+{
+  "model_type": "moss_tts_delay",
+  "architectures": [
+    "MossTTSDelayModel"
+  ],
+  "auto_map": {
+      "AutoConfig": "configuration_moss_tts.MossTTSDelayConfig",
+      "AutoModel": "modeling_moss_tts.MossTTSDelayModel"
+  },
+  "dtype": "bfloat16",
+  "initializer_range": 0.02,
+  "language_config": {
+    "_name_or_path": "Qwen/Qwen3-8B",
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": 151645,
+    "pad_token_id": 151643,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 40960,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 155648
+  },
+  "n_vq": 16,
+  "audio_vocab_size": 1024,
+  "audio_user_slot_token_id": 151654,
+  "audio_assistant_gen_slot_token_id": 151656,
+  "audio_assistant_delay_slot_token_id": 151662,
+  "audio_start_token_id": 151652,
+  "audio_end_token_id": 151653,
+  "audio_pad_code": 1024,
+  "sampling_rate": 24000,
+  "transformers_version": "4.57.1"
+}

configuration_moss_tts.py

@@ -0,0 +1,114 @@
+# coding=utf-8
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" MossTTSDelay model configuration """
+
+from typing import Optional, Union
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers.models.qwen3 import Qwen3Config
+
+logger = logging.get_logger(__name__)
+
+
+class MossTTSDelayConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossTTSDelayModel`]. It is used to instantiate an
+    MossTTSDelay model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the MossTTSDelay [MossTTSDelay-8B](https://huggingface.co/OpenMOSS/mosstts-8b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        language_config (`Union[Qwen3Config, dict]`, *optional*):
+            Configuration for the backbone language model (Qwen3).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        n_vq (`int`, *optional*, defaults to 16):
+            Number of additional VQ (Vector Quantization) heads/channels for audio. 
+            Determines the number of codebooks used in the audio representation.
+        audio_vocab_size (`int`, *optional*, defaults to 1024):
+            Vocabulary size for the audio tokens (codebooks 1 to N).
+        audio_user_slot_token_id (`int`, *optional*, defaults to 151654):
+            The specific token ID used as a placeholder/slot for user-side audio inputs in the prompt.
+        audio_assistant_gen_slot_token_id (`int`, *optional*, defaults to 151656):
+            The specific token ID representing the generation slot for the assistant's audio output. 
+            Acting as the trigger for the TTS generation process.
+        audio_assistant_delay_slot_token_id (`int`, *optional*, defaults to 151662):
+            The token ID used in the 'Delay Pattern' paradigm to represent the delayed/offset positions 
+            between different VQ channels.
+        audio_start_token_id (`int`, *optional*, defaults to 151652):
+            Special token ID used to denote the start of an audio sequence in the stream.
+        audio_end_token_id (`int`, *optional*, defaults to 151653):
+            Special token ID used to denote the end of an audio sequence (EOS for audio).
+        audio_pad_code (`int`, *optional*, defaults to 1024):
+            The padding value used within the audio VQ codebooks. Typically equals `audio_vocab_size`.
+    """
+    model_type = "moss_tts_delay"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        language_config: Optional[Union[Qwen3Config, dict]] = None,
+        initializer_range: float = 0.02,
+        n_vq: int = 16,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_vocab_size: int = 1024,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_gen_slot_token_id: int = 151656,
+        audio_assistant_delay_slot_token_id: int = 151662,
+        audio_start_token_id: int = 151652,
+        audio_end_token_id: int = 151653,
+        audio_pad_code: int = 1024,
+        sampling_rate: int = 24000,
+        **kwargs,
+    ):
+        if isinstance(language_config, dict):
+            self.language_config = Qwen3Config(**language_config)
+        elif language_config is None:
+            self.language_config = Qwen3Config()
+        else:
+            self.language_config = language_config
+
+        self.initializer_range = initializer_range
+        self.n_vq = n_vq
+        self.audio_vocab_size = audio_vocab_size
+        self.audio_user_slot_token_id = audio_user_slot_token_id
+        self.audio_assistant_gen_slot_token_id = audio_assistant_gen_slot_token_id
+        self.audio_assistant_delay_slot_token_id = audio_assistant_delay_slot_token_id
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.audio_pad_code = audio_pad_code
+        self.sampling_rate = sampling_rate
+
+        self.hidden_size = self.language_config.hidden_size
+        self.vocab_size = self.language_config.vocab_size
+        self.im_start_token_id = self.language_config
+        self.pad_token_id = pad_token_id
+        self.im_start_token_id = im_start_token_id
+        self.im_end_token_id = im_end_token_id
+
+        
+        super().__init__(**kwargs)
+
+    def to_dict(self):
+        output = super().to_dict()
+        if hasattr(self.language_config, "to_dict"):
+            output["language_config"] = self.language_config.to_dict()
+        else:
+            output["language_config"] = self.language_config
+        return output

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "do_sample": true,
+  "temperature": 1.1,
+  "top_p": 0.9,
+  "top_k": 50,
+  "repetition_penalty": 1.1,
+  "max_new_tokens": 8192
+}

inference_utils.py

@@ -0,0 +1,154 @@
+import torch
+import torchaudio
+import torch.nn.functional as F
+from typing import Optional, List, Tuple
+from tqdm import tqdm
+
+
+def apply_top_k(logits, top_k):
+    batch_size, vocab_size = logits.shape
+    top_k = min(top_k, vocab_size)
+    top_k_values, top_k_indices = torch.topk(logits, top_k, dim=-1)
+    filtered_logits = torch.full_like(logits, float("-inf"))
+    batch_indices = torch.arange(batch_size).unsqueeze(-1)
+    filtered_logits[batch_indices, top_k_indices] = top_k_values
+    return filtered_logits
+
+
+def apply_top_p(logits, top_p):
+    probs = F.softmax(logits, dim=-1)
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+    batch_size = logits.shape[0]
+    filtered_logits = logits.clone()
+    for i in range(batch_size):
+        indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]
+        filtered_logits[i, indices_to_remove] = float("-inf")
+    return filtered_logits
+
+
+def apply_top_p_optimized(logits, top_p):
+    probs = F.softmax(logits, dim=-1)
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+
+    indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(
+        dim=-1, index=sorted_indices, src=sorted_indices_to_remove
+    )
+
+    logits[indices_to_remove] = float("-inf")
+    return logits
+
+
+def apply_repetition_penalty_delay_pattern(
+    logits: torch.Tensor,
+    prev_tokens: torch.LongTensor,
+    penalty: float,
+):
+    """
+    logits: [B, H, V]  or [N, V]
+    prev_tokens: [B, T, H] or [N, T] or [B, H]
+
+    Apply the repetition penalty independently for each H (VQ head).
+    """
+    if penalty == 1.0 or prev_tokens is None:
+        return logits
+
+    vocab_size = logits.size(-1)
+
+    # Case 1: regular [N, V] (text layer)
+    if logits.dim() == 2:
+        prev_tokens_flat = prev_tokens.reshape(-1)
+        unique_tokens = torch.unique(prev_tokens_flat)
+
+        token_logits = logits[:, unique_tokens]
+        pos_mask = token_logits > 0
+        token_logits[pos_mask] /= penalty
+        token_logits[~pos_mask] *= penalty
+        logits[:, unique_tokens] = token_logits
+        return logits
+
+    # Case 2: Delay Pattern audio [B, H, V]
+    assert logits.dim() == 3, "Delay Pattern audio logits must be [B, H, V]"
+    B, H, V = logits.shape
+
+    for h in range(H):
+        # prev_tokens_h: [B, T] or [B]
+        prev_tokens_h = prev_tokens[..., h].reshape(-1)
+        unique_tokens = torch.unique(prev_tokens_h)
+
+        if unique_tokens.numel() == 0:
+            continue
+
+        token_logits = logits[:, h, unique_tokens]
+        pos_mask = token_logits > 0
+        token_logits[pos_mask] /= penalty
+        token_logits[~pos_mask] *= penalty
+        logits[:, h, unique_tokens] = token_logits
+
+    return logits
+
+
+def sample_token(
+    logits,
+    prev_tokens: Optional[torch.LongTensor] = None,
+    repetition_penalty: float = 1.0,
+    top_p=None,
+    top_k=None,
+    do_sample=True,
+):
+    vocab_size = logits.size(-1)
+
+    # ===== Repetition Penalty (before reshaping!) =====
+    if prev_tokens is not None and repetition_penalty != 1.0:
+        logits = apply_repetition_penalty_delay_pattern(
+            logits,
+            prev_tokens,
+            repetition_penalty,
+        )
+
+    if not do_sample:
+        return torch.argmax(logits, dim=-1)
+
+    # ===== Only flatten after this, for top-k / top-p / multinomial =====
+    original_shape = logits.shape
+    reshaped_logits = logits.view(-1, vocab_size)
+
+    if top_k is not None and top_k > 0:
+        reshaped_logits = apply_top_k(reshaped_logits, top_k)
+
+    if top_p is not None and top_p < 1.0:
+        reshaped_logits = apply_top_p_optimized(reshaped_logits, top_p)
+
+    probs = F.softmax(reshaped_logits, dim=-1)
+    next_tokens = torch.multinomial(probs, num_samples=1)
+
+    return next_tokens.view(original_shape[:-1])
+
+
+def find_last_equal_C(tensor, C):
+    """
+    tensor: torch.Tensor of shape [batch_size, seq_len]
+    C: scalar value to match
+    Returns: torch.Tensor of shape [batch_size] with last indices
+    """
+    mask = (tensor == C).int()  # Shape: [batch_size, seq_len], bool tensor
+    flipped_mask = mask.flip(dims=[1])  # Flip along sequence dimension
+    flipped_indices = flipped_mask.argmax(dim=1)  # First True in flipped
+    seq_len = tensor.shape[1]
+    last_indices = (seq_len - 1) - flipped_indices  # Convert to original indices
+
+    # Optional: Handle cases with no C (set to -1), though problem assumes existence
+    actual_values = tensor[torch.arange(tensor.shape[0]), last_indices]
+    no_match = actual_values != C
+    last_indices[no_match] = -1
+
+    return last_indices

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,439 @@
+{
+  "metadata": {
+    "total_parameters": 8355492864,
+    "total_size": 16710985728
+  },
+  "weight_map": {
+    "emb_ext.0.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.1.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.10.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.11.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.12.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.13.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.14.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.15.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.2.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.3.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.4.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.5.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.6.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.7.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.8.weight": "model-00004-of-00004.safetensors",
+    "emb_ext.9.weight": "model-00004-of-00004.safetensors",
+    "language_model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.20.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.21.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.22.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.22.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.32.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.33.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.34.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "language_model.layers.35.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.35.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "language_model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "language_model.layers.9.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.layers.9.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "language_model.norm.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.0.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.1.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.10.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.11.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.12.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.13.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.14.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.15.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.16.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.2.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.3.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.4.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.5.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.6.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.7.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.8.weight": "model-00004-of-00004.safetensors",
+    "lm_heads.9.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_moss_tts.py

@@ -0,0 +1,568 @@
+# coding=utf-8
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Modeling classes for MossTTSDelay. """
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.cache_utils import Cache
+from transformers.models.qwen3 import Qwen3Model
+from transformers import initialization as init
+
+from .configuration_moss_tts import MossTTSDelayConfig
+from .inference_utils import sample_token, find_last_equal_C
+
+try:
+    from .processing_moss_tts import UserMessage, AssistantMessage, MossTTSDelayProcessor
+except Exception:
+    UserMessage = None
+    AssistantMessage = None
+    MossTTSDelayProcessor = None
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MossTTSDelayConfig"
+
+
+@dataclass
+class MossTTSDelayOutputWithPast(ModelOutput):
+    """
+    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Weighted sum of channel losses.
+        all_sum_losses (`torch.FloatTensor` of shape `(batch_size, n_vq + 1)`, *optional*):
+            Sum of losses for each sample and each channel before averaging.
+        all_token_nums (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Number of non-masked tokens per sample.
+        sample_losses (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            Loss per sample.
+        channel_losses (`torch.FloatTensor` of shape `(n_vq + 1,)`, *optional*):
+            Loss per channel (text head + vq heads).
+        logits (`List[torch.FloatTensor]`, *optional*):
+            List of prediction scores from each head.
+        past_key_values (`Cache`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
+            Tuple of torch.FloatTensor (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer).
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed):
+            Tuple of torch.FloatTensor (one for each layer) of the attention weights.
+    """
+    loss: Optional[torch.FloatTensor] = None
+    all_sum_losses: Optional[torch.FloatTensor] = None
+    all_token_nums: Optional[torch.LongTensor] = None
+    sample_losses: Optional[torch.FloatTensor] = None
+    channel_losses: Optional[torch.FloatTensor] = None
+    logits: Optional[List[torch.FloatTensor]] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+
+
+class MossTTSDelayPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSDelayConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    def _init_weights(self, module):
+        """
+        Transformers 5.0+ safe init:
+        - MUST use transformers.initialization helpers
+        - MUST respect param._is_hf_initialized to avoid overwriting ckpt-loaded params
+        """
+        # Let HF handle its standard modules first (LayerNorm, Linear, Embedding, etc.)
+        super()._init_weights(module)
+
+        # Pick a std consistent with HF conventions
+        # Prefer model/text config initializer_range if present.
+        std = None
+        if hasattr(self.config, "initializer_range"):
+            std = self.config.initializer_range
+        elif hasattr(self.config, "language_config") and hasattr(self.config.language_config, "initializer_range"):
+            std = self.config.language_config.initializer_range
+        else:
+            std = 0.02
+
+        # Initialize extra audio embeddings
+        if isinstance(module, nn.Embedding):
+            # Only touch our extra embeddings (avoid double touching LM's embeddings if not desired)
+            # If you prefer, you can skip this check and rely on super()._init_weights for all embeddings.
+            if getattr(module, "num_embeddings", None) == self.config.audio_vocab_size + 1:
+                init.normal_(module.weight, mean=0.0, std=std)
+                # If you later set padding_idx, you must explicitly zero it (and respect _is_hf_initialized!)
+                # init.zeros_ will internally check param flags, but slicing needs manual care.
+
+        # Initialize multi-head projections you added
+        if isinstance(module, nn.Linear):
+            # For your lm_heads, super()._init_weights already covers typical Linear.
+            # This block is only needed if you have custom Linear variants later.
+            pass
+
+
+
+MOSSTTS_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`MossTTSDelayConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The MossTTSDelay Model architecture tailored for Text-to-Speech generation with multi-head VQ prediction.",
+    MOSSTTS_START_DOCSTRING,
+)
+class MossTTSDelayModel(MossTTSDelayPreTrainedModel):
+    UserMessage = UserMessage
+    AssistantMessage = AssistantMessage
+    Processor = MossTTSDelayProcessor
+
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        self.config = config
+
+        config.language_config.torch_dtype = config.torch_dtype
+        
+        self.language_model = Qwen3Model(config.language_config)
+
+        # Audio VQ Embeddings (Extra channels)
+        # Note: input_ids[..., 0] uses Qwen's embedding. 
+        # input_ids[..., 1:] use these extensions.
+        self.emb_ext = nn.ModuleList()
+        for vq_idx in range(self.config.n_vq):
+            # Add +1 for potential padding/special tokens logic if strictly required by upstream data prep
+            self.emb_ext.append(
+                nn.Embedding(self.config.audio_vocab_size + 1, config.language_config.hidden_size, padding_idx=None)
+            )
+
+        # Multi-Head Prediction Layers
+        # Head 0: Main language head
+        # Head 1..N: Audio VQ heads
+        self.lm_heads = nn.ModuleList([
+            nn.Linear(config.language_config.hidden_size, config.language_config.vocab_size, bias=False)
+        ])
+        for vq_idx in range(self.config.n_vq):
+            self.lm_heads.append(
+                nn.Linear(config.language_config.hidden_size, self.config.audio_vocab_size + 1, bias=False)
+            )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Computes the combined embeddings from text and multiple audio VQ channels.
+        
+        Args:
+            input_ids: Shape (Batch, Seq_Len, 1 + n_vq)
+        """
+        # Base Text/Content Embedding
+        # input_ids[..., 0] is standard text or semantic tokens
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids[..., 0])
+
+        # Add VQ Embeddings
+        for i, embed_layer in enumerate(self.emb_ext):
+            # i corresponds to channel i+1 in input_ids
+            # We assume the data pipeline ensures indices are within range
+            inputs_embeds = inputs_embeds + embed_layer(input_ids[..., i + 1])
+        
+        return inputs_embeds
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        # Returning a list of heads might break some HF utilities expecting a single head.
+        # However, for custom models, this is acceptable.
+        return self.lm_heads
+
+    @add_start_docstrings_to_model_forward(MOSSTTS_START_DOCSTRING)
+    @replace_return_docstrings(output_type=MossTTSDelayOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        hidden_out_layers: Optional[List[int]] = None,
+        channelwise_loss_weight: Optional[List[float]] = None,
+        **kwargs,
+    ) -> Union[Tuple, MossTTSDelayOutputWithPast]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`):
+                Indices of input sequence tokens in the vocabulary. 
+                Dimension 2 contains: [Text/Semantics, VQ_0, VQ_1, ..., VQ_N].
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`, *optional*):
+                Labels for computing the masked language modeling loss.
+            channelwise_loss_weight (`List[float]`, *optional*):
+                Manual weights for summing losses across different heads (Text vs Audio channels).
+
+        Returns:
+        """
+
+        if len(input_ids.shape) != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError("`Input_ids`'s shape should be exactly (batch_size, sequence_length, 1 + n_vq).")
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+        # 1. Prepare Embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+
+        # 2. Backbone Forward
+        # Qwen3Model outputs standard CausalLMOutputWithPast or similar
+        outputs = self.language_model(
+            input_ids=None, # Passed via inputs_embeds
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True, # Always need hidden states for multi-head projection
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # 3. Handle specific layer outputs if requested (Delay Pattern often requires features from specific layers)
+        last_hidden_state = outputs.last_hidden_state
+        if hidden_out_layers is None:
+            # Default to using the last layer for all heads
+            # In some architectures (like MusicGen), different codebooks come from different transformer layers.
+            # Here we default to the final layer as per original code behavior [-1] * (n + 1).
+            hidden_states_for_heads = [last_hidden_state] * (len(self.lm_heads))
+        else:
+            # If hidden_out_layers is provided (e.g. [-1, -2, -3...]), fetch them from all_hidden_states
+            # Note: outputs.hidden_states includes embedding output at index 0 usually.
+            all_hs = outputs.hidden_states
+            hidden_states_for_heads = [all_hs[idx] for idx in hidden_out_layers]
+
+        # 4. Project to Logits (Multi-Head)
+        layer_logits = []
+        for i, (hs, head) in enumerate(zip(hidden_states_for_heads, self.lm_heads)):
+            logits = head(hs)
+            # Original code logic: Mask the last token index for audio heads (indices > 0)
+            # This implies the vocab size is (N+1) but the model shouldn't predict the (N+1)-th token 
+            # (perhaps reserved for padding in the input but invalid for prediction).
+            if i > 0:
+                logits[..., -1] = float("-inf")
+            layer_logits.append(logits)
+
+        # 5. Loss Calculation
+        loss = None
+        all_sum_losses = None
+        all_token_nums = None
+        sample_losses = None
+        channel_losses = None
+
+        if labels is not None:
+            # Ensure labels match input shape rank (B, S, C)
+            if labels.dim() != 3:
+                raise ValueError(f"Labels must have rank 3 (B, S, C), got {labels.shape}")
+
+            batch_size = labels.size(0)
+            n_heads = len(layer_logits)
+            
+            # Container for per-sample, per-channel losses
+            # Shape: [Batch, n_heads]
+            all_sum_losses_list = [] 
+            
+            # Count valid tokens (not -100) per sample. 
+            # Note: Assuming mask is consistent across channels or we take sum over dim 1 (seq)
+            # Usually strict masking means checking one channel or all. 
+            # Original code: torch.sum(labels != -100, dim=1) -> [B, C]
+            all_token_nums = torch.sum(labels != -100, dim=1)
+
+            for i, logits in enumerate(layer_logits):
+                # logits: [B, S, V]
+                # cur_labels: [B, S]
+                cur_labels = labels[..., i]
+                
+                # Flatten for CrossEntropy
+                # logits: [B*S, V], labels: [B*S]
+                loss_fct = CrossEntropyLoss(reduction='none')
+                vocab_size = logits.size(-1)
+                
+                reshaped_logits = logits.view(-1, vocab_size)
+                reshaped_labels = cur_labels.contiguous().view(-1)
+                
+                # Calculate loss per token
+                per_token_loss = loss_fct(reshaped_logits, reshaped_labels)
+                
+                # Reshape back to [B, S] and sum over Sequence dimension to get per-sample loss
+                per_token_loss = per_token_loss.view(batch_size, -1)
+                per_sample_loss = torch.sum(per_token_loss, dim=-1) # [B]
+                
+                all_sum_losses_list.append(per_sample_loss)
+
+            # Stack to [B, n_heads]
+            all_sum_losses = torch.stack(all_sum_losses_list, dim=1)
+
+            # Weighted Loss Aggregation
+            if channelwise_loss_weight is not None:
+                if len(channelwise_loss_weight) != n_heads:
+                    raise ValueError(f"channelwise_loss_weight length {len(channelwise_loss_weight)} != {n_heads}")
+                
+                w_tensor = torch.tensor(channelwise_loss_weight, device=all_sum_losses.device, dtype=all_sum_losses.dtype)
+                
+                # Sample losses: Weighted sum over channels per sample / Total weight
+                # Normalize by token count per channel
+                # Avoid division by zero with epsilon or mask
+                token_counts_safe = all_token_nums.float().clamp(min=1.0)
+                
+                normalized_losses = all_sum_losses / token_counts_safe
+                sample_losses = (normalized_losses * w_tensor).sum(dim=1) / w_tensor.sum()
+                
+                # Channel losses: Sum over batch / Sum tokens over batch
+                total_loss_per_channel = all_sum_losses.sum(dim=0)
+                total_tokens_per_channel = all_token_nums.sum(dim=0).float().clamp(min=1.0)
+                channel_losses = total_loss_per_channel / total_tokens_per_channel
+                
+                # Final scalar loss
+                loss = (channel_losses * w_tensor).sum() / w_tensor.sum()
+            else:
+                # Default average if no weights provided
+                total_tokens = all_token_nums.sum().float().clamp(min=1.0)
+                loss = all_sum_losses.sum() / total_tokens
+                channel_losses = all_sum_losses.sum(dim=0) / all_token_nums.sum(dim=0).clamp(min=1.0)
+
+        return MossTTSDelayOutputWithPast(
+            loss=loss,
+            all_sum_losses=all_sum_losses,
+            all_token_nums=all_token_nums,
+            sample_losses=sample_losses,
+            channel_losses=channel_losses,
+            logits=layer_logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: Optional[int] = None,
+        text_temperature: float = 1.2,
+        text_top_p: float = 0.9,
+        text_top_k: int = 50,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+    ):
+        generation_config = getattr(self, "generation_config", None)
+
+        def _cfg_value(name: str, default_value: Union[int, float]) -> Union[int, float]:
+            if generation_config is None:
+                return default_value
+            value = getattr(generation_config, name, None)
+            if value is None:
+                return default_value
+            return value
+
+        if max_new_tokens is None:
+            try:
+                max_new_tokens = int(_cfg_value("max_new_tokens", 1000))
+            except (TypeError, ValueError):
+                max_new_tokens = 1000
+        if audio_temperature is None:
+            try:
+                audio_temperature = float(_cfg_value("temperature", 1.1))
+            except (TypeError, ValueError):
+                audio_temperature = 1.1
+        if audio_top_p is None:
+            try:
+                audio_top_p = float(_cfg_value("top_p", 0.9))
+            except (TypeError, ValueError):
+                audio_top_p = 0.9
+        if audio_top_k is None:
+            try:
+                audio_top_k = int(_cfg_value("top_k", 50))
+            except (TypeError, ValueError):
+                audio_top_k = 50
+        if audio_repetition_penalty is None:
+            try:
+                audio_repetition_penalty = float(_cfg_value("repetition_penalty", 1.1))
+            except (TypeError, ValueError):
+                audio_repetition_penalty = 1.1
+
+        if text_temperature > 0:
+            text_do_sample = True
+        else:
+            text_temperature = 1
+            text_do_sample = False
+        if audio_temperature > 0:
+            audio_do_sample = True
+        else:
+            audio_temperature = 1
+            audio_do_sample = False
+            
+        past_key_values = None
+        device = input_ids.device
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        batch_size, seq_len, n_vq = input_ids.shape
+        n_vq -= 1
+        
+        generation_ids = input_ids[:]
+        is_stopping = torch.zeros(batch_size, dtype=torch.bool, device=device)
+
+        # 三个阶段: 1. 非 audio; 2. audio not delay; 3. audio delay
+        audio_lengths = torch.zeros(batch_size, dtype=torch.int64, device=device) # 0 的时候表示阶段1;
+        torch_int64_max = torch.iinfo(torch.int64).max
+        delayed_lengths = torch.full((batch_size,), torch_int64_max, dtype=torch.int64, device=device) # 最大值的时候表示阶段2;
+        
+        # 考虑 continuation 时 audio_start 已经在 input_ids 中的情况;
+        # NOTE 注意我们目前不考虑任何输入已经开始 delay 的情况;
+        # 需要同时考虑 continuation 和直接生成的情况;
+        is_continuation = (input_ids[:, -1, 0] == self.config.audio_start_token_id) | (input_ids[:, -1, 0] == self.config.audio_assistant_gen_slot_token_id)
+        audio_start_indices = find_last_equal_C(input_ids[..., 0], self.config.audio_start_token_id)
+        audio_start_mask = is_continuation & (audio_start_indices != -1)
+        audio_lengths[audio_start_mask] = seq_len - audio_start_indices[audio_start_mask]
+        
+        is_audio = audio_start_mask.clone()
+        
+        pre_exclude_mask0 = torch.tensor([self.config.pad_token_id, self.config.audio_assistant_gen_slot_token_id, self.config.audio_assistant_delay_slot_token_id, self.config.audio_end_token_id], device=device)
+        pre_exclude_mask1 = torch.ones(self.config.language_config.vocab_size, device=device).bool()
+        pre_exclude_mask1[[self.config.audio_assistant_gen_slot_token_id, self.config.audio_assistant_delay_slot_token_id]] = False
+        
+        
+        # 注意 time_step 未必表示对于实际对话时，当前输出token的位置，因为有续写的情况;
+        for time_step in tqdm(range(max_new_tokens), desc=f"Generating bs{batch_size} ..."):
+            outputs = self(
+                input_ids=current_input_ids,
+                attention_mask=current_attention_mask,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+            past_key_values = outputs.past_key_values
+            
+            next_token_logits = [logit[:, -1, :] / text_temperature if logit_idx == 0 else logit[:, -1, :] / audio_temperature for logit_idx, logit in enumerate(outputs.logits)] # List, len=n_vq+1, [batch_size, 1, vocab_size]; 
+            next_token_logits[0] = next_token_logits[0].clone()
+            # 1. 先处理 text token;
+            next_text_token = torch.full((batch_size,), self.config.pad_token_id, device=device)
+            # 第二个 audio_assistant_delay_slot_token_id 和 audio_end 是不需要采样的，audio_start, 每一个 audio_assistant_gen_slot_token_ids 和第一个 audio_assistant_delay_slot_token_id 是需要采样的;
+            next_text_token[~is_stopping & (delayed_lengths < n_vq)] = self.config.audio_assistant_delay_slot_token_id
+            is_audio_eos = ~is_stopping & (delayed_lengths == n_vq)
+            next_text_token[is_audio_eos] = self.config.audio_end_token_id
+            is_audio[is_audio_eos] = False
+            sampling_text_mask = ~is_stopping & (delayed_lengths > n_vq)
+            next_token_logits[0][~is_audio] = next_token_logits[0][~is_audio].index_fill(-1, pre_exclude_mask0, float('-inf'))
+            next_token_logits[0][is_audio] = next_token_logits[0][is_audio].masked_fill(pre_exclude_mask1, float('-inf'))
+            if time_step == 0:
+                next_token_logits[0][..., 151662] = float('-inf')
+            if time_step <= n_vq:
+                next_token_logits[0][..., self.config.im_end_token_id] = float('-inf')
+                
+            # 文本层不使用重复惩罚;
+            next_text_token[sampling_text_mask] = sample_token(
+                logits=next_token_logits[0][sampling_text_mask],
+                top_p=text_top_p,
+                top_k=text_top_k,
+                do_sample=text_do_sample
+            )
+            is_audio[next_text_token == self.config.audio_start_token_id] = True
+            # 只存在一种停止逻辑，即 next_text_token = <|im_end|>;
+            is_stopping[next_text_token == self.config.im_end_token_id] = True
+
+            # 2. 再处理 audio tokens;
+            # audio_start 和 audio_end 之外的内容直接pad，默认是 pad，我们只需要填充有值的部分即可;
+            next_audio_tokens = torch.full((batch_size, n_vq), self.config.audio_pad_code, device=device)
+            
+            # 需要考虑的是与 audio_start 的距离;
+            # 先查看是否是pad的情况; true 表示有值;
+            pre_audio_mask = audio_lengths.unsqueeze(1) > torch.arange(n_vq, dtype=int, device=device).expand(batch_size, n_vq)
+            post_audio_mask = torch.arange(n_vq, dtype=int, device=device).expand(batch_size, n_vq) > delayed_lengths.unsqueeze(1) - 1
+            post_audio_mask[delayed_lengths == torch_int64_max] = True
+            sampling_audio_mask = pre_audio_mask & post_audio_mask
+            next_audio_tokens[~sampling_audio_mask] = self.config.audio_pad_code
+            
+            if sampling_audio_mask.sum() > 0:
+                audio_logits = torch.stack(next_token_logits[1:], dim=1)[sampling_audio_mask] # torch.stack -> [batch_size, n_vq - 1, vocab_size]
+                audio_logits[..., self.config.audio_pad_code] = float('-inf')
+                next_audio_tokens[sampling_audio_mask] = sample_token(
+                    logits=audio_logits,
+                    prev_tokens=generation_ids[:, :, 1:],
+                    repetition_penalty=audio_repetition_penalty,
+                    top_p=audio_top_p, 
+                    top_k=audio_top_k, 
+                    do_sample=audio_do_sample
+                )
+                
+            # 这里显示的是下一个时间步时可以直接使用的 audio_lengths 和 delayed_lengths 的状态;
+            # audio_lengths[(next_text_token == self.audio_start_token_id) & (audio_lengths > 0)] += 1
+            # audio_lengths[(next_text_token == self.audio_start_token_id) | (next_text_token == self.audio_assistant_gen_slot_token_id)] += 1
+            audio_lengths[(next_text_token == self.config.audio_start_token_id) | (next_text_token == self.config.audio_assistant_gen_slot_token_id) | (next_text_token == self.config.audio_assistant_delay_slot_token_id)] += 1
+            audio_lengths[next_text_token == self.config.audio_end_token_id] = 0
+            delayed_lengths[(delayed_lengths == torch_int64_max) & (next_text_token == self.config.audio_assistant_delay_slot_token_id)] = 0
+            delayed_lengths[delayed_lengths != torch_int64_max] += 1
+            delayed_lengths[delayed_lengths > n_vq] = torch_int64_max
+            
+            current_input_ids = torch.cat([next_text_token[:, None, None], next_audio_tokens[:, None, :]], dim=2) # [batch_size, 1, n_vq + 1]
+            current_attention_mask = torch.cat([current_attention_mask, (~is_stopping).unsqueeze(-1)], dim=-1)
+            generation_ids = torch.cat([generation_ids, current_input_ids], dim=1) # [batch_size, seq_len, n_vq + 1]
+            
+            if is_stopping.sum() == batch_size:
+                break
+                
+        start_indices = find_last_equal_C(input_ids[..., 0], self.config.im_start_token_id) + 3
+        start_lengths = seq_len - start_indices
+
+        output = []
+        for start_idx, start_length, cur_generation_ids in zip(start_indices, start_lengths, generation_ids):
+            output.append((start_length, cur_generation_ids[start_idx:]))
+        
+        return output

processing_moss_tts.py

@@ -0,0 +1,984 @@
+# coding=utf-8
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal, Final, cast
+from dataclasses import dataclass
+from pathlib import Path
+import re
+import torchaudio
+
+from transformers import processing_utils
+
+if hasattr(processing_utils, "MODALITY_TO_BASE_CLASS_MAPPING"):
+    processing_utils.MODALITY_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = (
+        "PreTrainedModel"
+    )
+elif hasattr(processing_utils, "AUTO_TO_BASE_CLASS_MAPPING"):
+    processing_utils.AUTO_TO_BASE_CLASS_MAPPING["audio_tokenizer"] = (
+        "PreTrainedModel"
+    )
+
+import torch
+from transformers import (
+    PreTrainedTokenizerBase,
+    BatchFeature,
+    ProcessorMixin,
+    logging,
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+)
+
+from .configuration_moss_tts import MossTTSDelayConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+AUDIO_PLACEHOLDER = "<|audio|>"
+
+
+@dataclass
+class Message:
+    def to_dict(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass
+class UserMessage(Message):
+    text: Optional[str] = None
+    reference: Optional[List[Optional[Union[str, torch.Tensor]]]] = None
+    instruction: Optional[str] = None
+    tokens: Optional[int] = None
+    quality: Optional[str] = None
+    sound_event: Optional[str] = None
+    ambient_sound: Optional[str] = None
+    language: Optional[str] = None
+    scene: Optional[str] = None
+
+    def __post_init__(self):
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+None
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- Scene:
+{scene}
+- Text:
+{text}
+</user_inst>"""
+
+        audio_codes_list = []
+        if self.reference is None:
+            reference = "None"
+        elif isinstance(self.reference, List):
+            reference = []
+            for speaker_idx, speaker_reference in enumerate(self.reference):
+                if speaker_reference is None:
+                    reference.append(f"[S{speaker_idx + 1}]: None")
+                else:
+                    reference.append(f"[S{speaker_idx + 1}]:\n{AUDIO_PLACEHOLDER}")
+                    audio_codes_list.append(speaker_reference)
+            reference = "\n".join(reference)
+        else:
+            raise TypeError("`reference` should be exactly a list when it is not None.")
+
+        content = (
+            template.replace("{reference}", str(reference))
+            .replace("{instruction}", str(self.instruction))
+            .replace("{quality}", str(self.quality))
+            .replace("{sound_event}", str(self.sound_event))
+            .replace("{ambient_sound}", str(self.ambient_sound))
+            .replace("{language}", str(self.language))
+            .replace("{scene}", "None")
+            .replace("{text}", str(self.text))
+        )
+
+        self._content = content
+        self._audio_codes_list = audio_codes_list
+
+    def to_dict(self):
+        return {
+            "role": "user",
+            "content": self._content,
+            "audio_codes_list": self._audio_codes_list,
+        }
+
+
+@dataclass
+class AssistantMessage(Message):
+    audio_codes_list: List[Union[str, torch.Tensor]]
+    content: str = AUDIO_PLACEHOLDER
+
+    def to_dict(self):
+        return {
+            "role": "assistant",
+            "content": self.content,
+            "audio_codes_list": self.audio_codes_list,
+        }
+
+
+USER_MESSAGE_FIELDS = (
+    "text",
+    "reference",
+    "instruction",
+    "tokens",
+    "quality",
+    "sound_event",
+    "ambient_sound",
+    "language",
+    "scene",
+)
+
+
+class MossTTSDelayProcessor(ProcessorMixin):
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "AutoModel"
+
+    tokenizer: PreTrainedTokenizerBase
+    audio_tokenizer: Any
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        audio_tokenizer: Any = None,
+        model_config: Optional[MossTTSDelayConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(tokenizer=tokenizer, audio_tokenizer=audio_tokenizer, **kwargs)
+
+        # Explicit assignments for type-checkers; ProcessorMixin sets these too.
+        self.tokenizer = tokenizer
+        self.audio_tokenizer = audio_tokenizer
+        if model_config is None:
+            model_config = MossTTSDelayConfig()
+        self.model_config = model_config
+
+        self.imstart_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
+        self.imend_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        self.newline_token_id = 198
+
+        def _id_to_token(token_id: int) -> str:
+            tok = tokenizer.convert_ids_to_tokens(int(token_id))
+            if isinstance(tok, list):
+                return tok[0] if len(tok) > 0 else ""
+            return cast(str, tok)
+
+        self.audio_user_slot_token = _id_to_token(
+            self.model_config.audio_user_slot_token_id
+        )
+        self.audio_assistant_gen_slot_token = _id_to_token(
+            self.model_config.audio_assistant_gen_slot_token_id
+        )
+        self.audio_assistant_delay_slot_token = _id_to_token(
+            self.model_config.audio_assistant_delay_slot_token_id
+        )
+        self.audio_start_token = _id_to_token(self.model_config.audio_start_token_id)
+        self.audio_end_token = _id_to_token(self.model_config.audio_end_token_id)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        kwargs.pop("_from_auto", None)
+
+        audio_tokenizer_name_or_path = kwargs.pop(
+            "codec_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"
+        )
+
+        pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+        model_config = cast(
+            MossTTSDelayConfig,
+            AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                *args,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            ),
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            *args,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        audio_tokenizer = AutoModel.from_pretrained(
+            audio_tokenizer_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+
+        return cls(
+            tokenizer=tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            model_config=model_config,
+            **kwargs,
+        )
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        conversations = args[0] if len(args) > 0 else kwargs.pop("conversations")
+        mode: str = kwargs.pop("mode", "generation")
+        apply_chat_template: bool = kwargs.pop("apply_chat_template", True)
+        n_vq: Optional[int] = kwargs.pop("n_vq", None)
+
+        # Common ProcessorMixin kwargs that we ignore because we always return torch tensors.
+        kwargs.pop("return_tensors", None)
+        kwargs.pop("padding", None)
+        kwargs.pop("truncation", None)
+
+        """
+        mode only works when a Message is converted to a dict.
+        """
+
+        if mode not in {"generation", "continuation"}:
+            raise RuntimeError
+
+        if isinstance(conversations, (Message, Dict)):
+            conversations = [conversations]
+
+        truncation = False
+        if mode == "continuation":
+            truncation = True
+
+        input_ids_list = []
+        for conversation in conversations:
+            if isinstance(conversation, (Message, Dict)):
+                conversation = [conversation]
+
+            # Normalize early so downstream logic always deals with dict messages.
+            conversation = [self._normalize_message(m) for m in conversation]
+
+            if (mode == "generation") ^ (len(conversation) % 2 != 0):
+                raise ValueError
+
+            if (mode == "generation") ^ (conversation[-1]["role"] == "user"):
+                raise ValueError
+
+            unified_codes = []
+            for message_idx, message in enumerate(conversation):
+                if apply_chat_template:
+                    add_generation_prompt = (
+                        mode == "generation" and message_idx == len(conversation) - 1
+                    )
+                    try:
+                        content = self.tokenizer.apply_chat_template(
+                            [{"role": message["role"], "content": message["content"]}],
+                            add_generation_prompt=add_generation_prompt,
+                            tokenize=False,
+                        )
+                    except TypeError:
+                        try:
+                            content = self.tokenizer.apply_chat_template(
+                                [
+                                    {
+                                        "role": message["role"],
+                                        "content": message["content"],
+                                    }
+                                ],
+                                add_generation_prompt=add_generation_prompt,
+                            )
+                        except Exception:
+                            logger.warning(
+                                "apply_chat_template failed; fallback to raw content."
+                            )
+                            content = message["content"]
+                else:
+                    content = message["content"]
+
+                if not isinstance(content, str):
+                    content = str(content)
+
+                # Batch-encode all path-based references in one call when possible.
+                # This ensures we actually exercise audio_tokenizer.batch_encode for multi-reference prompts,
+                # instead of repeatedly calling it with batch=1.
+                raw_audio_items = message.get("audio_codes_list", [])
+
+                audio_codes_list: List[torch.Tensor] = []
+                if len(raw_audio_items) > 0:
+                    encoded_items: List[Optional[torch.Tensor]] = [None] * len(
+                        raw_audio_items
+                    )
+                    paths: List[str] = []
+                    path_positions: List[int] = []
+
+                    for idx, item in enumerate(raw_audio_items):
+                        if isinstance(item, torch.Tensor):
+                            if n_vq is not None and item.shape[1] != n_vq:
+                                raise RuntimeError(
+                                    "audio_codes's n_vq is not equal to the parameter `n_vq`. Your can set the parameter `n_vq` as None if you have already tokenzied the wavs."
+                                )
+                            encoded_items[idx] = item
+                            continue
+
+                        if isinstance(item, (str, os.PathLike)):
+                            paths.append(str(item))
+                            path_positions.append(idx)
+                            continue
+
+                        raise TypeError(
+                            "Each audio item must be a torch.Tensor of codes or a path-like string."
+                        )
+
+                    if len(paths) > 0:
+                        encoded_from_paths = self.encode_audios_from_path(paths, n_vq)
+                        if len(encoded_from_paths) != len(paths):
+                            raise RuntimeError(
+                                "encode_audios_from_path returned an unexpected number of items."
+                            )
+                        for pos, codes in zip(path_positions, encoded_from_paths):
+                            encoded_items[pos] = codes
+
+                    audio_codes_list = [cast(torch.Tensor, t) for t in encoded_items]
+                unified_codes.append(
+                    self._get_unified_codes(
+                        message["role"], content, audio_codes_list, truncation
+                    )
+                )
+
+            unified_codes = torch.cat(unified_codes)
+            input_ids_list.append(unified_codes)
+
+        return BatchFeature(data=self._pad(input_ids_list))
+
+    @staticmethod
+    def build_user_message(
+        text: Optional[str] = None,
+        reference: Optional[List[Optional[Union[str, torch.Tensor]]]] = None,
+        instruction: Optional[str] = None,
+        tokens: Optional[int] = None,
+        quality: Optional[str] = None,
+        sound_event: Optional[str] = None,
+        ambient_sound: Optional[str] = None,
+        language: Optional[str] = None,
+        scene: Optional[str] = None,
+    ) -> Dict:
+        if reference is not None and not isinstance(reference, list):
+            reference = [reference]
+        return UserMessage(
+            text=text,
+            reference=reference,
+            instruction=instruction,
+            tokens=tokens,
+            quality=quality,
+            sound_event=sound_event,
+            ambient_sound=ambient_sound,
+            language=language,
+            scene=scene,
+        ).to_dict()
+
+    @staticmethod
+    def build_assistant_message(
+        audio_codes_list: List[Union[str, torch.Tensor]],
+        content: str = AUDIO_PLACEHOLDER,
+    ) -> Dict:
+        return AssistantMessage(
+            audio_codes_list=audio_codes_list,
+            content=content,
+        ).to_dict()
+
+    def _normalize_message(self, message: Union[Message, Dict]) -> Dict:
+        if isinstance(message, Message):
+            return message.to_dict()
+        if not isinstance(message, dict):
+            raise TypeError("Each message must be a Message or dict.")
+        if "role" not in message:
+            raise ValueError("Message dict must include a 'role' field.")
+        if "content" in message and "audio_codes_list" in message:
+            return message
+        role = message["role"]
+        if role == "user":
+            kwargs = {key: message.get(key) for key in USER_MESSAGE_FIELDS}
+            return self.build_user_message(**kwargs)
+        if role == "assistant":
+            return self.build_assistant_message(
+                audio_codes_list=message.get("audio_codes_list", []),
+                content=message.get("content", AUDIO_PLACEHOLDER),
+            )
+        raise ValueError(f"Unsupported role: {role}")
+
+    def _pad(self, input_ids_list: List[torch.Tensor]):
+        device = input_ids_list[0].device
+        lengths = torch.tensor([w.shape[0] for w in input_ids_list], device=device)
+        pad_input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids_list,
+            batch_first=True,
+            padding_value=self.model_config.audio_pad_code,
+            padding_side="left",
+        )
+        other_channel_mask = (pad_input_ids.shape[1] - lengths).unsqueeze(
+            1
+        ) > torch.arange(pad_input_ids.shape[1], device=device).unsqueeze(0)
+        pad_input_ids[..., 0][other_channel_mask] = self.model_config.pad_token_id
+        attention_mask = torch.zeros(
+            pad_input_ids.shape[0], pad_input_ids.shape[1], device=device
+        )
+        attention_mask[~other_channel_mask] = 1
+        attention_mask = attention_mask.bool()
+        return {
+            "input_ids": pad_input_ids,  # [batch_size, seqlen, n_vq]
+            "attention_mask": attention_mask,
+        }
+
+    @staticmethod
+    def _replace_audio_placeholders(
+        content: str,
+        lengths: List[int],
+        n_vq: int,
+        gen_slot_token: str,
+        delay_slot_token: str,
+        audio_start_token: str,
+        audio_end_token: str,
+    ) -> str:
+        if n_vq < 1:
+            raise ValueError(f"n_vq must be >= 1, got {n_vq}")
+
+        num_placeholders = content.count(AUDIO_PLACEHOLDER)
+        if num_placeholders != len(lengths):
+            raise ValueError(
+                f"Number of {AUDIO_PLACEHOLDER} ({num_placeholders}) "
+                f"does not match lengths ({len(lengths)})"
+            )
+
+        def build_audio_block(length: int) -> str:
+            if length < 0:
+                raise ValueError(f"length must be >= 0, got {length}")
+
+            if length == 0:
+                return f"{audio_start_token}{audio_end_token}"
+
+            step_tokens = gen_slot_token * length + (delay_slot_token * (n_vq - 1))
+            return f"{audio_start_token}{step_tokens}{audio_end_token}"
+
+        lengths_iter = iter(lengths)
+
+        def replacer(match: re.Match) -> str:
+            length = next(lengths_iter)
+            return build_audio_block(length)
+
+        result = re.sub(re.escape(AUDIO_PLACEHOLDER), replacer, content)
+
+        return result
+
+    @staticmethod
+    def _merge_consecutive_audio_placeholders(
+        content: str,
+        audio_codes_list: List[torch.Tensor],
+    ) -> Tuple[str, List[torch.Tensor]]:
+        matches = list(re.finditer(re.escape(AUDIO_PLACEHOLDER), content))
+        if len(matches) <= 1:
+            return content, audio_codes_list
+
+        if len(matches) != len(audio_codes_list):
+            raise ValueError(
+                "Audio placeholders do not match the provided audio codes list."
+            )
+
+        new_audio_codes_list = []
+        new_parts = []
+        last_pos = 0
+        i = 0
+        while i < len(matches):
+            j = i
+            while (
+                j + 1 < len(matches)
+                and content[matches[j].end() : matches[j + 1].start()].strip() == ""
+            ):
+                j += 1
+
+            new_parts.append(content[last_pos : matches[i].start()])
+            new_parts.append(AUDIO_PLACEHOLDER)
+            last_pos = matches[j].end()
+
+            if j == i:
+                new_audio_codes_list.append(audio_codes_list[i])
+            else:
+                new_audio_codes_list.append(
+                    torch.cat(audio_codes_list[i : j + 1], dim=0)
+                )
+
+            i = j + 1
+
+        new_parts.append(content[last_pos:])
+        return "".join(new_parts), new_audio_codes_list
+
+    @staticmethod
+    def apply_delay_pattern(codes: torch.Tensor, pad_code: int) -> torch.Tensor:
+        delayed_tokens = torch.full(
+            (codes.shape[0] + codes.shape[1] - 1, codes.shape[1]),
+            pad_code,
+            device=codes.device,
+            dtype=codes.dtype,
+        )
+        for i in range(codes.shape[1]):
+            delayed_tokens[i : i + codes.shape[0], i] = codes[:, i]
+        return delayed_tokens
+
+    @staticmethod
+    def apply_de_delay_pattern(delay_codes: torch.Tensor) -> torch.Tensor:
+        tokens = torch.full(
+            (delay_codes.shape[0] - delay_codes.shape[1] + 1, delay_codes.shape[1]),
+            0,
+            device=delay_codes.device,
+            dtype=delay_codes.dtype,
+        )
+        for i in range(delay_codes.shape[1]):
+            tokens[:, i] = delay_codes[i : i + tokens.shape[0], i]
+        return tokens
+
+    def _get_unified_codes(
+        self,
+        role: str,
+        content: str,
+        audio_codes_list: List[torch.Tensor],
+        truncation: bool,
+    ) -> torch.Tensor:
+        """
+        此时的 content 已经是带上了对话格式
+        """
+        if role == "user":
+            audio_gen_slot_token = audio_delay_slot_token = self.audio_user_slot_token
+            truncation= False
+        else:
+            audio_gen_slot_token = self.audio_assistant_gen_slot_token
+            audio_delay_slot_token = self.audio_assistant_delay_slot_token
+
+        # Always follow model RVQ channels (e.g. n_vq=16) and truncate
+        # tokenizer outputs (e.g. 32-layer RVQ) to the first n_vq layers.
+        n_vq = self.model_config.n_vq
+        normalized_audio_codes_list: List[torch.Tensor] = []
+        for audio_codes in audio_codes_list:
+            if audio_codes.dim() != 2:
+                raise RuntimeError(
+                    f"Expect audio codes with rank 2, got {tuple(audio_codes.shape)}"
+                )
+
+            # Handle possible [NQ, T] layout.
+            if audio_codes.shape[1] < n_vq and audio_codes.shape[0] >= n_vq:
+                audio_codes = audio_codes.transpose(0, 1)
+
+            if audio_codes.shape[1] < n_vq:
+                raise RuntimeError(
+                    f"audio_codes channels ({audio_codes.shape[1]}) < model n_vq ({n_vq})"
+                )
+
+            audio_codes = audio_codes[:, :n_vq]
+            normalized_audio_codes_list.append(audio_codes)
+
+        audio_codes_list = normalized_audio_codes_list
+
+        if len(audio_codes_list) > 1 and AUDIO_PLACEHOLDER in content:
+            content, audio_codes_list = self._merge_consecutive_audio_placeholders(
+                content, audio_codes_list
+            )
+        content = self._replace_audio_placeholders(
+            content=content,
+            lengths=[len(audio_codes) for audio_codes in audio_codes_list],
+            n_vq=n_vq,
+            gen_slot_token=audio_gen_slot_token,
+            delay_slot_token=audio_delay_slot_token,
+            audio_start_token=self.audio_start_token,
+            audio_end_token=self.audio_end_token,
+        )
+        text_codes = torch.tensor(
+            self.tokenizer.encode(content),
+            device=audio_codes_list[0].device if audio_codes_list else None,
+        )
+
+        audio_start_indices = torch.where(
+            text_codes == self.model_config.audio_start_token_id
+        )[0]
+        audio_end_indices = torch.where(
+            text_codes == self.model_config.audio_end_token_id
+        )[0]
+        if len(audio_start_indices) != len(audio_codes_list) or len(
+            audio_end_indices
+        ) != len(audio_codes_list):
+            raise ValueError(
+                "Audio placeholders do not match the provided audio codes list."
+            )
+
+        delay_audio_codes_list = []
+        if len(audio_codes_list) == 0:
+            delay_audio_codes_list = torch.full(
+                (len(text_codes), n_vq),
+                self.model_config.audio_pad_code,
+                device=text_codes.device,
+                dtype=text_codes.dtype,
+            )
+        else:
+            prefix_idx = 0
+            for audio_start_idx_t, audio_end_idx_t, audio_codes in zip(
+                audio_start_indices, audio_end_indices, audio_codes_list
+            ):
+                audio_start_idx = int(audio_start_idx_t.item())
+                audio_end_idx = int(audio_end_idx_t.item())
+                delay_audio_codes = self.apply_delay_pattern(
+                    audio_codes, self.model_config.audio_pad_code
+                )
+                pad_codes = torch.full(
+                    (audio_start_idx - prefix_idx + 1, n_vq),
+                    self.model_config.audio_pad_code,
+                    device=audio_codes.device,
+                    dtype=audio_codes.dtype,
+                )
+                delay_audio_codes_list.extend([pad_codes, delay_audio_codes])
+                prefix_idx = audio_end_idx
+
+            if truncation:
+                delay_audio_codes_list[-1] = delay_audio_codes_list[-1][
+                    : -(n_vq - 1), :
+                ]
+            else:
+                last_audio_end_idx = int(audio_end_indices[-1].item())
+                pad_codes = torch.full(
+                    (len(text_codes) - last_audio_end_idx, n_vq),
+                    self.model_config.audio_pad_code,
+                    device=audio_codes_list[0].device,
+                    dtype=audio_codes_list[0].dtype,
+                )
+                delay_audio_codes_list.append(pad_codes)
+
+            delay_audio_codes_list = torch.cat(delay_audio_codes_list)
+
+        if text_codes.shape[0] != delay_audio_codes_list.shape[0]:
+            text_codes = text_codes[: delay_audio_codes_list.shape[0]]
+
+        unified_codes = torch.cat(
+            [text_codes.unsqueeze(1), delay_audio_codes_list], dim=1
+        )
+        return unified_codes
+
+    def _parse_text_codes(self, start_length, text_codes):
+        text = cast(str, self.tokenizer.decode(text_codes))
+        prefix = cast(str, self.tokenizer.decode(text_codes[:start_length]))
+        text = text[len(prefix) :]
+
+        AUDIO_PATTERN = re.compile(
+            rf"(?:{self.audio_start_token})?"
+            rf"(?:{self.audio_assistant_gen_slot_token})*"
+            rf"(?:{self.audio_assistant_delay_slot_token})*"
+            rf"{self.audio_end_token}"
+        )
+
+        def normalize_audio_segments(text: str) -> str:
+            def repl(match: re.Match) -> str:
+                seg = match.group(0)
+                # Replace with <|audio|> if gen_slot is present in the segment;
+                if self.audio_assistant_gen_slot_token in seg:
+                    return AUDIO_PLACEHOLDER
+                # Otherwise, remove it.
+                return ""
+
+            return AUDIO_PATTERN.sub(repl, text)
+
+        return normalize_audio_segments(text)
+
+    def _parse_audio_codes(self, start_length, audio_codes):
+        # De-delay back to [T', n_vq]
+        audio_codes = self.apply_de_delay_pattern(audio_codes)
+
+        # Rows that are all pad are separators between real audio segments.
+        is_pad = (audio_codes == self.model_config.audio_pad_code).all(dim=1)
+        non_pad = ~is_pad
+        if not non_pad.any():
+            return []
+
+        idx = torch.nonzero(non_pad).squeeze(1)
+        breaks = torch.where(idx[1:] != idx[:-1] + 1)[0] + 1
+        if breaks.numel() == 0:
+            segments_idx = [idx]
+        else:
+            segments_idx = torch.split(idx, breaks.tolist())
+
+        audio_codes_list = [audio_codes[s] for s in segments_idx]
+
+        # Batch-decode all audio segments together.
+        decoded_audio_list = self.decode_audio_codes(audio_codes_list)
+
+        # Keep codec causal context by decoding the whole first segment first,
+        # then trim at waveform level according to start_length ratio.
+        if (
+            start_length > 0
+            and len(audio_codes_list) > 0
+            and len(decoded_audio_list) > 0
+        ):
+            first_codes_length = audio_codes_list[0].shape[0]
+            if first_codes_length > 0:
+                trim_ratio = max(
+                    0.0, min(float(start_length) / float(first_codes_length), 1.0)
+                )
+                first_audio = decoded_audio_list[0]
+                if trim_ratio >= 1.0:
+                    decoded_audio_list = decoded_audio_list[1:]
+                elif trim_ratio > 0.0:
+                    trim_samples = int(first_audio.shape[-1] * trim_ratio)
+                    decoded_audio_list[0] = first_audio[..., trim_samples:]
+
+        return decoded_audio_list
+
+    def decode(self, output: List[Tuple[int, torch.Tensor]]):
+        """
+        1. 这里不管怎样，都需要一个完整的 assistant generation ids;
+        2. 支持从任意位置进行截断；
+        """
+
+        genearted_messages = []
+        for start_length, generation_ids in output:
+            content = self._parse_text_codes(start_length, generation_ids[:, 0])
+            audio_codes_list = self._parse_audio_codes(
+                start_length, generation_ids[:, 1:]
+            )
+            if content == "":
+                message = None
+            else:
+                message = AssistantMessage(
+                    content=content,
+                    audio_codes_list=cast(
+                        List[Union[str, torch.Tensor]], audio_codes_list
+                    ),
+                )
+            genearted_messages.append(message)
+        return genearted_messages
+
+    @staticmethod
+    def loudness_normalize(
+        wav: torch.Tensor,
+        target_dbfs: float = -20,
+        gain_range: tuple[float, float] = (-3.0, 3.0),
+    ) -> torch.Tensor:
+        wav = wav.to(torch.float32)
+        if wav.numel() == 0:
+            return wav
+        current_dbfs = 10.0 * torch.log10(torch.mean(wav**2) + 1e-9)
+        gain = float(target_dbfs - current_dbfs)
+        gain = max(gain_range[0], min(gain, gain_range[1]))
+        factor = 10.0 ** (gain / 20.0)
+        return wav * factor
+
+    def _get_audio_tokenizer_device(self) -> torch.device:
+        """Best-effort device inference for `self.audio_tokenizer`.
+
+        Notes:
+        - Old TAC wrapper exposed `.device`, but standard `torch.nn.Module` does not.
+        - New MossAudioTokenizerModel is a `PreTrainedModel`; parameters define its device.
+        """
+
+        audio_tokenizer = getattr(self, "audio_tokenizer", None)
+        if audio_tokenizer is None:
+            logger.warning(
+                "audio_tokenizer is not set on processor. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+        device_attr = getattr(audio_tokenizer, "device", None)
+        if isinstance(device_attr, torch.device):
+            return device_attr
+
+        try:
+            return next(audio_tokenizer.parameters()).device
+        except StopIteration:
+            # No parameters (shouldn't happen for real models); default to CPU.
+            logger.warning(
+                "No parameters found on audio_tokenizer. Using CPU as default."
+            )
+            return torch.device("cpu")
+
+    def encode_audios_from_wav(
+        self,
+        wav_list: List[torch.Tensor],
+        sampling_rate: int,
+        n_vq: Optional[int] = None,
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
+        if isinstance(wav_list, torch.Tensor):
+            wav_list = [wav_list]
+        wav_list_ = []
+        resample = False
+        if sampling_rate != self.model_config.sampling_rate:
+            resample = True
+        device = self._get_audio_tokenizer_device()
+        for wav in wav_list:
+            if wav.shape[0] > 1:
+                wav = torch.mean(wav, dim=0, keepdim=True)
+            if resample:
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=sampling_rate,
+                    new_freq=self.model_config.sampling_rate,
+                )
+            wav = wav.to(device)
+            wav_list_.append(self.loudness_normalize(wav.squeeze(0)))
+
+        # New MossAudioTokenizerModel API: prefer batch_encode(list[wav])
+        if hasattr(audio_tokenizer, "batch_encode"):
+            enc = audio_tokenizer.batch_encode(wav_list_, num_quantizers=n_vq)
+            audio_codes = enc.audio_codes  # (NQ, B, T)
+            audio_codes_lengths = enc.audio_codes_lengths  # (B,)
+        else:
+            # Fallback: use encode() with explicit padding.
+            max_len = max(int(wav.shape[-1]) for wav in wav_list_)
+            input_values = torch.zeros(
+                len(wav_list_), 1, max_len, device=device, dtype=torch.float32
+            )
+            padding_mask = torch.zeros(
+                len(wav_list_), max_len, device=device, dtype=torch.bool
+            )
+            for i, wav in enumerate(wav_list_):
+                this_len = int(wav.shape[-1])
+                input_values[i, 0, :this_len] = wav
+                padding_mask[i, :this_len] = True
+            enc = audio_tokenizer.encode(
+                input_values,
+                padding_mask=padding_mask,
+                num_quantizers=n_vq,
+                return_dict=True,
+            )
+            audio_codes = enc.audio_codes
+            audio_codes_lengths = enc.audio_codes_lengths
+
+        if audio_codes is None or audio_codes_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.encode() returned empty outputs (audio_codes/audio_codes_lengths)."
+            )
+
+        # Keep processor's historical contract: list[Tensor] with shape (T, NQ)
+        # and on CPU (so downstream text/audio packing remains device-agnostic).
+        codes_list: List[torch.Tensor] = []
+        for i in range(int(audio_codes.shape[1])):
+            length_i = int(audio_codes_lengths[i].item())
+            codes_i = (
+                audio_codes[:, i, :length_i]
+                .transpose(0, 1)
+                .contiguous()
+                .to(torch.long)
+                .cpu()
+            )
+            codes_list.append(codes_i)
+        return codes_list
+
+    def encode_audios_from_path(
+        self, wav_path_list: Union[str, List[str]], n_vq: Optional[int] = None
+    ):
+        if isinstance(wav_path_list, str):
+            wav_path_list = [wav_path_list]
+
+        if len(wav_path_list) == 0:
+            raise ValueError("Empty wav_path_list")
+
+        # Load + (if needed) resample each wav independently, so callers can
+        # pass a heterogeneous batch of files while still benefiting from
+        # audio_tokenizer.batch_encode.
+        target_sr = int(self.model_config.sampling_rate)
+        wav_list: List[torch.Tensor] = []
+        for wav_path in wav_path_list:
+            wav, sr = torchaudio.load(wav_path)
+            if int(sr) != target_sr:
+                wav = torchaudio.functional.resample(
+                    waveform=wav,
+                    orig_freq=int(sr),
+                    new_freq=target_sr,
+                )
+            wav_list.append(wav)
+
+        return self.encode_audios_from_wav(wav_list, target_sr, n_vq)
+
+    def decode_audio_codes(
+        self, audio_tokens_list: Union[torch.Tensor, List[torch.Tensor]]
+    ):
+        if self.audio_tokenizer is None:
+            raise RuntimeError("audio_tokenizer is not set on processor.")
+        audio_tokenizer = self.audio_tokenizer
+
+        if isinstance(audio_tokens_list, torch.Tensor):
+            audio_tokens_list = [audio_tokens_list]
+        if len(audio_tokens_list) == 0:
+            return []
+
+        device = self._get_audio_tokenizer_device()
+
+        # Processor uses (T, NQ); MossAudioTokenizer expects (NQ, T) (or (NQ, B, T)).
+        codes_list = [
+            codes.transpose(0, 1).contiguous().to(device=device, dtype=torch.long)
+            for codes in audio_tokens_list
+        ]
+
+        # Align with legacy behavior: decode each sample with chunk_duration=8.0.
+        # Streaming chunk decode currently supports batch_size=1 in MossAudioTokenizer.
+        if hasattr(audio_tokenizer, "decode"):
+            wav_list: List[torch.Tensor] = []
+            for codes in codes_list:
+                try:
+                    dec = audio_tokenizer.decode(
+                        codes,
+                        return_dict=True,
+                        chunk_duration=8.0,
+                    )
+                except TypeError:
+                    # Compatibility fallback for tokenizers without chunk_duration arg.
+                    dec = audio_tokenizer.decode(
+                        codes,
+                        return_dict=True,
+                    )
+
+                audio = dec.audio
+                audio_lengths = dec.audio_lengths
+                if audio is None:
+                    raise RuntimeError("audio_tokenizer.decode() returned empty audio.")
+
+                if audio_lengths is None:
+                    cur_len = int(audio.shape[-1])
+                else:
+                    cur_len = int(audio_lengths[0].item())
+
+                if audio.ndim == 3:
+                    wav = audio[0, 0, :cur_len]
+                elif audio.ndim == 2:
+                    wav = audio[0, :cur_len]
+                else:
+                    raise RuntimeError(
+                        f"Unexpected audio shape from decode: {tuple(audio.shape)}"
+                    )
+                wav_list.append(wav.contiguous().to(torch.float32).cpu())
+            return wav_list
+
+        if hasattr(audio_tokenizer, "batch_decode"):
+            dec = audio_tokenizer.batch_decode(codes_list)
+            audio = dec.audio
+            audio_lengths = dec.audio_lengths
+            if audio is None or audio_lengths is None:
+                raise RuntimeError(
+                    "audio_tokenizer.batch_decode() returned empty outputs (audio/audio_lengths)."
+                )
+            wav_list = []
+            for i in range(int(audio.shape[0])):
+                length_i = int(audio_lengths[i].item())
+                wav = audio[i, 0, :length_i].contiguous().to(torch.float32).cpu()
+                wav_list.append(wav)
+            return wav_list
+
+        raise RuntimeError("audio_tokenizer has neither decode() nor batch_decode().")

processor_config.json

@@ -0,0 +1,6 @@
+{
+    "processor_class": "MossTTSDelayProcessor",
+    "auto_map": {
+        "AutoProcessor": "processing_moss_tts.MossTTSDelayProcessor"
+    }
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_user_slot|>",
+    "<|image_pad|>",
+    "<|audio_assistant_gen_slot|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|audio_user_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|audio_assistant_gen_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|audio_assistant_delay_slot|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_user_slot|>",
+    "<|image_pad|>",
+    "<|audio_assistant_gen_slot|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "AsteroidProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}