Upload folder using huggingface_hub

.gitattributes

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

README.md

@@ -0,0 +1,212 @@
+---
+license: apache-2.0
+tags:
+- text-to-speech
+- gptq
+- quantized
+language:
+- zh
+- en
+- de
+- es
+- fr
+- ja
+- it
+- he
+- ko
+- ru
+- fa
+- ar
+- pl
+- pt
+- cs
+- da
+- sv
+- hu
+- el
+- tr
+base_model: OpenMOSS-Team/MOSS-TTS
+---
+
+# MOSS-TTS-GPTQ
+
+This is a **GPTQ 4-bit quantized** version of [MOSS-TTS (MossTTSDelay 8B)](https://huggingface.co/OpenMOSS-Team/MOSS-TTS) by the OpenMOSS Team, quantized using [GPTQModel](https://github.com/ModelCloud/GPTQModel). The quantization targets the LLM backbone only; the non-backbone weights (`emb_ext`, `lm_heads`, `language_model.norm`) are preserved in their original precision and merged back into a single self-contained repo.
+
+This repo is intended to be used together with the float16 audio tokenizer: 🤗 [blazingbhavneek/MOSS-Audio-Tokenizer-FP16](https://huggingface.co/blazingbhavneek/MOSS-Audio-Tokenizer-FP16).
+
+**VRAM usage: ~9 GB** (vs. ~23 GB for the original bf16 model + tokenizer), making this accessible on a single consumer GPU.
+
+The quantization scripts and a working environment `requirements.txt` are available in the [fork used to produce this model](https://github.com/blazingbhavneek/MOSS-TTS/tree/feature/gptq-quant-model-support).
+
+**Original model resources:**
+- 🤗 [OpenMOSS-Team/MOSS-TTS](https://huggingface.co/OpenMOSS-Team/MOSS-TTS)
+- 💻 [OpenMOSS/MOSS-TTS (GitHub)](https://github.com/OpenMOSS/MOSS-TTS)
+
+---
+
+## Model Description
+
+MOSS-TTS is a production-grade TTS foundation model focused on **zero-shot voice cloning**, **ultra-long stable speech generation** (up to 1 hour), **token-level duration control**, **phoneme-level pronunciation control** (Pinyin / IPA), and **multilingual & code-switched synthesis**. It is built on a clean autoregressive discrete-token recipe with a large-scale Transformer backbone (MossTTSDelay architecture).
+
+### Supported Languages
+
+20 languages are supported:
+
+| Language | Code | | Language | Code | | Language | Code |
+|---|---|---|---|---|---|---|---|
+| Chinese | zh | 🇨🇳 | English | en | 🇺🇸 | German | de | 🇩🇪 |
+| Spanish | es | 🇪🇸 | French | fr | 🇫🇷 | Japanese | ja | 🇯🇵 |
+| Italian | it | 🇮🇹 | Hebrew | he | 🇮🇱 | Korean | ko | 🇰🇷 |
+| Russian | ru | 🇷🇺 | Persian (Farsi) | fa | 🇮🇷 | Arabic | ar | 🇸🇦 |
+| Polish | pl | 🇵🇱 | Portuguese | pt | 🇵🇹 | Czech | cs | 🇨🇿 |
+| Danish | da | 🇩🇰 | Swedish | sv | 🇸🇪 | Hungarian | hu | 🇭🇺 |
+| Greek | el | 🇬🇷 | Turkish | tr | 🇹🇷 | | | |
+
+---
+
+## Quantization Details
+
+| Property | Value |
+|---|---|
+| Base model | MossTTSDelay-8B |
+| Quantization method | GPTQ |
+| Bits | 4 |
+| Group size | 128 |
+| Calibration dataset | wikitext2 |
+| Quantization library | [GPTQModel](https://github.com/ModelCloud/GPTQModel) |
+| VRAM (this model + FP16 tokenizer) | ~9 GB |
+| VRAM (original bf16 model + tokenizer) | ~23 GB |
+
+---
+
+## Evaluation — Quality Degradation
+
+Benchmarks on [seed-tts-eval](https://github.com/BytedanceSpeech/seed-tts-eval). Lower WER is better.
+
+| Model | EN WER (%) ↓ |
+|---|---:|
+| MossTTSDelay-8B (original fp32/bf16) | 1.79 |
+| **MossTTSDelay-8B (this GPTQ 4-bit)** | 2.585 |
+
+---
+
+## Installation
+
+Install [GPTQModel](https://github.com/ModelCloud/GPTQModel) in addition to the standard MOSS-TTS dependencies — it is required to load the quantized backbone at inference time:
+
+```bash
+pip install gptqmodel
+```
+
+Then install MOSS-TTS from the fork that added GPTQ support:
+
+```bash
+git clone -b feature/gptq-quant-model-support https://github.com/blazingbhavneek/MOSS-TTS.git
+cd MOSS-TTS
+pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e .
+```
+
+> **Note:** When loading this model you will see warnings about unexpected keys (e.g. `model.layers.X.mlp.gate_proj.qweight`). These are normal — they come from the GPTQModel library's quantized layer format and do not affect inference.
+
+---
+
+## Inference
+
+```python
+from pathlib import Path
+import torch
+import torchaudio
+from moss_tts_delay.modeling_moss_tts import MossTTSDelayModel
+from moss_tts_delay.processing_moss_tts import MossTTSDelayProcessor
+
+# Disable problematic SDP backends
+torch.backends.cuda.enable_cudnn_sdp(False)
+torch.backends.cuda.enable_flash_sdp(False)
+torch.backends.cuda.enable_mem_efficient_sdp(False)
+torch.backends.cuda.enable_math_sdp(True)
+
+MERGED_PATH = "blazingbhavneek/MOSS-TTS-GPTQ"
+AUDIO_TOK_PATH = "blazingbhavneek/MOSS-Audio-Tokenizer-FP16"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = MossTTSDelayModel.from_pretrained(
+    MERGED_PATH,
+    gptq_device=device,
+    trust_remote_code=True,
+).eval()
+
+processor = MossTTSDelayProcessor.from_pretrained(
+    MERGED_PATH,
+    codec_path=AUDIO_TOK_PATH,
+    trust_remote_code=True,
+)
+
+# --- Example texts ---
+text_en = "We stand on the threshold of the AI era. Artificial intelligence is no longer just a concept in laboratories, but is entering every industry, every creative endeavor, and every decision."
+text_zh = "亲爱的你，你好呀。今天，我想用最认真、最温柔的声音，对你说一些重要的话。"
+text_pinyin = "nin2 hao3，qing3 wen4 nin2 lai2 zi4 na3 zuo4 cheng2 shi4？"
+text_ipa = "/həloʊ, meɪ aɪ æsk wɪtʃ sɪti juː ɑːr frʌm?/"
+
+# Reference audio for voice cloning (URLs or local paths)
+ref_audio_zh = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_zh.wav"
+ref_audio_en = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_en.m4a"
+
+conversations = [
+    # Direct TTS (no reference)
+    [processor.build_user_message(text=text_zh)],
+    [processor.build_user_message(text=text_en)],
+    # Pronunciation control
+    [processor.build_user_message(text=text_pinyin)],
+    [processor.build_user_message(text=text_ipa)],
+    # Voice cloning
+    [processor.build_user_message(text=text_zh, reference=[ref_audio_zh])],
+    [processor.build_user_message(text=text_en, reference=[ref_audio_en])],
+    # Duration control (1s ≈ 12.5 tokens)
+    [processor.build_user_message(text=text_en, tokens=325)],
+    [processor.build_user_message(text=text_en, tokens=600)],
+]
+
+save_dir = Path("inference_root")
+save_dir.mkdir(exist_ok=True, parents=True)
+
+with torch.no_grad():
+    for idx, conversation in enumerate(conversations):
+        batch = processor([conversation], mode="generation")
+        outputs = model.generate(
+            input_ids=batch["input_ids"].to(device),
+            attention_mask=batch["attention_mask"].to(device),
+            max_new_tokens=4096,
+        )
+        for message in processor.decode([(sl, ids.long()) for sl, ids in outputs]):
+            audio = message.audio_codes_list[0]
+            out_path = save_dir / f"sample{idx}.wav"
+            torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+            print(f"Saved {out_path}")
+```
+
+### Input reference
+
+**UserMessage fields**
+
+| Field | Type | Required | Description |
+|---|---|:---:|---|
+| `text` | `str` | ✅ | Text to synthesize. Supports all 20 languages, raw text, Pinyin, IPA, or any mix. |
+| `reference` | `List[str]` | ❌ | Reference audio path(s) or URL(s) for zero-shot voice cloning. One audio expected. |
+| `tokens` | `int` | ❌ | Target audio token count for duration control. **1 second ≈ 12.5 tokens.** |
+
+### Generation hyperparameters
+
+| Parameter | Default | Description |
+|---|---|---|
+| `max_new_tokens` | — | Total audio tokens to generate. Use the 1s ≈ 12.5 tokens rule. |
+| `audio_temperature` | 1.7 | Higher = more variation; lower = more stable prosody. |
+| `audio_top_p` | 0.8 | Nucleus sampling cutoff. |
+| `audio_top_k` | 25 | Top-K sampling. |
+| `audio_repetition_penalty` | 1.0 | Values > 1.0 discourage repeating patterns. |
+
+---
+
+## License
+
+Apache 2.0, consistent with the original [MOSS-TTS](https://huggingface.co/OpenMOSS-Team/MOSS-TTS) release.

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": 32,
+  "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 32):
+            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 = 32,
+        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,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": 151645,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 151643,
+  "transformers_version": "5.2.0",
+  "use_cache": true
+}

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-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

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

modeling_moss_tts.py

@@ -0,0 +1,512 @@
+# 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:
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids[..., 0])
+        for i, embed_layer in enumerate(self.emb_ext):
+            inputs_embeds = inputs_embeds + embed_layer(input_ids[..., i + 1])
+        return inputs_embeds.to(torch.float16)  # ← force fp16
+
+    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: int = 1000,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+    ):
+        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)
+
+        audio_lengths = torch.zeros(batch_size, dtype=torch.int64, device=device)
+        torch_int64_max = torch.iinfo(torch.int64).max
+        delayed_lengths = torch.full((batch_size,), torch_int64_max, dtype=torch.int64, device=device)
+        
+        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
+        
+        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()
+            next_text_token = torch.full((batch_size,), self.config.pad_token_id, device=device)
+            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
+            is_stopping[next_text_token == self.config.im_end_token_id] = True
+
+            next_audio_tokens = torch.full((batch_size, n_vq), self.config.audio_pad_code, device=device)
+            
+            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_ch0_logits = next_token_logits[1][sampling_audio_mask[:, 0]]
+                audio_logits = torch.stack(next_token_logits[2:], dim=1)[sampling_audio_mask[:, 1:]]
+                audio_ch0_logits[..., self.config.audio_pad_code] = float('-inf')
+                audio_logits[..., self.config.audio_pad_code] = float('-inf')
+                next_audio_tokens[:, 0][sampling_audio_mask[:, 0]] = sample_token(
+                    logits=audio_ch0_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
+                )
+                next_audio_tokens[:, 1:][sampling_audio_mask[:, 1:]] = sample_token(
+                    logits=audio_logits,
+                    prev_tokens=generation_ids[:, :, 2:],
+                    repetition_penalty=audio_repetition_penalty,
+                    top_p=audio_top_p, 
+                    top_k=audio_top_k, 
+                    do_sample=audio_do_sample
+                )
+                
+            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)
+            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)
+            
+            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

moss_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": 32,
+  "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"
+}

moss_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 32):
+            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 = 32,
+        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

moss_extras.safetensors

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

moss_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

moss_modeling_moss_tts.py

@@ -0,0 +1,512 @@
+# 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:
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids[..., 0])
+        for i, embed_layer in enumerate(self.emb_ext):
+            inputs_embeds = inputs_embeds + embed_layer(input_ids[..., i + 1])
+        return inputs_embeds.to(torch.float16)  # ← force fp16
+
+    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: int = 1000,
+        text_temperature: float = 1.5,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: float = 1.7,
+        audio_top_p: float = 0.8,
+        audio_top_k: int = 25,
+        audio_repetition_penalty: float = 1.0,
+    ):
+        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)
+
+        audio_lengths = torch.zeros(batch_size, dtype=torch.int64, device=device)
+        torch_int64_max = torch.iinfo(torch.int64).max
+        delayed_lengths = torch.full((batch_size,), torch_int64_max, dtype=torch.int64, device=device)
+        
+        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
+        
+        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()
+            next_text_token = torch.full((batch_size,), self.config.pad_token_id, device=device)
+            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
+            is_stopping[next_text_token == self.config.im_end_token_id] = True
+
+            next_audio_tokens = torch.full((batch_size, n_vq), self.config.audio_pad_code, device=device)
+            
+            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_ch0_logits = next_token_logits[1][sampling_audio_mask[:, 0]]
+                audio_logits = torch.stack(next_token_logits[2:], dim=1)[sampling_audio_mask[:, 1:]]
+                audio_ch0_logits[..., self.config.audio_pad_code] = float('-inf')
+                audio_logits[..., self.config.audio_pad_code] = float('-inf')
+                next_audio_tokens[:, 0][sampling_audio_mask[:, 0]] = sample_token(
+                    logits=audio_ch0_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
+                )
+                next_audio_tokens[:, 1:][sampling_audio_mask[:, 1:]] = sample_token(
+                    logits=audio_logits,
+                    prev_tokens=generation_ids[:, :, 2:],
+                    repetition_penalty=audio_repetition_penalty,
+                    top_p=audio_top_p, 
+                    top_k=audio_top_k, 
+                    do_sample=audio_do_sample
+                )
+                
+            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)
+            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)
+            
+            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

moss_processing_moss_tts.py

@@ -0,0 +1,930 @@
+# 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
+
+processing_utils.MODALITY_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
+
+    def __post_init__(self):
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- 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 not None:
+                    reference.append(f"[S{speaker_idx+1}]:\n{AUDIO_PLACEHOLDER}")
+            reference = "\n".join(reference)
+            audio_codes_list = [
+                speaker_reference
+                for speaker_reference in self.reference
+                if speaker_reference is not None
+            ]
+        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("{tokens}", str(self.tokens))
+            .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("{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",
+)
+
+
+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,
+    ) -> 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,
+        ).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
+
+        if len(audio_codes_list):
+            n_vq = audio_codes_list[0].shape[1]
+        else:
+            n_vq = self.model_config.n_vq
+
+        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
+        ]
+    
+        # Fallback: pad to (NQ, B, T) + mask, then decode.
+        nq = int(codes_list[0].shape[0])
+        max_t = max(int(c.shape[1]) for c in codes_list)
+        audio_codes = torch.zeros(
+            nq, len(codes_list), max_t, device=device, dtype=torch.long
+        )
+        padding_mask = torch.zeros(
+            len(codes_list), max_t, device=device, dtype=torch.bool
+        )
+        for i, c in enumerate(codes_list):
+            t = int(c.shape[1])
+            audio_codes[:, i, :t] = c
+            padding_mask[i, :t] = True
+        dec = audio_tokenizer.decode(
+            audio_codes, padding_mask=padding_mask, return_dict=True, chunk_duration=8
+        )
+        audio = dec.audio
+        audio_lengths = dec.audio_lengths
+
+        if audio is None or audio_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.decode() returned empty outputs (audio/audio_lengths)."
+            )
+
+        # Return historical contract: list of 1D waveforms (T,)
+        wav_list: List[torch.Tensor] = []
+        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

processing_moss_tts.py

@@ -0,0 +1,930 @@
+# 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
+
+processing_utils.MODALITY_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
+
+    def __post_init__(self):
+        template = """<user_inst>
+- Reference(s):
+{reference}
+- Instruction:
+{instruction}
+- Tokens:
+{tokens}
+- Quality:
+{quality}
+- Sound Event:
+{sound_event}
+- Ambient Sound:
+{ambient_sound}
+- Language:
+{language}
+- 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 not None:
+                    reference.append(f"[S{speaker_idx+1}]:\n{AUDIO_PLACEHOLDER}")
+            reference = "\n".join(reference)
+            audio_codes_list = [
+                speaker_reference
+                for speaker_reference in self.reference
+                if speaker_reference is not None
+            ]
+        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("{tokens}", str(self.tokens))
+            .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("{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",
+)
+
+
+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,
+    ) -> 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,
+        ).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
+
+        if len(audio_codes_list):
+            n_vq = audio_codes_list[0].shape[1]
+        else:
+            n_vq = self.model_config.n_vq
+
+        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
+        ]
+    
+        # Fallback: pad to (NQ, B, T) + mask, then decode.
+        nq = int(codes_list[0].shape[0])
+        max_t = max(int(c.shape[1]) for c in codes_list)
+        audio_codes = torch.zeros(
+            nq, len(codes_list), max_t, device=device, dtype=torch.long
+        )
+        padding_mask = torch.zeros(
+            len(codes_list), max_t, device=device, dtype=torch.bool
+        )
+        for i, c in enumerate(codes_list):
+            t = int(c.shape[1])
+            audio_codes[:, i, :t] = c
+            padding_mask[i, :t] = True
+        dec = audio_tokenizer.decode(
+            audio_codes, padding_mask=padding_mask, return_dict=True, chunk_duration=8
+        )
+        audio = dec.audio
+        audio_lengths = dec.audio_lengths
+
+        if audio is None or audio_lengths is None:
+            raise RuntimeError(
+                "audio_tokenizer.decode() returned empty outputs (audio/audio_lengths)."
+            )
+
+        # Return historical contract: list of 1D waveforms (T,)
+        wav_list: List[torch.Tensor] = []
+        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

quant_log.csv

@@ -0,0 +1,253 @@
+layer,module,loss,samples,damp,time
+0,self_attn.v_proj,0.0000000010,0.01000,1.829
+0,self_attn.q_proj,0.0000000062,0.01000,1.845
+0,self_attn.k_proj,0.0000000016,0.01000,1.847
+0,self_attn.o_proj,0.0000000027,0.01000,0.577
+0,mlp.gate_proj,0.0000005435,0.01000,1.155
+0,mlp.up_proj,0.0000004846,0.01000,1.162
+0,mlp.down_proj,0.0000002644,0.01000,2.252
+1,self_attn.v_proj,0.0000000130,0.01000,1.642
+1,self_attn.q_proj,0.0000000625,0.01000,1.648
+1,self_attn.k_proj,0.0000000153,0.01000,1.652
+1,self_attn.o_proj,0.0000000898,0.01000,0.568
+1,mlp.up_proj,0.0000072422,0.01000,1.143
+1,mlp.gate_proj,0.0000082405,0.01000,1.145
+1,mlp.down_proj,0.0000001166,0.01000,2.235
+2,self_attn.q_proj,0.0000000749,0.01000,1.535
+2,self_attn.v_proj,0.0000000177,0.01000,1.563
+2,self_attn.k_proj,0.0000000193,0.01000,1.566
+2,self_attn.o_proj,0.0000000921,0.01000,0.570
+2,mlp.up_proj,0.0000127681,0.01000,1.150
+2,mlp.gate_proj,0.0000122131,0.01000,1.151
+2,mlp.down_proj,0.0000003517,0.01000,2.241
+3,self_attn.v_proj,0.0000000529,0.01000,1.574
+3,self_attn.k_proj,0.0000000404,0.01000,1.573
+3,self_attn.q_proj,0.0000001661,0.01000,1.543
+3,self_attn.o_proj,0.0000002439,0.01000,0.574
+3,mlp.up_proj,0.0000191999,0.01000,1.128
+3,mlp.gate_proj,0.0000186392,0.01000,1.133
+3,mlp.down_proj,0.0000030536,0.01000,2.273
+4,self_attn.v_proj,0.0000000613,0.01000,1.539
+4,self_attn.q_proj,0.0000002399,0.01000,1.545
+4,self_attn.k_proj,0.0000000552,0.01000,1.554
+4,self_attn.o_proj,0.0000003474,0.01000,0.573
+4,mlp.gate_proj,0.0000143707,0.01000,1.140
+4,mlp.up_proj,0.0000126579,0.01000,1.146
+4,mlp.down_proj,0.0000051150,0.01000,2.245
+5,self_attn.v_proj,0.0000000678,0.01000,1.541
+5,self_attn.q_proj,0.0000002662,0.01000,1.547
+5,self_attn.k_proj,0.0000000669,0.01000,1.552
+5,self_attn.o_proj,0.0000001778,0.01000,0.571
+5,mlp.gate_proj,0.0000023585,0.01000,1.141
+5,mlp.up_proj,0.0000018802,0.01000,1.149
+5,mlp.down_proj,0.0000004271,0.01000,2.224
+6,self_attn.q_proj,0.0000002649,0.01000,1.560
+6,self_attn.k_proj,0.0000000600,0.01000,1.563
+6,self_attn.v_proj,0.0000000583,0.01000,1.566
+6,self_attn.o_proj,0.0000001879,0.01000,0.583
+6,mlp.gate_proj,0.0000024982,0.01000,1.135
+6,mlp.up_proj,0.0000023775,0.01000,1.138
+6,mlp.down_proj,0.0000003491,0.01000,2.253
+7,self_attn.v_proj,0.0000000447,0.01000,1.562
+7,self_attn.k_proj,0.0000000380,0.01000,1.567
+7,self_attn.q_proj,0.0000001416,0.01000,1.576
+7,self_attn.o_proj,0.0000002175,0.01000,0.576
+7,mlp.up_proj,0.0000024159,0.01000,1.150
+7,mlp.gate_proj,0.0000028353,0.01000,1.154
+7,mlp.down_proj,0.0000003850,0.01000,2.244
+8,self_attn.q_proj,0.0000001304,0.01000,1.590
+8,self_attn.v_proj,0.0000000394,0.01000,1.592
+8,self_attn.k_proj,0.0000000326,0.01000,1.604
+8,self_attn.o_proj,0.0000001448,0.01000,0.575
+8,mlp.gate_proj,0.0000018485,0.01000,1.137
+8,mlp.up_proj,0.0000019048,0.01000,1.141
+8,mlp.down_proj,0.0000003137,0.01000,2.253
+9,self_attn.v_proj,0.0000000571,0.01000,1.579
+9,self_attn.q_proj,0.0000001757,0.01000,1.582
+9,self_attn.k_proj,0.0000000423,0.01000,1.585
+9,self_attn.o_proj,0.0000002859,0.01000,0.584
+9,mlp.gate_proj,0.0000019399,0.01000,1.186
+9,mlp.up_proj,0.0000018338,0.01000,1.193
+9,mlp.down_proj,0.0000003901,0.01000,2.240
+10,self_attn.v_proj,0.0000000418,0.01000,1.544
+10,self_attn.k_proj,0.0000000374,0.01000,1.549
+10,self_attn.q_proj,0.0000001466,0.01000,1.563
+10,self_attn.o_proj,0.0000003586,0.01000,0.568
+10,mlp.up_proj,0.0000012885,0.01000,1.135
+10,mlp.gate_proj,0.0000015318,0.01000,1.140
+10,mlp.down_proj,0.0000001714,0.01000,2.251
+11,self_attn.q_proj,0.0000000464,0.01000,1.541
+11,self_attn.v_proj,0.0000000131,0.01000,1.542
+11,self_attn.k_proj,0.0000000114,0.01000,1.546
+11,self_attn.o_proj,0.0000001012,0.01000,0.568
+11,mlp.gate_proj,0.0000006189,0.01000,1.147
+11,mlp.up_proj,0.0000005842,0.01000,1.161
+11,mlp.down_proj,0.0000001167,0.01000,2.239
+12,self_attn.q_proj,0.0000000243,0.01000,1.589
+12,self_attn.v_proj,0.0000000072,0.01000,1.600
+12,self_attn.k_proj,0.0000000058,0.01000,1.604
+12,self_attn.o_proj,0.0000000608,0.01000,0.578
+12,mlp.gate_proj,0.0000005512,0.01000,1.144
+12,mlp.up_proj,0.0000005652,0.01000,1.152
+12,mlp.down_proj,0.0000000815,0.01000,2.234
+13,self_attn.q_proj,0.0000000319,0.01000,1.542
+13,self_attn.v_proj,0.0000000089,0.01000,1.546
+13,self_attn.k_proj,0.0000000074,0.01000,1.549
+13,self_attn.o_proj,0.0000000500,0.01000,0.566
+13,mlp.gate_proj,0.0000004800,0.01000,1.142
+13,mlp.up_proj,0.0000004858,0.01000,1.147
+13,mlp.down_proj,0.0000000940,0.01000,2.240
+14,self_attn.k_proj,0.0000000096,0.01000,1.534
+14,self_attn.q_proj,0.0000000405,0.01000,1.539
+14,self_attn.v_proj,0.0000000106,0.01000,1.544
+14,self_attn.o_proj,0.0000000484,0.01000,0.575
+14,mlp.up_proj,0.0000005409,0.01000,1.146
+14,mlp.gate_proj,0.0000005316,0.01000,1.150
+14,mlp.down_proj,0.0000000756,0.01000,2.261
+15,self_attn.k_proj,0.0000000103,0.01000,1.544
+15,self_attn.q_proj,0.0000000466,0.01000,1.564
+15,self_attn.v_proj,0.0000000119,0.01000,1.567
+15,self_attn.o_proj,0.0000000454,0.01000,0.566
+15,mlp.up_proj,0.0000005134,0.01000,1.148
+15,mlp.gate_proj,0.0000005055,0.01000,1.153
+15,mlp.down_proj,0.0000001016,0.01000,2.260
+16,self_attn.k_proj,0.0000000187,0.01000,1.558
+16,self_attn.q_proj,0.0000000835,0.01000,1.570
+16,self_attn.v_proj,0.0000000236,0.01000,1.574
+16,self_attn.o_proj,0.0000000851,0.01000,0.576
+16,mlp.gate_proj,0.0000006101,0.01000,1.139
+16,mlp.up_proj,0.0000006160,0.01000,1.142
+16,mlp.down_proj,0.0000004110,0.01000,2.231
+17,self_attn.k_proj,0.0000000142,0.01000,1.549
+17,self_attn.v_proj,0.0000000171,0.01000,1.560
+17,self_attn.q_proj,0.0000000693,0.01000,1.566
+17,self_attn.o_proj,0.0000000767,0.01000,0.569
+17,mlp.gate_proj,0.0000005645,0.01000,1.157
+17,mlp.up_proj,0.0000005859,0.01000,1.164
+17,mlp.down_proj,0.0000000809,0.01000,2.231
+18,self_attn.q_proj,0.0000001014,0.01000,1.529
+18,self_attn.v_proj,0.0000000304,0.01000,1.539
+18,self_attn.k_proj,0.0000000222,0.01000,1.541
+18,self_attn.o_proj,0.0000000898,0.01000,0.566
+18,mlp.gate_proj,0.0000004769,0.01000,1.132
+18,mlp.up_proj,0.0000005321,0.01000,1.137
+18,mlp.down_proj,0.0000001979,0.01000,2.247
+19,self_attn.q_proj,0.0000001481,0.01000,1.541
+19,self_attn.k_proj,0.0000000314,0.01000,1.548
+19,self_attn.v_proj,0.0000000349,0.01000,1.551
+19,self_attn.o_proj,0.0000000803,0.01000,0.573
+19,mlp.gate_proj,0.0000004752,0.01000,1.147
+19,mlp.up_proj,0.0000004978,0.01000,1.149
+19,mlp.down_proj,0.0000000639,0.01000,2.242
+20,self_attn.q_proj,0.0000001669,0.01000,1.539
+20,self_attn.k_proj,0.0000000367,0.01000,1.551
+20,self_attn.v_proj,0.0000000457,0.01000,1.553
+20,self_attn.o_proj,0.0000000969,0.01000,0.567
+20,mlp.up_proj,0.0000004723,0.01000,1.139
+20,mlp.gate_proj,0.0000004650,0.01000,1.143
+20,mlp.down_proj,0.0000000665,0.01000,2.265
+21,self_attn.q_proj,0.0000001604,0.01000,1.552
+21,self_attn.v_proj,0.0000000494,0.01000,1.559
+21,self_attn.k_proj,0.0000000325,0.01000,1.566
+21,self_attn.o_proj,0.0000000864,0.01000,0.566
+21,mlp.gate_proj,0.0000004808,0.01000,1.142
+21,mlp.up_proj,0.0000005138,0.01000,1.148
+21,mlp.down_proj,0.0000002167,0.01000,2.243
+22,self_attn.q_proj,0.0000005798,0.01000,1.572
+22,self_attn.k_proj,0.0000001053,0.01000,1.580
+22,self_attn.v_proj,0.0000001031,0.01000,1.582
+22,self_attn.o_proj,0.0000001522,0.01000,0.563
+22,mlp.up_proj,0.0000006384,0.01000,1.142
+22,mlp.gate_proj,0.0000005433,0.01000,1.146
+22,mlp.down_proj,0.0000000918,0.01000,2.259
+23,self_attn.q_proj,0.0000006371,0.01000,1.548
+23,self_attn.k_proj,0.0000001273,0.01000,1.554
+23,self_attn.v_proj,0.0000001431,0.01000,1.557
+23,self_attn.o_proj,0.0000002528,0.01000,0.573
+23,mlp.gate_proj,0.0000006702,0.01000,1.130
+23,mlp.up_proj,0.0000007339,0.01000,1.137
+23,mlp.down_proj,0.0000000830,0.01000,2.247
+24,self_attn.k_proj,0.0000001151,0.01000,1.554
+24,self_attn.v_proj,0.0000001562,0.01000,1.558
+24,self_attn.q_proj,0.0000008313,0.01000,1.564
+24,self_attn.o_proj,0.0000002746,0.01000,0.572
+24,mlp.up_proj,0.0000009086,0.01000,1.142
+24,mlp.gate_proj,0.0000008630,0.01000,1.148
+24,mlp.down_proj,0.0000002657,0.01000,2.239
+25,self_attn.q_proj,0.0000007934,0.01000,1.552
+25,self_attn.v_proj,0.0000001541,0.01000,1.554
+25,self_attn.k_proj,0.0000001487,0.01000,1.561
+25,self_attn.o_proj,0.0000003166,0.01000,0.572
+25,mlp.gate_proj,0.0000008579,0.01000,1.151
+25,mlp.up_proj,0.0000009756,0.01000,1.156
+25,mlp.down_proj,0.0000002044,0.01000,2.242
+26,self_attn.q_proj,0.0000014293,0.01000,1.558
+26,self_attn.k_proj,0.0000002228,0.01000,1.562
+26,self_attn.v_proj,0.0000002375,0.01000,1.563
+26,self_attn.o_proj,0.0000004076,0.01000,0.567
+26,mlp.gate_proj,0.0000009329,0.01000,1.151
+26,mlp.up_proj,0.0000011207,0.01000,1.154
+26,mlp.down_proj,0.0000003374,0.01000,2.246
+27,self_attn.k_proj,0.0000004542,0.01000,1.571
+27,self_attn.q_proj,0.0000023284,0.01000,1.599
+27,self_attn.v_proj,0.0000008071,0.01000,1.606
+27,self_attn.o_proj,0.0000011174,0.01000,0.568
+27,mlp.up_proj,0.0000015285,0.01000,1.143
+27,mlp.gate_proj,0.0000011756,0.01000,1.147
+27,mlp.down_proj,0.0000007143,0.01000,2.246
+28,self_attn.v_proj,0.0000006405,0.01000,1.542
+28,self_attn.k_proj,0.0000006730,0.01000,1.546
+28,self_attn.q_proj,0.0000030264,0.01000,1.549
+28,self_attn.o_proj,0.0000012090,0.01000,0.572
+28,mlp.up_proj,0.0000021163,0.01000,1.138
+28,mlp.gate_proj,0.0000018347,0.01000,1.146
+28,mlp.down_proj,0.0000017199,0.01000,2.240
+29,self_attn.v_proj,0.0000015738,0.01000,1.545
+29,self_attn.q_proj,0.0000077582,0.01000,1.558
+29,self_attn.k_proj,0.0000012017,0.01000,1.566
+29,self_attn.o_proj,0.0000032131,0.01000,0.565
+29,mlp.up_proj,0.0000028664,0.01000,1.148
+29,mlp.gate_proj,0.0000025585,0.01000,1.151
+29,mlp.down_proj,0.0000256958,0.01000,2.225
+30,self_attn.k_proj,0.0000013016,0.01000,1.539
+30,self_attn.q_proj,0.0000079752,0.01000,1.542
+30,self_attn.v_proj,0.0000017943,0.01000,1.547
+30,self_attn.o_proj,0.0000069710,0.01000,0.563
+30,mlp.up_proj,0.0000044374,0.01000,1.152
+30,mlp.gate_proj,0.0000036016,0.01000,1.156
+30,mlp.down_proj,0.0000099136,0.01000,2.248
+31,self_attn.k_proj,0.0000022095,0.01000,1.536
+31,self_attn.q_proj,0.0000119859,0.01000,1.540
+31,self_attn.v_proj,0.0000024236,0.01000,1.543
+31,self_attn.o_proj,0.0000090073,0.01000,0.565
+31,mlp.gate_proj,0.0000042423,0.01000,1.140
+31,mlp.up_proj,0.0000057331,0.01000,1.150
+31,mlp.down_proj,0.0000113583,0.01000,2.248
+32,self_attn.k_proj,0.0000032344,0.01000,1.553
+32,self_attn.q_proj,0.0000177639,0.01000,1.556
+32,self_attn.v_proj,0.0000057728,0.01000,1.588
+32,self_attn.o_proj,0.0000294938,0.01000,0.575
+32,mlp.gate_proj,0.0000079738,0.01000,1.134
+32,mlp.up_proj,0.0000090901,0.01000,1.134
+32,mlp.down_proj,0.0001718242,0.01000,2.271
+33,self_attn.k_proj,0.0000067906,0.01000,1.549
+33,self_attn.q_proj,0.0000625890,0.01000,1.553
+33,self_attn.v_proj,0.0000118461,0.01000,1.557
+33,self_attn.o_proj,0.0000358326,0.01000,0.568
+33,mlp.up_proj,0.0000198184,0.01000,1.151
+33,mlp.gate_proj,0.0000165273,0.01000,1.159
+33,mlp.down_proj,0.0001491802,0.01000,2.239
+34,self_attn.q_proj,0.0000543061,0.01000,1.567
+34,self_attn.v_proj,0.0000131741,0.01000,1.573
+34,self_attn.k_proj,0.0000087816,0.01000,1.579
+34,self_attn.o_proj,0.0000818334,0.01000,0.569
+34,mlp.gate_proj,0.0000169936,0.01000,1.145
+34,mlp.up_proj,0.0000183243,0.01000,1.151
+34,mlp.down_proj,0.0000622494,0.01000,2.241
+35,self_attn.q_proj,0.0000360077,0.01000,1.567
+35,self_attn.k_proj,0.0000084222,0.01000,1.570
+35,self_attn.v_proj,0.0000127596,0.01000,1.587
+35,self_attn.o_proj,0.0000860495,0.01000,0.571
+35,mlp.gate_proj,0.0000566747,0.01000,1.142
+35,mlp.up_proj,0.0000673418,0.01000,1.145
+35,mlp.down_proj,0.0008014260,0.01000,2.267

quantize_config.json

@@ -0,0 +1,46 @@
+{
+  "bits": 4,
+  "group_size": 128,
+  "desc_act": false,
+  "lm_head": false,
+  "quant_method": "gptq",
+  "checkpoint_format": "gptq",
+  "pack_dtype": "int32",
+  "meta": {
+    "quantizer": [
+      "gptqmodel:5.7.1"
+    ],
+    "uri": "https://github.com/modelcloud/gptqmodel",
+    "damp_percent": 0.01,
+    "damp_auto_increment": 0.01,
+    "static_groups": false,
+    "true_sequential": true,
+    "mse": 0.0,
+    "gptaq": null,
+    "act_group_aware": true,
+    "failsafe": {
+      "strategy": "rtn",
+      "threshold": "0.5%",
+      "smooth": {
+        "type": "mad",
+        "group_size_threshold": 128,
+        "k": 2.75
+      }
+    },
+    "offload_to_disk": true,
+    "offload_to_disk_path": "./gptqmodel_offload/hqdpgrum-rkaakpxx/",
+    "pack_impl": "cpu",
+    "mock_quantization": false,
+    "gc_mode": "interval",
+    "wait_for_submodule_finalizers": false,
+    "auto_forward_data_parallel": true,
+    "hessian": {
+      "chunk_size": null,
+      "chunk_bytes": null,
+      "staging_dtype": "float32"
+    },
+    "vram_strategy": "exclusive"
+  },
+  "sym": true,
+  "format": "gptq"
+}

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
+}