Duplicate from OpenMOSS-Team/MOSS-TTS-v1.5

Co-authored-by: YWMditto <YWMditto@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -0,0 +1,260 @@
+---
+license: apache-2.0
+tags:
+- text-to-speech
+language:
+- zh
+- yue
+- en
+- ar
+- cs
+- da
+- de
+- nl
+- es
+- fr
+- fi
+- el
+- he
+- hi
+- hu
+- ja
+- it
+- ko
+- mk
+- ms
+- ru
+- fa
+- pl
+- pt
+- sv
+- ro
+- sw
+- tl
+- th
+- tr
+- vi
+---
+# MOSS-TTS Family
+
+
+<br>
+
+<p align="center">
+  &nbsp;&nbsp;&nbsp;&nbsp;
+  <img src="https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_imgaes_demo/openmoss_x_mosi" height="50" align="middle" />
+</p>
+
+
+
+<div align="center">
+  <a href="https://github.com/OpenMOSS/MOSS-TTS/tree/main"><img src="https://img.shields.io/badge/Project%20Page-GitHub-blue"></a>
+  <a href="https://modelscope.cn/collections/OpenMOSS-Team/MOSS-TTS"><img src="https://img.shields.io/badge/ModelScope-Models-lightgrey?logo=modelscope&amp"></a>
+  <a href="https://mosi.cn/#models"><img src="https://img.shields.io/badge/Blog-View-blue?logo=internet-explorer&amp"></a>
+  <a href="https://arxiv.org/abs/2603.18090"><img src="https://img.shields.io/badge/Arxiv-2603.18090-red?logo=Arxiv&amp"></a>
+
+  <a href="https://studio.mosi.cn"><img src="https://img.shields.io/badge/AIStudio-Try-green?logo=internet-explorer&amp"></a>
+  <a href="https://studio.mosi.cn/docs/moss-tts"><img src="https://img.shields.io/badge/API-Docs-00A3FF?logo=fastapi&amp"></a>
+  <a href="https://x.com/Open_MOSS"><img src="https://img.shields.io/badge/Twitter-Follow-black?logo=x&amp"></a>
+  <a href="https://discord.gg/fvm5TaWjU3"><img src="https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&amp"></a>
+</div>
+
+
+# MOSS-TTS-v1.5
+
+**MOSS-TTS-v1.5** is continued from [MOSS-TTS 1.0](https://huggingface.co/OpenMOSS-Team/MOSS-TTS). It preserves the main 1.0 capabilities, including zero-shot voice cloning, long-form speech generation, token-level duration control, Pinyin/IPA pronunciation control, multilingual synthesis, and code-switching. For the full 1.0 feature walkthrough, input schema, decoding hyperparameters, and evaluation tables, please refer to the [MOSS-TTS 1.0 README](https://huggingface.co/OpenMOSS-Team/MOSS-TTS).
+
+Compared with MOSS-TTS 1.0, v1.5 focuses on the following improvements:
+
+- **Stronger multilingual synthesis with language tags**: when the `language` field is omitted, v1.5 may improve some languages and regress slightly on others compared with 1.0. When the language is specified, v1.5 is stronger than 1.0 on almost all supported languages. Set the tag when building the user message, for example `processor.build_user_message(text=text_fr, language="French")`.
+- **More stable voice cloning**: v1.5 improves speaker similarity and reduces cloning variance, making repeated generations more consistent.
+- **Better long-reference, short-text cloning**: v1.5 handles scenarios where the reference audio is much longer than the target text more reliably than 1.0.
+- **More stable punctuation-following prosody**: v1.5 follows punctuation-driven pauses more closely, especially in long sentences.
+- **Explicit pause control**: v1.5 supports inline pause markers such as `"[pause 3.2s]"`. For example, `我今天学习了一首中国的古诗，它的名字是[pause 3.2s]静夜思！` inserts an explicit 3.2s pause before `静夜思`.
+
+## Supported Languages
+
+MOSS-TTS-v1.5 currently supports **31 languages**. It keeps the 20 languages supported by [MOSS-TTS 1.0](https://huggingface.co/OpenMOSS-Team/MOSS-TTS) and extends multilingual continued training to additional languages including Cantonese, Dutch, Finnish, Hindi, Macedonian, Malay, Romanian, Swahili, Tagalog, Thai, and Vietnamese.
+
+| Language | Code | Flag | Language | Code | Flag | Language | Code | Flag |
+|---|---|---|---|---|---|---|---|---|
+| Chinese | zh | 🇨🇳 | Cantonese | yue | 🇭🇰 | English | en | 🇺🇸 |
+| Arabic | ar | 🇸🇦 | Czech | cs | 🇨🇿 | Danish | da | 🇩🇰 |
+| Dutch | nl | 🇳🇱 | Finnish | fi | 🇫🇮 | French | fr | 🇫🇷 |
+| German | de | 🇩🇪 | Greek | el | 🇬🇷 | Hebrew | he | 🇮🇱 |
+| Hindi | hi | 🇮🇳 | Hungarian | hu | 🇭🇺 | Italian | it | 🇮🇹 |
+| Japanese | ja | 🇯🇵 | Korean | ko | 🇰🇷 | Macedonian | mk | 🇲🇰 |
+| Malay | ms | 🇲🇾 | Persian (Farsi) | fa | 🇮🇷 | Polish | pl | 🇵🇱 |
+| Portuguese | pt | 🇵🇹 | Romanian | ro | 🇷🇴 | Russian | ru | 🇷🇺 |
+| Spanish | es | 🇪🇸 | Swahili | sw | 🇹🇿 | Swedish | sv | 🇸🇪 |
+| Tagalog | tl | 🇵🇭 | Thai | th | 🇹🇭 | Turkish | tr | 🇹🇷 |
+| Vietnamese | vi | 🇻🇳 | | | | | | |
+
+
+## Quick Start
+
+### Environment Setup
+
+We recommend a clean, isolated Python environment with **Transformers 5.0.0** to avoid dependency conflicts.
+
+```bash
+conda create -n moss-tts python=3.12 -y
+conda activate moss-tts
+```
+
+Install all required dependencies:
+
+```bash
+git clone https://github.com/OpenMOSS/MOSS-TTS.git
+cd MOSS-TTS
+pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e .
+```
+
+#### (Optional) Install FlashAttention 2
+
+For better speed and lower GPU memory usage, you can install FlashAttention 2 if your hardware supports it.
+
+```bash
+pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e ".[flash-attn]"
+```
+
+If your machine has limited RAM and many CPU cores, you can cap build parallelism:
+
+```bash
+MAX_JOBS=4 pip install --extra-index-url https://download.pytorch.org/whl/cu128 -e ".[flash-attn]"
+```
+
+Notes:
+- Dependencies are managed in `pyproject.toml`, which currently pins `torch==2.9.1+cu128` and `torchaudio==2.9.1+cu128`.
+- If FlashAttention 2 fails to build on your machine, you can skip it and use the default attention backend.
+- FlashAttention 2 is only available on supported GPUs and is typically used with `torch.float16` or `torch.bfloat16`.
+
+
+### Basic Usage
+
+> Tip: MOSS-TTS-v1.5 uses the same generation API as the 1.0 **MossTTSDelay-8B** checkpoint. For multilingual inputs, set `language` whenever the language is known.
+
+MOSS-TTS provides a convenient `generate` interface for rapid usage. The examples below cover:
+1. Direct generation (Chinese / English / multilingual text with language tags / Pinyin / IPA)
+2. Voice cloning
+3. Duration control
+4. Explicit pause control with `[pause X.Ys]`
+
+```python
+from pathlib import Path
+import importlib.util
+import torch
+import torchaudio
+from transformers import AutoModel, AutoProcessor
+# Disable the broken cuDNN SDPA backend
+torch.backends.cuda.enable_cudnn_sdp(False)
+# Keep these enabled as fallbacks
+torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_mem_efficient_sdp(True)
+torch.backends.cuda.enable_math_sdp(True)
+
+
+pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-TTS-v1.5"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+def resolve_attn_implementation() -> str:
+    # Prefer FlashAttention 2 when package + device conditions are met.
+    if (
+        device == "cuda"
+        and importlib.util.find_spec("flash_attn") is not None
+        and dtype in {torch.float16, torch.bfloat16}
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        if major >= 8:
+            return "flash_attention_2"
+
+    # CUDA fallback: use PyTorch SDPA kernels.
+    if device == "cuda":
+        return "sdpa"
+
+    # CPU fallback.
+    return "eager"
+
+
+attn_implementation = resolve_attn_implementation()
+print(f"[INFO] Using attn_implementation={attn_implementation}")
+
+processor = AutoProcessor.from_pretrained(
+    pretrained_model_name_or_path,
+    trust_remote_code=True,
+)
+processor.audio_tokenizer = processor.audio_tokenizer.to(device)
+
+text_1 = "亲爱的你，\n你好呀。\n\n今天，我想用最认真、最温柔的声音，对你说一些重要的话。\n这些话，像一颗小小的星星，希望能在你的心里慢慢发光。\n\n首先，我想祝你——\n每天都能平平安安、快快乐乐。\n\n希望你早上醒来的时候，\n窗外有光，屋子里很安静，\n你的心是轻轻的，没有着急，也没有害怕。\n\n希望你吃饭的时候胃口很好，\n走路的时候脚步稳稳，\n晚上睡觉的时候，能做一个又一个甜甜的梦。\n\n我希望你能一直保持好奇心。\n对世界充满问题，\n对天空、星星、花草、书本和故事感兴趣。\n当你问“为什么”的时候，\n希望总有人愿意认真地听你说话。\n\n我也希望你学会温柔。\n温柔地对待朋友，\n温柔地对待小动物，\n也温柔地对待自己。\n\n如果有一天你犯了错，\n请不要太快责怪自己，\n因为每一个认真成长的人，\n都会在路上慢慢学会更好的方法。\n\n愿你拥有勇气。\n当你站在陌生的地方时，\n当你第一次举手发言时，\n当你遇到困难、感到害怕的时候，\n希望你能轻轻地告诉自己：\n“我可以试一试。”\n\n就算没有一次成功，也没有关系。\n失败不是坏事，\n它只是告诉你，你正在努力。\n\n我希望你学会分享快乐。\n把开心的事情告诉别人，\n把笑声送给身边的人，\n因为快乐被分享的时候，\n会变得更大、更亮。\n\n如果有一天你感到难过，\n我希望你知道——\n难过并不丢脸，\n哭泣也不是软弱。\n\n愿你能找到一个安全的地方，\n慢慢把心里的话说出来，\n然后再一次抬起头，看见希望。\n\n我还希望你能拥有梦想。\n这个梦想也许很大，\n也许很小，\n也许现在还说不清楚。\n\n没关系。\n梦想会和你一起长大，\n在时间里慢慢变得清楚。\n\n最后，我想送你一个最最重要的祝福：\n\n愿你被世界温柔对待，\n也愿你成为一个温柔的人。\n\n愿你的每一天，\n都值得被记住，\n都值得被珍惜。\n\n亲爱的你，\n请记住，\n你是独一无二的，\n你已经很棒了，\n而你的未来，\n一定会慢慢变得闪闪发光。\n\n祝你健康、勇敢、幸福，\n祝你永远带着笑容向前走。"
+text_2 = "We stand on the threshold of the AI era.\nArtificial intelligence is no longer just a concept in laboratories, but is entering every industry, every creative endeavor, and every decision. It has learned to see, hear, speak, and think, and is beginning to become an extension of human capabilities. AI is not about replacing humans, but about amplifying human creativity, making knowledge more equitable, more efficient, and allowing imagination to reach further. A new era, jointly shaped by humans and intelligent systems, has arrived."
+text_3 = "nin2 hao3，qing3 wen4 nin2 lai2 zi4 na3 zuo4 cheng2 shi4？"
+text_4 = "nin2 hao3，qing4 wen3 nin2 lai2 zi4 na4 zuo3 cheng4 shi3？"
+text_5 = "您好，请问您来自哪 zuo4 cheng2 shi4？"
+text_6 = "/həloʊ, meɪ aɪ æsk wɪtʃ sɪti juː ɑːr frʌm?/"
+text_7 = "Bonjour, je voudrais essayer une voix française naturelle et stable."
+text_8 = "我今天学习了一首中国的古诗，它的名字是[pause 3.2s]静夜思！"
+
+# Use audio from ./assets/audio to avoid downloading from the cloud.
+ref_audio_1 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_zh.wav"
+ref_audio_2 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_en.m4a"
+
+conversations = [
+    # Direct TTS (no reference). Language tags are recommended in v1.5.
+    [processor.build_user_message(text=text_1)],
+    [processor.build_user_message(text=text_2)],
+    # Direct TTS (no reference). For languages ​​other than Chinese and English, it is recommended to use language tags.
+    [processor.build_user_message(text=text_7, language="French")],
+    # Pinyin or IPA input
+    [processor.build_user_message(text=text_3)],
+    [processor.build_user_message(text=text_4)],
+    [processor.build_user_message(text=text_5)],
+    [processor.build_user_message(text=text_6)],
+    # Explicit pause control. Use [pause X.Ys], such as [pause 3.2s].
+    [processor.build_user_message(text=text_8)],
+    # Voice cloning (with reference)
+    [processor.build_user_message(text=text_1, reference=[ref_audio_1])],
+    [processor.build_user_message(text=text_2, reference=[ref_audio_2])],
+    # Duration control
+    [processor.build_user_message(text=text_2, tokens=325)],
+    [processor.build_user_message(text=text_2, tokens=600)],
+]
+
+model = AutoModel.from_pretrained(
+    pretrained_model_name_or_path,
+    trust_remote_code=True,
+    # If FlashAttention 2 is installed, you can set attn_implementation="flash_attention_2"
+    attn_implementation=attn_implementation,
+    torch_dtype=dtype,
+).to(device)
+model.eval()
+
+batch_size = 1
+
+save_dir = Path("inference_root")
+save_dir.mkdir(exist_ok=True, parents=True)
+sample_idx = 0
+with torch.no_grad():
+    for start in range(0, len(conversations), batch_size):
+        batch_conversations = conversations[start : start + batch_size]
+        batch = processor(batch_conversations, mode="generation")
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+
+        outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=4096,
+        )
+
+        for message in processor.decode(outputs):
+            audio = message.audio_codes_list[0]
+            out_path = save_dir / f"sample{sample_idx}.wav"
+            sample_idx += 1
+            torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+
+```
+
+## More Usage
+
+MOSS-TTS-v1.5 is API-compatible with MOSS-TTS 1.0. For continuation with prefix audio, detailed `UserMessage` and `AssistantMessage` fields, generation hyperparameters, Pinyin/IPA preprocessing examples, and evaluation results, see the [MOSS-TTS 1.0 README](https://huggingface.co/OpenMOSS-Team/MOSS-TTS).

added_tokens.json

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

chat_template.jinja

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

config.json

@@ -0,0 +1,89 @@
+{
+  "model_type": "moss_tts_delay",
+  "architectures": [
+    "MossTTSDelayModel"
+  ],
+  "auto_map": {
+      "AutoConfig": "configuration_moss_tts.MossTTSDelayConfig",
+      "AutoModel": "modeling_moss_tts.MossTTSDelayModel"
+  },
+  "dtype": "bfloat16",
+  "initializer_range": 0.02,
+  "language_config": {
+    "_name_or_path": "Qwen/Qwen3-8B",
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": 151645,
+    "pad_token_id": 151643,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 40960,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 155648
+  },
+  "n_vq": 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

inference_utils.py

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

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

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

modeling_moss_tts.py

@@ -0,0 +1,525 @@
+# coding=utf-8
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Modeling classes for MossTTSDelay. """
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.cache_utils import Cache
+from transformers.models.qwen3 import Qwen3Model
+from transformers import initialization as init
+
+from .configuration_moss_tts import MossTTSDelayConfig
+from .inference_utils import sample_token, find_last_equal_C
+
+try:
+    from .processing_moss_tts import UserMessage, AssistantMessage, MossTTSDelayProcessor
+except Exception:
+    UserMessage = None
+    AssistantMessage = None
+    MossTTSDelayProcessor = None
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MossTTSDelayConfig"
+
+
+@dataclass
+class MossTTSDelayOutputWithPast(ModelOutput):
+    """
+    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Weighted sum of channel losses.
+        all_sum_losses (`torch.FloatTensor` of shape `(batch_size, n_vq + 1)`, *optional*):
+            Sum of losses for each sample and each channel before averaging.
+        all_token_nums (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Number of non-masked tokens per sample.
+        sample_losses (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            Loss per sample.
+        channel_losses (`torch.FloatTensor` of shape `(n_vq + 1,)`, *optional*):
+            Loss per channel (text head + vq heads).
+        logits (`List[torch.FloatTensor]`, *optional*):
+            List of prediction scores from each head.
+        past_key_values (`Cache`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
+            Tuple of torch.FloatTensor (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer).
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed):
+            Tuple of torch.FloatTensor (one for each layer) of the attention weights.
+    """
+    loss: Optional[torch.FloatTensor] = None
+    all_sum_losses: Optional[torch.FloatTensor] = None
+    all_token_nums: Optional[torch.LongTensor] = None
+    sample_losses: Optional[torch.FloatTensor] = None
+    channel_losses: Optional[torch.FloatTensor] = None
+    logits: Optional[List[torch.FloatTensor]] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+
+
+class MossTTSDelayPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSDelayConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    def _init_weights(self, module):
+        """
+        Transformers 5.0+ safe init:
+        - MUST use transformers.initialization helpers
+        - MUST respect param._is_hf_initialized to avoid overwriting ckpt-loaded params
+        """
+        # Let HF handle its standard modules first (LayerNorm, Linear, Embedding, etc.)
+        super()._init_weights(module)
+
+        # Pick a std consistent with HF conventions
+        # Prefer model/text config initializer_range if present.
+        std = None
+        if hasattr(self.config, "initializer_range"):
+            std = self.config.initializer_range
+        elif hasattr(self.config, "language_config") and hasattr(self.config.language_config, "initializer_range"):
+            std = self.config.language_config.initializer_range
+        else:
+            std = 0.02
+
+        # Initialize extra audio embeddings
+        if isinstance(module, nn.Embedding):
+            # Only touch our extra embeddings (avoid double touching LM's embeddings if not desired)
+            # If you prefer, you can skip this check and rely on super()._init_weights for all embeddings.
+            if getattr(module, "num_embeddings", None) == self.config.audio_vocab_size + 1:
+                init.normal_(module.weight, mean=0.0, std=std)
+                # If you later set padding_idx, you must explicitly zero it (and respect _is_hf_initialized!)
+                # init.zeros_ will internally check param flags, but slicing needs manual care.
+
+        # Initialize multi-head projections you added
+        if isinstance(module, nn.Linear):
+            # For your lm_heads, super()._init_weights already covers typical Linear.
+            # This block is only needed if you have custom Linear variants later.
+            pass
+
+
+
+MOSSTTS_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`MossTTSDelayConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The MossTTSDelay Model architecture tailored for Text-to-Speech generation with multi-head VQ prediction.",
+    MOSSTTS_START_DOCSTRING,
+)
+class MossTTSDelayModel(MossTTSDelayPreTrainedModel):
+    UserMessage = UserMessage
+    AssistantMessage = AssistantMessage
+    Processor = MossTTSDelayProcessor
+
+    def __init__(self, config: MossTTSDelayConfig):
+        super().__init__(config)
+        self.config = config
+
+        config.language_config.torch_dtype = config.torch_dtype
+        
+        self.language_model = Qwen3Model(config.language_config)
+
+        # Audio VQ Embeddings (Extra channels)
+        # Note: input_ids[..., 0] uses Qwen's embedding. 
+        # input_ids[..., 1:] use these extensions.
+        self.emb_ext = nn.ModuleList()
+        for vq_idx in range(self.config.n_vq):
+            # Add +1 for potential padding/special tokens logic if strictly required by upstream data prep
+            self.emb_ext.append(
+                nn.Embedding(self.config.audio_vocab_size + 1, config.language_config.hidden_size, padding_idx=None)
+            )
+
+        # Multi-Head Prediction Layers
+        # Head 0: Main language head
+        # Head 1..N: Audio VQ heads
+        self.lm_heads = nn.ModuleList([
+            nn.Linear(config.language_config.hidden_size, config.language_config.vocab_size, bias=False)
+        ])
+        for vq_idx in range(self.config.n_vq):
+            self.lm_heads.append(
+                nn.Linear(config.language_config.hidden_size, self.config.audio_vocab_size + 1, bias=False)
+            )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Computes the combined embeddings from text and multiple audio VQ channels.
+        
+        Args:
+            input_ids: Shape (Batch, Seq_Len, 1 + n_vq)
+        """
+        # Base Text/Content Embedding
+        # input_ids[..., 0] is standard text or semantic tokens
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids[..., 0])
+
+        # Add VQ Embeddings
+        for i, embed_layer in enumerate(self.emb_ext):
+            # i corresponds to channel i+1 in input_ids
+            # We assume the data pipeline ensures indices are within range
+            inputs_embeds = inputs_embeds + embed_layer(input_ids[..., i + 1])
+        
+        return inputs_embeds
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        # Returning a list of heads might break some HF utilities expecting a single head.
+        # However, for custom models, this is acceptable.
+        return self.lm_heads
+
+    @add_start_docstrings_to_model_forward(MOSSTTS_START_DOCSTRING)
+    @replace_return_docstrings(output_type=MossTTSDelayOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        hidden_out_layers: Optional[List[int]] = None,
+        channelwise_loss_weight: Optional[List[float]] = None,
+        **kwargs,
+    ) -> Union[Tuple, MossTTSDelayOutputWithPast]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`):
+                Indices of input sequence tokens in the vocabulary. 
+                Dimension 2 contains: [Text/Semantics, VQ_0, VQ_1, ..., VQ_N].
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length, 1 + n_vq)`, *optional*):
+                Labels for computing the masked language modeling loss.
+            channelwise_loss_weight (`List[float]`, *optional*):
+                Manual weights for summing losses across different heads (Text vs Audio channels).
+
+        Returns:
+        """
+
+        if len(input_ids.shape) != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError("`Input_ids`'s shape should be exactly (batch_size, sequence_length, 1 + n_vq).")
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+        # 1. Prepare Embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+
+        # 2. Backbone Forward
+        # Qwen3Model outputs standard CausalLMOutputWithPast or similar
+        outputs = self.language_model(
+            input_ids=None, # Passed via inputs_embeds
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True, # Always need hidden states for multi-head projection
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # 3. Handle specific layer outputs if requested (Delay Pattern often requires features from specific layers)
+        last_hidden_state = outputs.last_hidden_state
+        if hidden_out_layers is None:
+            # Default to using the last layer for all heads
+            # In some architectures (like MusicGen), different codebooks come from different transformer layers.
+            # Here we default to the final layer as per original code behavior [-1] * (n + 1).
+            hidden_states_for_heads = [last_hidden_state] * (len(self.lm_heads))
+        else:
+            # If hidden_out_layers is provided (e.g. [-1, -2, -3...]), fetch them from all_hidden_states
+            # Note: outputs.hidden_states includes embedding output at index 0 usually.
+            all_hs = outputs.hidden_states
+            hidden_states_for_heads = [all_hs[idx] for idx in hidden_out_layers]
+
+        # 4. Project to Logits (Multi-Head)
+        layer_logits = []
+        for i, (hs, head) in enumerate(zip(hidden_states_for_heads, self.lm_heads)):
+            logits = head(hs)
+            # Original code logic: Mask the last token index for audio heads (indices > 0)
+            # This implies the vocab size is (N+1) but the model shouldn't predict the (N+1)-th token 
+            # (perhaps reserved for padding in the input but invalid for prediction).
+            if i > 0:
+                logits[..., -1] = float("-inf")
+            layer_logits.append(logits)
+
+        # 5. Loss Calculation
+        loss = None
+        all_sum_losses = None
+        all_token_nums = None
+        sample_losses = None
+        channel_losses = None
+
+        if labels is not None:
+            # Ensure labels match input shape rank (B, S, C)
+            if labels.dim() != 3:
+                raise ValueError(f"Labels must have rank 3 (B, S, C), got {labels.shape}")
+
+            batch_size = labels.size(0)
+            n_heads = len(layer_logits)
+            
+            # Container for per-sample, per-channel losses
+            # Shape: [Batch, n_heads]
+            all_sum_losses_list = [] 
+            
+            # Count valid tokens (not -100) per sample. 
+            # Note: Assuming mask is consistent across channels or we take sum over dim 1 (seq)
+            # Usually strict masking means checking one channel or all. 
+            # Original code: torch.sum(labels != -100, dim=1) -> [B, C]
+            all_token_nums = torch.sum(labels != -100, dim=1)
+
+            for i, logits in enumerate(layer_logits):
+                # logits: [B, S, V]
+                # cur_labels: [B, S]
+                cur_labels = labels[..., i]
+                
+                # Flatten for CrossEntropy
+                # logits: [B*S, V], labels: [B*S]
+                loss_fct = CrossEntropyLoss(reduction='none')
+                vocab_size = logits.size(-1)
+                
+                reshaped_logits = logits.view(-1, vocab_size)
+                reshaped_labels = cur_labels.contiguous().view(-1)
+                
+                # Calculate loss per token
+                per_token_loss = loss_fct(reshaped_logits, reshaped_labels)
+                
+                # Reshape back to [B, S] and sum over Sequence dimension to get per-sample loss
+                per_token_loss = per_token_loss.view(batch_size, -1)
+                per_sample_loss = torch.sum(per_token_loss, dim=-1) # [B]
+                
+                all_sum_losses_list.append(per_sample_loss)
+
+            # Stack to [B, n_heads]
+            all_sum_losses = torch.stack(all_sum_losses_list, dim=1)
+
+            # Weighted Loss Aggregation
+            if channelwise_loss_weight is not None:
+                if len(channelwise_loss_weight) != n_heads:
+                    raise ValueError(f"channelwise_loss_weight length {len(channelwise_loss_weight)} != {n_heads}")
+                
+                w_tensor = torch.tensor(channelwise_loss_weight, device=all_sum_losses.device, dtype=all_sum_losses.dtype)
+                
+                # Sample losses: Weighted sum over channels per sample / Total weight
+                # Normalize by token count per channel
+                # Avoid division by zero with epsilon or mask
+                token_counts_safe = all_token_nums.float().clamp(min=1.0)
+                
+                normalized_losses = all_sum_losses / token_counts_safe
+                sample_losses = (normalized_losses * w_tensor).sum(dim=1) / w_tensor.sum()
+                
+                # Channel losses: Sum over batch / Sum tokens over batch
+                total_loss_per_channel = all_sum_losses.sum(dim=0)
+                total_tokens_per_channel = all_token_nums.sum(dim=0).float().clamp(min=1.0)
+                channel_losses = total_loss_per_channel / total_tokens_per_channel
+                
+                # Final scalar loss
+                loss = (channel_losses * w_tensor).sum() / w_tensor.sum()
+            else:
+                # Default average if no weights provided
+                total_tokens = all_token_nums.sum().float().clamp(min=1.0)
+                loss = all_sum_losses.sum() / total_tokens
+                channel_losses = all_sum_losses.sum(dim=0) / all_token_nums.sum(dim=0).clamp(min=1.0)
+
+        return MossTTSDelayOutputWithPast(
+            loss=loss,
+            all_sum_losses=all_sum_losses,
+            all_token_nums=all_token_nums,
+            sample_losses=sample_losses,
+            channel_losses=channel_losses,
+            logits=layer_logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: 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

processing_moss_tts.py

@@ -0,0 +1,949 @@
+# 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
+from .tts_robust_normalizer_single_script import normalize_tts_text
+
+
+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", None)
+        if audio_tokenizer_name_or_path is None:
+            processor_lookup_kwargs = dict(kwargs)
+            try:
+                processor_dict, _ = cls.get_processor_dict(
+                    pretrained_model_name_or_path,
+                    **processor_lookup_kwargs,
+                )
+                audio_tokenizer_name_or_path = processor_dict.get(
+                    "audio_tokenizer_name_or_path"
+                )
+                audio_tokenizer_dict = processor_dict.get("audio_tokenizer", {})
+                if isinstance(audio_tokenizer_dict, dict):
+                    audio_tokenizer_name_or_path = audio_tokenizer_dict.get(
+                        "audio_tokenizer_name_or_path"
+                    ) or audio_tokenizer_name_or_path
+            except Exception:
+                audio_tokenizer_name_or_path = None
+        if audio_tokenizer_name_or_path is None:
+            audio_tokenizer_name_or_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", "computing_loss"}:
+            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]
+        text = normalize_tts_text(text)
+        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

processor_config.json

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

special_tokens_map.json

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

tokenizer.json

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

tokenizer_config.json

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

tts_robust_normalizer_single_script.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+TTS 输入鲁棒性正则化器（非语义 TN）
+
+目标
+----
+1. 只做“鲁棒性清洗”，不做数字/单位/日期/金额等语义展开。
+2. 优先保护高风险 token，避免把 `.map`、`app.js.map`、`v2.3.1`、URL、Email、@mention、#hashtag 清坏。
+3. 保留 `[]` / `{}` 的内容，为后续声音事件、控制标签留接口。
+4. 对结构性符号做“替换而非删除”：
+   - `【】 / 〖〗 / 『』 / 「」` 在结构位置转成句边界。
+   - `《》` 只在“独立标题/栏目名”场景拆开；嵌入式标题保持不变。
+   - `—— / -- / ——...` 转成句边界。
+   - `-> / => / →` 等流程连接符转成中文逗号，避免 TTS 读入崩溃。
+5. 对社交平台常见噪声做弱归一化：
+   - `...... / ……` -> `。`
+   - `？？？！！！` -> `？！`
+   - `！！！` -> `！`
+6. 空格按脚本类型处理：
+   - 西文片段内部：连续空格压缩为 1 个。
+   - 汉字 / 日文假名片段内部：删除空格。
+   - 汉字 / 日文假名 与“拉丁字母类 token / 受保护 token”相邻：保留或补 1 个空格。
+   - 汉字 / 日文假名 与纯数字相邻：不强行补空格。
+7. 轻量处理 Markdown 与换行：
+   - `[text](url)` -> `text url`
+   - 去掉标题 `#`、引用 `>`、列表前缀
+   - 换行转句边界 `。`
+
+非目标
+------
+1. 不决定“应该怎么读”。
+2. 不删除 `[] / {}` 内容。
+3. 不做 HTML/SSML/语义标签解释。
+"""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+
+
+# ---------------------------
+# 基础常量与正则
+# ---------------------------
+
+# 不依赖空格分词的脚本：汉字 + 日文假名
+_CJK_CHARS = r"\u3400-\u4dbf\u4e00-\u9fff\u3040-\u30ff"
+_CJK = f"[{_CJK_CHARS}]"
+
+# 保护占位符
+_PROT = r"___PROT\d+___"
+
+# 需要保护的高风险 token
+_URL_RE = re.compile(r"https?://[^\s\u3000，。！？；、）】》〉」』]+")
+_EMAIL_RE = re.compile(r"(?<![\w.+-])[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}(?![\w.-])")
+_MENTION_RE = re.compile(r"(?<![A-Za-z0-9_])@[A-Za-z0-9_]{1,32}")
+_REDDIT_RE = re.compile(r"(?<![A-Za-z0-9_])(?:u|r)/[A-Za-z0-9_]+")
+_HASHTAG_RE = re.compile(r"(?<![A-Za-z0-9_])#(?!\s)[^\s#]+")
+
+# `.map` / `.env` / `.gitignore`
+_DOT_TOKEN_RE = re.compile(r"(?<![A-Za-z0-9_])\.(?=[A-Za-z0-9._-]*[A-Za-z0-9])[A-Za-z0-9._-]+")
+
+# `app.js.map` / `index.d.ts` / `v2.3.1` / `foo/bar-baz.py` 等
+_FILELIKE_RE = re.compile(
+    r"(?<![A-Za-z0-9_])"
+    r"(?=[A-Za-z0-9._/+:-]*[A-Za-z])"
+    r"(?=[A-Za-z0-9._/+:-]*[._/+:-])"
+    r"[A-Za-z0-9](?:[A-Za-z0-9._/+:-]*[A-Za-z0-9])?"
+    r"(?![A-Za-z0-9_])"
+)
+
+# 参与“中英混排边界补空格”的 token：必须至少含 1 个拉丁字母，或本身就是受保护 token
+_LATINISH = rf"(?:{_PROT}|(?=[A-Za-z0-9._/+:-]*[A-Za-z])[A-Za-z0-9][A-Za-z0-9._/+:-]*)"
+
+# 零宽字符
+_ZERO_WIDTH_RE = re.compile(r"[\u200b-\u200d\ufeff]")
+
+
+# ---------------------------
+# 主函数
+# ---------------------------
+
+def normalize_tts_text(text: str) -> str:
+    """对 TTS 输入做鲁棒性正则化。"""
+    text = _base_cleanup(text)
+    text = _normalize_markdown_and_lines(text)
+    text, protected = _protect_spans(text)
+
+    text = _normalize_spaces(text)
+    text = _normalize_structural_punctuation(text)
+    text = _normalize_repeated_punctuation(text)
+    text = _normalize_spaces(text)
+
+    text = _restore_spans(text, protected)
+    return text.strip()
+
+
+# ---------------------------
+# 具体规则
+# ---------------------------
+
+def _base_cleanup(text: str) -> str:
+    text = text.replace("\r\n", "\n").replace("\r", "\n").replace("\u3000", " ")
+    text = _ZERO_WIDTH_RE.sub("", text)
+
+    cleaned = []
+    for ch in text:
+        cat = unicodedata.category(ch)
+        if ch in "\n\t " or not cat.startswith("C"):
+            cleaned.append(ch)
+    return "".join(cleaned)
+
+
+def _normalize_markdown_and_lines(text: str) -> str:
+    # Markdown 链接：[text](url) -> text url
+    text = re.sub(r"\[([^\[\]]+?)\]\((https?://[^)\s]+)\)", r"\1 \2", text)
+
+    lines = []
+    for raw in text.splitlines():
+        line = raw.strip()
+        if not line:
+            continue
+
+        line = re.sub(r"^#{1,6}\s+", "", line)   # 标题
+        line = re.sub(r"^>\s+", "", line)        # 引用
+        line = re.sub(r"^[-*+]\s+", "", line)    # 无序列表
+        line = re.sub(r"^\d+[.)]\s+", "", line)  # 有序列表
+        lines.append(line)
+
+    return "。".join(lines) if lines else ""
+
+
+def _protect_spans(text: str) -> tuple[str, list[str]]:
+    protected: list[str] = []
+
+    def repl(match: re.Match[str]) -> str:
+        idx = len(protected)
+        protected.append(match.group(0))
+        return f"___PROT{idx}___"
+
+    for pattern in (
+        _URL_RE,
+        _EMAIL_RE,
+        _MENTION_RE,
+        _REDDIT_RE,
+        _HASHTAG_RE,
+        _DOT_TOKEN_RE,
+        _FILELIKE_RE,
+    ):
+        text = pattern.sub(repl, text)
+
+    return text, protected
+
+
+def _restore_spans(text: str, protected: list[str]) -> str:
+    for idx, original in enumerate(protected):
+        text = text.replace(f"___PROT{idx}___", original)
+    return text
+
+
+def _normalize_spaces(text: str) -> str:
+    # 统一空白
+    text = re.sub(r"[ \t\r\f\v]+", " ", text)
+
+    # 汉字 / 日文片段内部：删除空格
+    text = re.sub(rf"({_CJK})\s+(?={_CJK})", r"\1", text)
+
+    # 汉字 / 日文 与纯数字之间：删除空格（不强行保留）
+    text = re.sub(rf"({_CJK})\s+(?=\d)", r"\1", text)
+    text = re.sub(rf"(\d)\s+(?={_CJK})", r"\1", text)
+
+    # 汉字 / 日文 与拉丁字母类 token / protected token 相邻：保留或补 1 个空格
+    text = re.sub(rf"({_CJK})(?=({_LATINISH}))", r"\1 ", text)
+    text = re.sub(rf"(({_LATINISH}))(?={_CJK})", r"\1 ", text)
+
+    # 再压一遍连续空格
+    text = re.sub(r" {2,}", " ", text)
+
+    # 中文标点前后不保留空格
+    text = re.sub(r"\s+([，。！？；：、”’」』】）》])", r"\1", text)
+    text = re.sub(r"([（【「『《“‘])\s+", r"\1", text)
+    text = re.sub(r"([，。！？；：、])\s*", r"\1", text)
+
+    # ASCII 标点前不留空格；后面的英文空格不强改
+    text = re.sub(r"\s+([,.;!?])", r"\1", text)
+
+    return re.sub(r" {2,}", " ", text).strip()
+
+
+def _normalize_structural_punctuation(text: str) -> str:
+    # 结构性括号：在“结构位置”解包并转成句边界
+    # 连续块要支持收敛，因此做两轮
+    for _ in range(2):
+        text = re.sub(
+            r"(^|[。！？!?；;]\s*)[【〖『「]([^】〗』」]+)[】〗』」]\s*",
+            r"\1\2。",
+            text,
+        )
+
+    # 《》只处理独立标题，不处理嵌入式标题
+    # 例：重磅。《新品发布》——现在开始！ -> 重磅。新品发布。现在开始！
+    text = re.sub(
+        r"(^|[。！？!?；;]\s*)《([^》]+)》(?=\s*(?:___PROT\d+___|[—–―-]{2,}|$|[。！？!?；;，,]))",
+        r"\1\2",
+        text,
+    )
+
+    # 流程 / 映射箭头：转成中文逗号，保留链路结构但避免把 `->` 原样喂给 TTS。
+    text = re.sub(
+        r"\s*(?:<[-=]+>|[-=]+>|<[-=]+|[→←↔⇒⇐⇔⟶⟵⟷⟹⟸⟺↦↤↪↩])\s*",
+        "，",
+        text,
+    )
+
+    # 长破折号 / 多连字符：转句边界
+    text = re.sub(r"\s*(?:—|–|―|-){2,}\s*", "。", text)
+
+    return text
+
+
+def _normalize_repeated_punctuation(text: str) -> str:
+    # 省略号 / 连续句点
+    text = re.sub(r"(?:\.{3,}|…{2,}|……+)", "。", text)
+
+    # 同类重复标点
+    text = re.sub(r"[。．]{2,}", "。", text)
+    text = re.sub(r"[，,]{2,}", "，", text)
+    text = re.sub(r"[!！]{2,}", "！", text)
+    text = re.sub(r"[?？]{2,}", "？", text)
+
+    # 混合问叹号：收敛到 ？！
+    def _mixed_qe(match: re.Match[str]) -> str:
+        s = match.group(0)
+        has_q = any(ch in s for ch in "?？")
+        has_e = any(ch in s for ch in "!！")
+        if has_q and has_e:
+            return "？！"
+        return "？" if has_q else "！"
+
+    text = re.sub(r"[!?！？]{2,}", _mixed_qe, text)
+    return text
+
+
+# ---------------------------
+# 测试
+# ---------------------------
+
+TEST_CASES = [
+    # 1) .map / dot-leading token / 文件名 / 版本号
+    (
+        "dot_map_sentence",
+        "2026 年 3 月 31 日，安全研究员 Chaofan Shou (@Fried_rice) 发现 Anthropic 的 npm 包中暴露了 .map 文件，",
+        "2026年3月31日，安全研究员 Chaofan Shou (@Fried_rice) 发现 Anthropic 的 npm 包中暴露了 .map 文件，",
+    ),
+    ("dot_tokens", "别把 .env、.npmrc、.gitignore 提交上去。", "别把 .env、.npmrc、.gitignore 提交上去。"),
+    ("file_names", "请检查 bundle.min.js、package.json 和 processing_moss_tts.py。", "请检查 bundle.min.js、package.json 和 processing_moss_tts.py。"),
+    ("index_d_ts", "index.d.ts 里也有同样的问题。", "index.d.ts 里也有同样的问题。"),
+    ("version_build", "Bug 的讨论可以精确到 v2.3.1 (Build 15)。", "Bug 的讨论可以精确到 v2.3.1 (Build 15)。"),
+    ("version_rc", "3.0.0-rc.1 还不能上生产。", "3.0.0-rc.1 还不能上生产。"),
+    ("jar_name", "fabric-api-0.91.3+1.20.2.jar 需要单独下载。", "fabric-api-0.91.3+1.20.2.jar 需要单独下载。"),
+
+    # 2) URL / Email / mention / hashtag / Reddit
+    ("url", "仓库地址是 https://github.com/instructkr/claude-code", "仓库地址是 https://github.com/instructkr/claude-code"),
+    ("email", "联系邮箱：ops+tts@example.ai", "联系邮箱：ops+tts@example.ai"),
+    ("mention", "@Fried_rice 说这是 source map 暴露。", "@Fried_rice 说这是 source map 暴露。"),
+    ("reddit", "去 r/singularity 看讨论。", "去 r/singularity 看讨论。"),
+    ("hashtag_chain", "#张雪峰#张雪峰[话题]#张雪峰事件", "#张雪峰#张雪峰[话题]#张雪峰事件"),
+    ("mention_hashtag_boundary", "关注@biscuit0228_��转发#thetime_tbs", "关注 @biscuit0228_ 并转发 #thetime_tbs"),
+
+    # 3) bracket / 控制 token：保留，不删除
+    ("speaker_bracket", "[S1]你好。[S2]收到。", "[S1]你好。[S2]收到。"),
+    ("event_bracket", "请模仿 {whisper} 的语气说“别出声”。", "请模仿 {whisper} 的语气说“别出声”。"),
+    ("order_bracket", "订单号：[AB-1234-XYZ]", "订单号：[AB-1234-XYZ]"),
+
+    # 4) 结构性符号：替换成边界，而不是直接删除
+    ("struct_headline", "〖重磅〗《新品发布》——现在开始！", "重磅。新品发布。现在开始！"),
+    ("struct_notice", "【公告】今天 20:00 维护——预计 30 分钟。", "公告。今天20:00维护。预计30分钟。"),
+    ("struct_quote_chain", "『特别提醒』「不要外传」", "特别提醒。不要外传。"),
+    ("flow_arrow_chain", "请求接入 -> 身份与策略判定 -> 域服务处理", "请求接入，身份与策略判定，域服务处理"),
+    ("flow_arrow_no_space", "A->B", "A，B"),
+    ("flow_arrow_unicode", "配置中心→推理编排→运行时执行", "配置中心，推理编排，运行时执行"),
+    (
+        "flow_arrow_maas_example",
+        "MaaS 主链遵循请求接入 -> 身份与策略判定 -> 域服务处理 -> 推理编排 -> 运行时执行的在线数据面结构。Dashboard 不直接承载单次在线推理请求，而是负责统一配置、统一治理、统一运营和统一展示。",
+        "MaaS 主链遵循请求接入，身份与策略判定，域服务处理，推理编排，运行时执行的在线数据面结构。Dashboard 不直接承载单次在线推理请求，而是负责统一配置、统一治理、统一运营和统一展示。",
+    ),
+
+    # 5) 嵌入式标题：保留
+    ("embedded_title", "我喜欢《哈姆雷特》这本书。", "我喜欢《哈姆雷特》这本书。"),
+
+    # 6) 重复标点 / 社交噪声
+    ("noise_qe", "真的假的？？？！！！", "真的假的？！"),
+    ("noise_ellipsis", "这个包把 app.js.map 也发上去了......太离谱了！！！", "这个包把 app.js.map 也发上去了。太离谱了！"),
+    ("noise_ellipsis_cn", "【系统提示】请模仿{sad}低沉语气，说“今天下雨了……”", "系统提示。请模仿{sad}低沉语气，说“今天下雨了。”"),
+
+    # 7) 空格规则：英文压缩、中文删除、中英混排保留边界
+    ("english_spaces", "This   is   a   test.", "This is a test."),
+    ("chinese_spaces", "这 是　一 段  含有多种空白的文本。", "这是一段含有多种空白的文本。"),
+    ("mixed_spaces_1", "这是Anthropic的npm包", "这是 Anthropic 的 npm 包"),
+    ("mixed_spaces_2", "今天update到v2.3.1了", "今天 update 到 v2.3.1 了"),
+    ("mixed_spaces_3", "处理app.js.map文件", "处理 app.js.map 文件"),
+
+    # 8) Markdown / 列表 / 换行
+    ("markdown_link", "详情见 [release note](https://github.com/example/release)", "详情见 release note https://github.com/example/release"),
+    ("markdown_heading", "# I made a free open source app to help with markdown files", "I made a free open source app to help with markdown files"),
+    ("list_lines", "- 修复 .map 泄露\n- 发布 v2.3.1", "修复 .map 泄露。发布 v2.3.1"),
+    ("numbered_lines", "1. 安装依赖\n2. 运行测试\n3. 发布 v2.3.1", "安装依赖。运行测试。发布 v2.3.1"),
+    ("newlines", "第一行\n第二行\n第三行", "第一行。第二行。第三行"),
+
+    # 9) 零宽字符 / 幂等性
+    ("zero_width_url", "详见 https://x.com/\u200bSafety", "详见 https://x.com/Safety"),
+]
+
+
+def run_tests(verbose: bool = True) -> None:
+    failed = []
+
+    for name, text, expected in TEST_CASES:
+        actual = normalize_tts_text(text)
+        if actual != expected:
+            failed.append((name, text, expected, actual))
+            continue
+
+        # 幂等性：第二次归一化不应继续改动结果
+        second = normalize_tts_text(actual)
+        if second != actual:
+            failed.append((name + "_idempotence", actual, actual, second))
+
+    if failed:
+        lines = ["\nTEST FAILED:\n"]
+        for name, text, expected, actual in failed:
+            lines.append(f"[{name}]")
+            lines.append(f"input   : {text}")
+            lines.append(f"expected: {expected}")
+            lines.append(f"actual  : {actual}")
+            lines.append("")
+        raise AssertionError("\n".join(lines))
+
+    if verbose:
+        print(f"All {len(TEST_CASES)} tests passed.")
+
+
+if __name__ == "__main__":
+    run_tests()