init

.gitattributes

@@ -1,35 +1,2 @@
-*.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
-*.wasm 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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|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,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,83 @@
+{
+  "architectures": [
+    "MossSpeechForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "audio_pad_token_id": 512,
+  "audio_vocab_size": 16512,
+  "auto_map": {
+    "AutoConfig": "configuration_moss_speech.MossSpeechConfig",
+    "AutoModel": "modeling_moss_speech.MossSpeechForCausalLM"
+  },
+  "channels": 2,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "eosp_token_id": 16384,
+  "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": null,
+  "modality_pad_token_id": 151667,
+  "model_type": "moss_speech",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "num_modality_layers": 4,
+  "num_shared_layers": 32,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "sosp_token_id": 151646,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.0.dev0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151680,
+  "vocab_size_list": [
+    151680,
+    16512
+  ]
+}

configuration_moss_speech.py

@@ -0,0 +1,258 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/moss_speech/modular_moss_speech.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_moss_speech.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
+# Copyright 2025 the HuggingFace 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.
+
+
+from transformers.configuration_utils import PretrainedConfig, layer_type_validation
+from transformers.modeling_rope_utils import rope_config_validation
+
+
+class MossSpeechConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossSpeechModel`]. It is used to instantiate a
+    MossSpeech model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    MossSpeech-8B [Qwen/MossSpeech-8B](https://huggingface.co/Qwen/MossSpeech-8B).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151680):
+            Vocabulary size of the MossSpeech model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`MossSpeechModel`].
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 12288):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 36):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 40960):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*):
+            Sliding window attention (SWA) window size. Only used when `use_sliding_window=True`.
+        max_window_layers (`int`, *optional*):
+            Number of initial layers using full attention. Layers after `max_window_layers` (if any) use SWA.
+        layer_types (`list`, *optional*):
+            Attention pattern for each layer.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+        audio_vocab_size (`int`, *optional*):
+            Vocabulary size for the audio branch used by the audio language modeling head.
+        modality_pad_token_id (`int`, *optional*, defaults to 0):
+            Token id used to pad the inactive modality channel when generating.
+        sosp_token_id (`int`, *optional*):
+            Start-of-speech special token id used by the processor and generation logic.
+        eosp_token_id (`int`, *optional*):
+            End-of-speech special token id used by the processor and generation logic.
+        num_shared_layers (`int`, *optional*, defaults to 32):
+            Number of shared decoder layers applied before splitting into modality-specific blocks.
+        num_modality_layers (`int`, *optional*, defaults to 4):
+            Number of decoder layers per modality after the shared block.
+
+    ```python
+    >>> from transformers import MossSpeechModel, MossSpeechConfig
+
+    >>> # Initializing a MossSpeech style configuration
+    >>> configuration = MossSpeechConfig()
+
+    >>> # Initializing a model from the MossSpeech-8B style configuration
+    >>> model = MossSpeechModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "moss_speech"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `MossSpeech`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151680,
+        hidden_size=4096,
+        intermediate_size=12288,
+        num_hidden_layers=36,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=40960,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=None,
+        max_window_layers=None,
+        layer_types=None,
+        attention_dropout=0.0,
+        audio_vocab_size=None,
+        modality_pad_token_id=0,
+        sosp_token_id=None,
+        eosp_token_id=None,
+        num_shared_layers=32,
+        num_modality_layers=4,
+        **kwargs,
+    ):
+        assert num_shared_layers + num_modality_layers == num_hidden_layers, (
+            f"num_shared_layers ({num_shared_layers}) + num_modality_layers ({num_modality_layers}) must equal to num_hidden_layers ({num_hidden_layers})"
+        )
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if self.use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        self.audio_vocab_size = audio_vocab_size
+        self.modality_pad_token_id = int(modality_pad_token_id)
+        self.sosp_token_id = None if sosp_token_id is None else int(sosp_token_id)
+        self.eosp_token_id = None if eosp_token_id is None else int(eosp_token_id)
+        self.num_shared_layers = int(num_shared_layers)
+        self.num_modality_layers = int(num_modality_layers)
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                (
+                    "sliding_attention"
+                    if self.sliding_window is not None and i >= self.max_window_layers
+                    else "full_attention"
+                )
+                for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["MossSpeechConfig"]

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.0.dev0"
+}

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,454 @@
+{
+  "metadata": {
+    "total_parameters": 9095694336,
+    "total_size": 18191388672
+  },
+  "weight_map": {
+    "audio_lm_head.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.0.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.1.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.2.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_block.layers.3.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.audio_embed.weight": "model-00001-of-00004.safetensors",
+    "model.audio_norm.weight": "model-00004-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.20.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.21.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.22.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.shared_block.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.shared_block.layers.9.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.text_block.layers.0.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.0.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.1.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.2.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.3.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.text_block.layers.3.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.text_norm.weight": "model-00004-of-00004.safetensors",
+    "text_lm_head.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_moss_speech.py

@@ -0,0 +1,1132 @@
+# coding=utf-8
+# Copyright 2025 OpenMOSS and HuggingFace Inc. teams. 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 copy
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Union
+
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+from transformers.generation.stopping_criteria import StoppingCriteriaList
+from transformers.generation.streamers import BaseStreamer
+from transformers.generation.utils import GenerateDecoderOnlyOutput
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import (
+    ModelOutput,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.pytorch_utils import isin_mps_friendly
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import check_model_inputs
+
+from .configuration_moss_speech import MossSpeechConfig
+
+
+@dataclass
+class MossSpeechModelOutputWithPast(ModelOutput):
+    """MossSpeech model output that includes per-modality last_hidden_state dict
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed
+            or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with
+            each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`.
+        last_hidden_state_dict (`Dict[str, torch.FloatTensor]`, *optional*):
+            Dictionary containing hidden states for each modality. Keys are modality names (e.g., "text", "audio")
+            and values are tensors of shape `(batch_size, sequence_length, hidden_size)`.
+        attentions (`tuple(torch.FloatTensor)`, *optional*):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape
+            `(batch_size, num_heads, sequence_length, sequence_length)`.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    past_key_values: Optional[tuple] = None
+    hidden_states: Optional[tuple] = None
+    attentions: Optional[tuple] = None
+    last_hidden_state_dict: Optional[dict[str, torch.Tensor]] = None
+    past_key_values_dict: Optional[dict] = None
+
+
+@dataclass
+class MossSpeechCausalLMOutputWithPast(ModelOutput):
+    """MossSpeech causal language modeling output, includes per-modality hidden_states dict
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        hidden_states (`Dict[str, torch.FloatTensor]`, *optional*):
+            Dictionary containing hidden states for each modality. Keys are modality names (e.g., "text", "audio")
+            and values are tensors of shape `(batch_size, sequence_length, hidden_size)`.
+        attentions (`tuple(torch.FloatTensor)`, *optional*):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
+    """
+
+    hidden_states: Optional[dict[str, torch.Tensor]] = None
+    past_key_values: Optional[tuple] = None
+    attentions: Optional[tuple] = None
+    last_hidden_state_dict: Optional[dict[str, torch.Tensor]] = None
+    audio_loss: Optional[torch.FloatTensor] = None
+    audio_logits: Optional[torch.FloatTensor] = None
+    text_loss: Optional[torch.FloatTensor] = None
+    text_logits: Optional[torch.FloatTensor] = None
+    text_hidden_states: Optional[torch.FloatTensor] = None
+    audio_hidden_states: Optional[torch.FloatTensor] = None
+    logits_all: Optional[tuple] = None
+    past_key_values_dict: Optional[dict] = None
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class MossSpeechRMSNorm(nn.Module):
+    # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MossSpeech
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """Root Mean Square LayerNorm used in MossSpeech."""
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class MossSpeechMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    # Copied from transformers.models.llama.modeling_llama.rotate_half
+    """Rotate half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+    """Apply Rotary Position Embeddings to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*): Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            Dimension along which to unsqueeze cos[position_ids] and sin[position_ids] for broadcasting to `q`/`k`.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the rotated query and key tensors.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    # Copied from transformers.models.llama.modeling_llama.repeat_kv
+    """Repeat key/value heads to match attention heads.
+
+    Equivalent to `torch.repeat_interleave(x, dim=1, repeats=n_rep)` transforming
+    (batch, num_key_value_heads, seqlen, head_dim) -> (batch, num_attention_heads, seqlen, head_dim).
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    # Copied from transformers.models.llama.modeling_llama.eager_attention_forward
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class MossSpeechAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need'."""
+
+    def __init__(self, config: MossSpeechConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = MossSpeechRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = MossSpeechRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class MossSpeechDecoderLayer(GradientCheckpointingLayer):
+    """Single decoder layer used in the MossSpeech transformer."""
+
+    def __init__(self, config: MossSpeechConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = MossSpeechAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = MossSpeechMLP(config)
+        self.input_layernorm = MossSpeechRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MossSpeechRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MossSpeechTransformerBlock(nn.Module):
+    """A contiguous stack of decoder layers that handles attention types and cache offsets."""
+
+    def __init__(self, config: MossSpeechConfig, start_idx: int, num_layers: int):
+        super().__init__()
+        self.start_idx = start_idx
+        self.num_layers = num_layers
+        self.layers = nn.ModuleList(
+            [MossSpeechDecoderLayer(config, layer_idx=start_idx + i) for i in range(num_layers)]
+        )
+        self.config = config
+        self.has_sliding_layers = "sliding_attention" in config.layer_types
+
+    def _mask_for_layer(self, causal_masks: dict[str, torch.Tensor], layer_type: str) -> torch.Tensor:
+        return causal_masks[("sliding_attention" if layer_type == "sliding_attention" else "full_attention")]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        causal_masks: dict[str, torch.Tensor],
+        position_ids: torch.LongTensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        past_key_values: Optional[Cache],
+        use_cache: bool,
+        cache_position: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        for i, decoder_layer in enumerate(self.layers):
+            layer_type = self.config.layer_types[self.start_idx + i]
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=self._mask_for_layer(causal_masks, layer_type),
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        return hidden_states
+
+
+@auto_docstring
+class MossSpeechPreTrainedModel(PreTrainedModel):
+    config: MossSpeechConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MossSpeechDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": MossSpeechDecoderLayer,
+        "attentions": MossSpeechAttention,
+    }
+
+
+class MossSpeechRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: MossSpeechConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        seq_len = position_ids_expanded.max().item() + 1
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = int(seq_len)
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class MossSpeechModel(MossSpeechPreTrainedModel):
+    """The bare MossSpeech decoder-only Transformer with shared and modality-specific blocks.
+
+    This model outputs a dictionary of per-modality hidden states in addition to the shared hidden states.
+    """
+
+    def __init__(self, config: MossSpeechConfig):
+        super().__init__(config)
+        pad_idx = (
+            config.pad_token_id if getattr(config, "pad_token_id", None) is not None else config.modality_pad_token_id
+        )
+        self.padding_idx = pad_idx
+        self.vocab_size = config.vocab_size
+
+        # Shared embed + trunk
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=self.padding_idx)
+        self.audio_embed = nn.Embedding(config.audio_vocab_size, config.hidden_size)
+        self.rotary_emb = MossSpeechRotaryEmbedding(config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in config.layer_types
+
+        self.shared_block = MossSpeechTransformerBlock(config, start_idx=0, num_layers=config.num_shared_layers)
+
+        # Modality-specific blocks (route after the shared trunk)
+        self.text_block = MossSpeechTransformerBlock(
+            config,
+            start_idx=0,
+            num_layers=config.num_modality_layers,
+        )
+        self.audio_block = MossSpeechTransformerBlock(
+            config,
+            start_idx=0,
+            num_layers=config.num_modality_layers,
+        )
+
+        # Modality-specific final norms
+        self.text_norm = MossSpeechRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.audio_norm = MossSpeechRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # register
+        self.modality_blocks = nn.ModuleDict(
+            {
+                "text": self.text_block,
+                "audio": self.audio_block,
+            }
+        )
+        self.modality_norms = nn.ModuleDict(
+            {
+                "text": self.text_norm,
+                "audio": self.audio_norm,
+            }
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _route_block(self, modality: str):
+        """Return the Transformer block for the given modality."""
+        if modality in self.modality_blocks:
+            return self.modality_blocks[modality]
+        else:
+            raise ValueError(f"Unknown modality: {modality}")
+
+    def _route_final_layer_norm(self, modality: str):
+        """Return the final layer normalization for the given modality."""
+        if modality in self.modality_norms:
+            return self.modality_norms[modality]
+        else:
+            raise ValueError(f"Unknown modality: {modality}")
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        modalities: list[str],
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[dict] = None,
+        past_key_values_dict: Optional[dict] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MossSpeechModelOutputWithPast:
+        r"""
+        Args:
+            modalities (List[str]): Modalities to compute in this forward pass.
+            past_key_values_dict (Optional[dict]): KV cache per block when using `use_cache=True`.
+        """
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("require input_ids (packed or text-only)")
+
+            if input_ids.dim() == 3 and input_ids.shape[1] == 2:
+                text_ids = input_ids[:, 0, :]
+                audio_ids_packed = input_ids[:, 1, :]
+                pad_id = int(getattr(self.config, "modality_pad_token_id", 0))
+                if pad_id == 0:
+                    raise ValueError("Expected a non-zero modality_pad_token_id for packed inputs")
+                # Safely replace pad tokens before embedding lookup
+                text_ids_safe = text_ids.masked_fill(text_ids == pad_id, 0)
+
+                text_embeds = self.embed_tokens(text_ids_safe)
+                audio_embeds = self.audio_embed(audio_ids_packed)
+                selection_mask = (text_ids != pad_id).unsqueeze(-1).to(dtype=text_embeds.dtype)
+                inputs_embeds = text_embeds * selection_mask + audio_embeds * (1 - selection_mask)
+            elif input_ids.dim() == 2:
+                # text-only input
+                inputs_embeds = self.embed_tokens(input_ids)
+            else:
+                raise ValueError("invalid inputs for embedding construction")
+
+        past_key_values_dict = {} if past_key_values_dict is None else dict(past_key_values_dict)
+
+        shared_cache: Optional[Cache] = None
+        if isinstance(past_key_values, dict):
+            # Accept dictionaries passed via either `past_key_values` or `past_key_values_dict`.
+            past_key_values_dict.update(past_key_values)
+            shared_cache = past_key_values.get("shared")
+        elif isinstance(past_key_values, Cache):
+            shared_cache = past_key_values
+        elif past_key_values is not None:
+            raise TypeError("past_key_values must be either a Cache instance, a dict of caches, or None")
+
+        if shared_cache is None:
+            shared_cache = past_key_values_dict.get("shared")
+
+        if use_cache:
+            if shared_cache is None:
+                shared_cache_config = copy.deepcopy(self.config)
+                shared_cache_config.layer_types = shared_cache_config.layer_types[: self.config.num_shared_layers]
+                shared_cache = DynamicCache(config=shared_cache_config)
+
+            past_key_values_dict = dict(past_key_values_dict)
+            past_key_values_dict["shared"] = shared_cache
+
+            for modality in self.modality_blocks:
+                cache = past_key_values_dict.get(modality)
+                if not isinstance(cache, Cache):
+                    modality_config = copy.deepcopy(self.config)
+                    modality_config.layer_types = modality_config.layer_types[-self.config.num_modality_layers :]
+                    past_key_values_dict[modality] = DynamicCache(config=modality_config)
+        else:
+            shared_cache = None
+
+        if cache_position is None:
+            past_seen_tokens = 0
+            if isinstance(shared_cache, Cache):
+                past_seen_tokens = shared_cache.get_seq_length()
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            # For mask creation, we use the shared cache or None
+            mask_past_key_values = shared_cache if isinstance(shared_cache, Cache) else None
+            # attention_mask = None
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": mask_past_key_values,
+                "position_ids": position_ids,
+            }
+
+            # Create the masks
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+            }
+            # The sliding window alternating layers are not always activated depending on the config
+            if self.has_sliding_layers:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # shared block
+        hidden_states = self.shared_block(
+            hidden_states,
+            causal_masks=causal_mask_mapping,
+            position_ids=position_ids,
+            position_embeddings=position_embeddings,
+            past_key_values=shared_cache if use_cache else None,
+            use_cache=bool(use_cache),
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # Compute hidden states for each modality
+        last_hidden_state_dict = {}
+        for modality in modalities:
+            mod_block = self._route_block(modality)
+            mod_norm = self._route_final_layer_norm(modality)
+            mod_cache = past_key_values_dict.get(modality) if use_cache else None
+
+            # Build modality-specific hidden_states starting from the same shared input
+            mod_hidden_states = mod_block(
+                hidden_states,
+                causal_masks=causal_mask_mapping,
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+                past_key_values=mod_cache,
+                use_cache=bool(use_cache),
+                cache_position=cache_position,
+                **kwargs,
+            )
+            mod_hidden_states = mod_norm(mod_hidden_states)
+            last_hidden_state_dict[modality] = mod_hidden_states
+
+        return MossSpeechModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            last_hidden_state_dict=last_hidden_state_dict,
+            past_key_values=shared_cache if use_cache else None,
+            past_key_values_dict=past_key_values_dict if use_cache else None,
+        )
+
+
+class MossSpeechGenerationMixin(GenerationMixin):
+    """Generation mixin for MossSpeech model with two-channel (text/audio) support."""
+
+    def _setup_processors(self, generation_config: GenerationConfig, modalities: int) -> list[LogitsProcessorList]:
+        """Setup per-channel logits processors based on the generation config."""
+        realprocessor = [LogitsProcessorList() for _ in range(modalities)]
+
+        if hasattr(generation_config, "layers"):
+            for i, layer_config in enumerate(generation_config.layers):
+                if i >= len(realprocessor):
+                    break
+
+                if layer_config.get("repetition_penalty") is not None:
+                    realprocessor[i].append(
+                        RepetitionPenaltyLogitsProcessor(penalty=layer_config.get("repetition_penalty"))
+                    )
+                if layer_config.get("temperature") is not None:
+                    realprocessor[i].append(TemperatureLogitsWarper(temperature=layer_config.get("temperature")))
+                if layer_config.get("top_k") is not None:
+                    realprocessor[i].append(TopKLogitsWarper(top_k=layer_config.get("top_k")))
+                if layer_config.get("top_p") is not None:
+                    realprocessor[i].append(TopPLogitsWarper(top_p=layer_config.get("top_p")))
+
+        return realprocessor
+
+    def _generate_next_tokens_with_scores(
+        self,
+        logits_all: tuple[torch.Tensor, ...],
+        input_ids: torch.LongTensor,
+        realprocessor: list[LogitsProcessorList],
+        do_samples: list[bool],
+        generation_config: GenerationConfig,
+        generating_length: int,
+    ) -> tuple[torch.LongTensor, tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]:
+        """Generate next tokens for all channels with scores and logits."""
+        # Get next token logits
+        next_token_logits = tuple(logits[:, -1, :].clone().float().to(input_ids.device) for logits in logits_all)
+
+        # Apply audio-channel-specific constraints
+        next_token_logits[1][:, 16385:] = -torch.inf
+        if hasattr(generation_config, "min_new_tokens") and generating_length < generation_config.min_new_tokens:
+            next_token_logits[1][:, 16384] = -torch.inf
+
+        # Process logits
+        next_token_scores = tuple(
+            realprocessor[i](input_ids[:, :, i], logits) for i, logits in enumerate(next_token_logits)
+        )
+
+        # Sample or select tokens
+        next_tokens = []
+        for i, channel_score in enumerate(next_token_scores):
+            if do_samples[i]:
+                channel_ntk = torch.multinomial(nn.functional.softmax(channel_score, dim=-1), num_samples=1).squeeze(1)
+            else:
+                channel_ntk = torch.argmax(channel_score, dim=-1)
+            next_tokens.append(channel_ntk)
+
+        return torch.stack(next_tokens, dim=-1), next_token_scores, next_token_logits
+
+    def _process_multi_modality_tokens(
+        self,
+        next_tokens: torch.LongTensor,
+        current_modality: torch.Tensor,
+        modality_pad_token: Union[int, torch.Tensor],
+    ) -> torch.LongTensor:
+        """Process tokens for MossSpeech generation."""
+
+        mask = current_modality == 1
+        if mask.any():
+            pad_value = modality_pad_token.item() if torch.is_tensor(modality_pad_token) else modality_pad_token
+            next_tokens[mask, 0] = int(pad_value)
+
+        return next_tokens
+
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool,
+        streamer: Optional[BaseStreamer],
+        **model_kwargs,
+    ) -> Union[GenerateDecoderOnlyOutput, torch.LongTensor]:
+        """Sampling implementation for MossSpeech with text and audio modalities."""
+
+        # Determine the pad token used to mask the text channel when the audio modality is active.
+        modality_pad_token = generation_config._pad_token_tensor
+        if modality_pad_token is None:
+            pad_fallback = getattr(self.config, "modality_pad_token_id", None)
+            if pad_fallback is None:
+                pad_fallback = getattr(self.config, "pad_token_id", None)
+            if pad_fallback is None:
+                raise ValueError(
+                    "MossSpeech generation requires a pad token id; please set it on the config or generation config."
+                )
+            modality_pad_token = torch.tensor(
+                pad_fallback,
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+        else:
+            modality_pad_token = modality_pad_token.to(device=input_ids.device, dtype=input_ids.dtype)
+
+        audio_pad_token_id = getattr(self.config, "audio_pad_token_id", None)
+        sosp_token_id = getattr(self.config, "sosp_token_id", None)
+        eosp_token_id = getattr(self.config, "eosp_token_id", None)
+
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        global_do_sample = generation_config.do_sample
+
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        batch_size, cur_len, input_modalities = input_ids.shape
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+
+        model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+
+        if hasattr(generation_config, "do_samples") and generation_config.do_samples is not None:
+            per_modality_do_sample = generation_config.do_samples
+            logits_processors = self._setup_processors(generation_config, input_modalities)
+        else:
+            per_modality_do_sample = [global_do_sample for _ in range(input_modalities)]
+            logits_processors = [logits_processor for _ in range(input_modalities)]
+
+        current_modality = torch.zeros((batch_size,), dtype=torch.long, device=input_ids.device)
+
+        # Infer starting modality: text channel uses index 0, audio channel index 1
+        if self.config.modality_pad_token_id is not None:
+            audio_mask = input_ids[:, -1, 0] == self.config.modality_pad_token_id
+            current_modality[audio_mask] = 1
+        if audio_pad_token_id is not None:
+            text_mask = input_ids[:, -1, 1] == audio_pad_token_id
+            current_modality[text_mask] = 0
+
+        generating_length = 0
+        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+            generating_length += 1
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            if output_attentions:
+                model_inputs["output_attentions"] = output_attentions
+            if output_hidden_states:
+                model_inputs["output_hidden_states"] = output_hidden_states
+            if "past_key_values_dict" in model_kwargs:
+                model_inputs["past_key_values_dict"] = model_kwargs["past_key_values_dict"]
+
+            if sosp_token_id is not None:
+                text_mode_mask = current_modality == 0
+                text_to_audio_mask = text_mode_mask & (input_ids[:, -1, 0] == sosp_token_id)
+                current_modality[text_to_audio_mask] = 1
+            if eosp_token_id is not None:
+                audio_mode_mask = current_modality == 1
+                audio_to_text_mask = audio_mode_mask & (input_ids[:, -1, 1] == eosp_token_id)
+                current_modality[audio_to_text_mask] = 0
+
+            outputs = self(**model_inputs, return_dict=True)
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder=False)
+            if outputs.past_key_values_dict is not None:
+                model_kwargs["past_key_values_dict"] = outputs.past_key_values_dict
+
+            if synced_gpus and this_peer_finished:
+                continue
+
+            next_tokens, next_token_scores, next_token_logits = self._generate_next_tokens_with_scores(
+                outputs.logits_all,
+                input_ids,
+                logits_processors,
+                per_modality_do_sample,
+                generation_config,
+                generating_length,
+            )
+            next_tokens = self._process_multi_modality_tokens(
+                next_tokens,
+                current_modality,
+                modality_pad_token,
+            )
+
+            input_ids = torch.cat([input_ids, next_tokens[:, None, :]], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens[:, 0].cpu())
+
+            stopping = stopping_criteria(input_ids[:, :, 0], scores)
+
+            unfinished_sequences = unfinished_sequences & ~stopping
+            this_peer_finished = unfinished_sequences.max() == 0
+
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (outputs.attentions,)
+                if output_hidden_states:
+                    decoder_hidden_states += (outputs.hidden_states,)
+
+            cur_len += 1
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                attentions=decoder_attentions,
+                hidden_states=decoder_hidden_states,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+
+        return input_ids
+
+    def generate(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        output_only: bool = True,
+        **kwargs,
+    ):
+        batch_size, seq_len, modalities = input_ids.shape
+        start_id = seq_len
+        outputs = super().generate(input_ids, **kwargs)
+        return_dict_in_generate = kwargs.get("return_dict_in_generate", False)
+        if return_dict_in_generate:
+            output_ids = outputs["sequences"]
+        else:
+            output_ids = outputs
+        if output_only:
+            output_ids = output_ids[:, start_id:, :]
+        if return_dict_in_generate:
+            outputs["sequences"] = output_ids
+        else:
+            outputs = output_ids
+        return outputs
+
+    def _prepare_attention_mask_for_generation(
+        self,
+        inputs_tensor: torch.Tensor,
+        generation_config: GenerationConfig,
+        model_kwargs: dict[str, Any],
+    ) -> torch.LongTensor:
+        pad_token_id = generation_config._pad_token_tensor
+        eos_token_id = generation_config._eos_token_tensor
+
+        # `input_ids` may be present in the model kwargs, instead of being the main input (e.g. multimodal model)
+        if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+            inputs_tensor = model_kwargs["input_ids"]
+
+        # No information for attention mask inference -> return default attention mask
+        if len(inputs_tensor.shape) == 3 and inputs_tensor.shape[1] == 2:
+            # For (B, 2, seq_len) inputs, create a (B, seq_len) attention mask
+            default_attention_mask = torch.ones(
+                (inputs_tensor.shape[0], inputs_tensor.shape[2]),  # (B, seq_len)
+                dtype=torch.long,
+                device=inputs_tensor.device,
+            )
+        else:
+            # Keep the original logic for (B, seq_len) inputs
+            default_attention_mask = torch.ones(
+                inputs_tensor.shape[:2],  # (B, seq_len)
+                dtype=torch.long,
+                device=inputs_tensor.device,
+            )
+        if pad_token_id is None:
+            return default_attention_mask
+
+        is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+        if not is_input_ids:
+            return default_attention_mask
+
+        is_pad_token_in_inputs = (pad_token_id is not None) and (
+            isin_mps_friendly(elements=inputs_tensor, test_elements=pad_token_id).any()
+        )
+        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or ~(
+            isin_mps_friendly(elements=eos_token_id, test_elements=pad_token_id).any()
+        )
+        can_infer_attention_mask = is_pad_token_in_inputs * is_pad_token_not_equal_to_eos_token_id
+        attention_mask_from_padding = inputs_tensor.ne(pad_token_id).long()
+
+        attention_mask = (
+            attention_mask_from_padding * can_infer_attention_mask + default_attention_mask * ~can_infer_attention_mask
+        )
+        return attention_mask
+
+
+@auto_docstring
+class MossSpeechForCausalLM(MossSpeechPreTrainedModel, MossSpeechGenerationMixin):
+    _tied_weights_keys = ["text_lm_head.weight", "audio_lm_head.weight"]
+    _tp_plan = {"text_lm_head": "colwise_rep", "audio_lm_head": "colwise_rep"}
+    _pp_plan = {
+        "text_lm_head": (["hidden_states"], ["logits"]),
+        "audio_lm_head": (["hidden_states"], ["logits"]),
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MossSpeechModel(config)
+        self.vocab_size = config.vocab_size
+        self.text_lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.audio_lm_head = nn.Linear(config.hidden_size, config.audio_vocab_size, bias=False)
+        self.modality_lm_head = {
+            "text": self.text_lm_head,
+            "audio": self.audio_lm_head,
+        }
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    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[torch.Tensor] = None,
+        past_key_values_dict: Optional[dict] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MossSpeechCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, 2, sequence_length)` or `(batch_size, 2*sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        past_key_values_dict (Optional[dict]): KV cache for each block.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, MossSpeechForCausalLM
+
+        >>> model = MossSpeechForCausalLM.from_pretrained("Qwen/MossSpeech-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/MossSpeech-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```
+        """
+        if input_ids is not None and input_ids.dim() == 2:
+            B = input_ids.shape[0]
+            input_ids = input_ids.reshape([B, 2, -1])
+
+        input_ids = input_ids.transpose(1, 2)
+        outputs: MossSpeechModelOutputWithPast = self.model(
+            modalities=["text", "audio"],
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            past_key_values_dict=past_key_values_dict,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        text_hidden_states = outputs.last_hidden_state_dict["text"]
+        audio_hidden_states = outputs.last_hidden_state_dict["audio"]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        text_logits = self.modality_lm_head["text"](text_hidden_states[:, slice_indices, :])
+        audio_logits = self.modality_lm_head["audio"](audio_hidden_states[:, slice_indices, :])
+
+        text_loss = None
+        audio_loss = None
+        if labels is not None:
+            if labels.dim() == 2:
+                B = labels.shape[0]
+                labels.reshape([B, 2, -1])
+            text_labels = labels[:, 0, :]
+            audio_labels = labels[:, 1, :]
+            text_loss = self.loss_function(
+                logits=text_logits,
+                labels=text_labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+            audio_loss = self.loss_function(
+                logits=audio_logits,
+                labels=audio_labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+        return MossSpeechCausalLMOutputWithPast(
+            text_loss=text_loss,
+            audio_loss=audio_loss,
+            text_logits=text_logits,
+            audio_logits=audio_logits,
+            logits_all=(text_logits, audio_logits),
+            past_key_values=outputs.past_key_values,
+            past_key_values_dict=outputs.past_key_values_dict,
+            text_hidden_states=text_hidden_states,
+            audio_hidden_states=audio_hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class MossSpeechForSequenceClassification(GenericForSequenceClassification, MossSpeechPreTrainedModel):
+    pass
+
+
+class MossSpeechForTokenClassification(GenericForTokenClassification, MossSpeechPreTrainedModel):
+    pass
+
+
+class MossSpeechForQuestionAnswering(GenericForQuestionAnswering, MossSpeechPreTrainedModel):
+    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`
+
+
+__all__ = [
+    "MossSpeechForCausalLM",
+    "MossSpeechForQuestionAnswering",
+    "MossSpeechPreTrainedModel",
+    "MossSpeechModel",
+    "MossSpeechForSequenceClassification",
+    "MossSpeechForTokenClassification",
+]

processing_moss_speech.py

@@ -0,0 +1,419 @@
+# coding=utf-8
+# Copyright 2025 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.
+"""Processor class for MossSpeech."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Mapping, Optional, Sequence, Union
+
+from transformers import AutoTokenizer
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import OptionalDependencyNotAvailable, is_torch_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+if is_torch_available():
+    import torch
+else:  # pragma: no cover
+    torch = None  # type: ignore[assignment]
+
+MossSpeechCodec = None
+_MOSS_CODEC_IMPORT_ERROR: Optional[Exception] = None
+
+_TEXT_PLACEHOLDER_TOKEN_ID = 151667
+_AUDIO_PAD_TOKEN_ID = 512
+_SOSP_TOKEN_ID = 151646
+_EOSP_TOKEN_ID = 16384
+
+
+def _ensure_dependencies() -> None:
+    global MossSpeechCodec, _MOSS_CODEC_IMPORT_ERROR
+
+    if torch is None:
+        raise OptionalDependencyNotAvailable()
+
+    if MossSpeechCodec is None:
+        try:
+            from MossSpeechCodec import MossSpeechCodec as _Codec  # type: ignore
+        except ImportError as import_error:  # pragma: no cover
+            _MOSS_CODEC_IMPORT_ERROR = import_error
+            raise OptionalDependencyNotAvailable() from import_error
+        else:
+            MossSpeechCodec = _Codec
+            _MOSS_CODEC_IMPORT_ERROR = None
+
+
+class MossSpeechProcessorKwargs(ProcessingKwargs, total=False):
+    """Default keyword argument groups supported by :class:`MossSpeechProcessor`."""
+
+    _defaults = {
+        "common_kwargs": {
+            "return_tensors": "pt",
+            "padding": True,
+        }
+    }
+
+
+@dataclass
+class MossSpeechChatSample:
+    """Per-sample MossSpeech input with two-channel token grid."""
+
+    input_ids_2d: "torch.LongTensor"
+    label_ids_2d: Optional["torch.LongTensor"] = None
+
+
+@dataclass
+class MossSpeechBatchInput:
+    """Batched MossSpeech tensors returned by the processor."""
+
+    input_ids: "torch.LongTensor"
+    attention_mask: "torch.LongTensor"
+    labels: Optional["torch.LongTensor"] = None
+
+
+@dataclass
+class MossSpeechResponse:
+    """Decoded MossSpeech output item containing text and optional audio."""
+
+    audio: Optional["torch.Tensor"] = None
+    generated_text: str = ""
+    sampling_rate: Optional[int] = None
+
+
+@dataclass
+class _MossSpeechInputSegment:
+    """Internal helper representing either text or audio tokens."""
+
+    text: Optional[str] = None
+    audio_tokens: Optional["torch.Tensor"] = None
+    tokenized_text: Optional["torch.Tensor"] = None
+
+    def __post_init__(self) -> None:
+        if self.text is None and self.tokenized_text is None:
+            raise ValueError("A segment must provide text or tokenized_text content.")
+        if self.audio_tokens is not None and self.audio_tokens.dim() != 1:
+            raise ValueError("`audio_tokens` must be a 1D tensor of codec token ids.")
+
+    def to_tensor(self, tokenizer: AutoTokenizer) -> "torch.Tensor":
+        _ensure_dependencies()
+        if self.tokenized_text is None:
+            tokenized = tokenizer(
+                self.text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=999999,
+                padding=False,
+                add_special_tokens=False,
+            )["input_ids"].to(dtype=torch.long)
+        else:
+            tokenized = self.tokenized_text.unsqueeze(0).to(dtype=torch.long)
+
+        if self.audio_tokens is None:
+            audio_channel = torch.full_like(tokenized, _AUDIO_PAD_TOKEN_ID)
+            return torch.cat([tokenized, audio_channel], dim=0)
+
+        audio_tokens = self.audio_tokens.reshape(1, -1).to(dtype=torch.long)
+        text_channel = torch.full((1, audio_tokens.shape[1]), _TEXT_PLACEHOLDER_TOKEN_ID, dtype=torch.long)
+
+        sosp = torch.tensor([[_SOSP_TOKEN_ID]], dtype=torch.long)
+        text_pad = torch.tensor([[_TEXT_PLACEHOLDER_TOKEN_ID]], dtype=torch.long)
+        eosp = torch.tensor([[_EOSP_TOKEN_ID]], dtype=torch.long)
+        audio_pad = torch.tensor([[_AUDIO_PAD_TOKEN_ID]], dtype=torch.long)
+
+        text_channel = torch.cat([sosp, text_channel, text_pad], dim=1)
+        audio_channel = torch.cat([audio_pad, audio_tokens, eosp], dim=1)
+        return torch.cat([text_channel, audio_channel], dim=0)
+
+
+class MossSpeechSampleProcessor:
+    """Formats a structured conversation into MossSpeech grid tokens."""
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        audio_codec: MossSpeechCodec,
+        default_system_prompts: Mapping[str, str],
+    ) -> None:
+        self.tokenizer = tokenizer
+        self.audio_codec = audio_codec
+        self.default_system_prompts = default_system_prompts
+
+    def prepare_sample(
+        self,
+        conversation: Sequence[Mapping[str, Any]],
+        output_modality: str,
+    ) -> MossSpeechChatSample:
+        _ensure_dependencies()
+        if len(conversation) == 0:
+            raise ValueError("`conversation` must contain at least one turn.")
+
+        segments: list[_MossSpeechInputSegment] = []
+        for turn in conversation:
+            role = turn.get("role")
+            if role not in {"user", "assistant", "system"}:
+                raise ValueError(f"Unsupported role `{role}` detected.")
+
+            segments.append(_MossSpeechInputSegment(text=f"<|im_start|>{role}\n"))
+            content = turn.get("content")
+            if isinstance(content, Mapping):
+                audio_path = content.get("path")
+                if audio_path is None:
+                    raise ValueError("Audio turn content must include a `path` entry.")
+                encoded = self.audio_codec.encode([audio_path])[0]
+                segments.append(_MossSpeechInputSegment(audio_tokens=torch.tensor(encoded, dtype=torch.long)))
+            else:
+                segments.append(_MossSpeechInputSegment(text=str(content)))
+            segments.append(_MossSpeechInputSegment(text="<|im_end|>\n"))
+
+        if conversation[0].get("role") != "system":
+            system_prompt = self.default_system_prompts.get(output_modality)
+            if system_prompt is None:
+                raise KeyError(f"Missing default system prompt for modality `{output_modality}`.")
+            segments.extend(
+                [
+                    _MossSpeechInputSegment(text="<|im_start|>system\n"),
+                    _MossSpeechInputSegment(text=system_prompt),
+                    _MossSpeechInputSegment(text="<|im_end|>\n"),
+                ]
+            )
+
+        if output_modality == "text":
+            segments.append(_MossSpeechInputSegment(text="<|im_start|>assistant\n"))
+        elif output_modality == "audio":
+            segments.append(_MossSpeechInputSegment(text="<|im_start|>assistant\n<|object_ref_start|>"))
+        else:
+            raise NotImplementedError("Supported modalities are `text` and `audio`.")
+
+        input_tensors = [segment.to_tensor(self.tokenizer) for segment in segments]
+        input_ids = torch.cat(input_tensors, dim=1)
+        return MossSpeechChatSample(input_ids_2d=input_ids)
+
+    def collate(
+        self,
+        samples: Sequence[MossSpeechChatSample],
+        *,
+        pad_token_id: int,
+        audio_pad_token_id: int,
+    ) -> MossSpeechBatchInput:
+        _ensure_dependencies()
+        if len(samples) == 0:
+            raise ValueError("`samples` must not be empty.")
+
+        channel_count = samples[0].input_ids_2d.shape[0]
+        max_length = max(sample.input_ids_2d.shape[1] for sample in samples)
+
+        padded_inputs: list["torch.Tensor"] = []
+        attention_masks: list["torch.Tensor"] = []
+
+        for sample in samples:
+            seq_len = sample.input_ids_2d.shape[1]
+            pad_len = max_length - seq_len
+
+            pad_grid = torch.full((channel_count, pad_len), audio_pad_token_id, dtype=torch.long)
+            pad_grid[0] = pad_grid[0].fill_(pad_token_id)
+            padded_inputs.append(torch.cat([pad_grid, sample.input_ids_2d], dim=1))
+
+            attention_prefix = torch.zeros(pad_len, dtype=torch.long)
+            attention_body = torch.ones(seq_len, dtype=torch.long)
+            attention_masks.append(torch.cat([attention_prefix, attention_body], dim=0))
+
+        input_ids = torch.stack(padded_inputs).permute(0, 2, 1)
+        attention_mask = torch.stack(attention_masks)
+
+        return MossSpeechBatchInput(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,
+        )
+
+
+class MossSpeechProcessor(ProcessorMixin):
+    r"""Combines MossSpeech tokenizer and codec for unified text/audio processing."""
+
+    tokenizer_class = "AutoTokenizer"
+    audio_codec_class = "PreTrainedModel"
+    attributes = ["tokenizer", "audio_codec"]
+
+    def __init__(
+        self,
+        tokenizer,
+        audio_codec,
+        **kwargs,
+    ) -> None:
+        _ensure_dependencies()
+        super().__init__(tokenizer=tokenizer, audio_codec=audio_codec, **kwargs)
+        self.default_system_prompts = {
+            "text": "You are a helpful assistant. Respond with text outputs.",
+            "audio": "You are a helpful assistant. Respond with spoken outputs.",
+        }
+        self.sample_processor = MossSpeechSampleProcessor(
+            tokenizer=self.tokenizer,
+            audio_codec=self.audio_codec,
+            default_system_prompts=self.default_system_prompts,
+        )
+        self.sosp_token_id = _SOSP_TOKEN_ID
+        self.eosp_token_id = _EOSP_TOKEN_ID
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike[str]],
+        trust_remote_code: bool = True,
+        **kwargs: Any,
+    ) -> "MossSpeechProcessor":
+        kwargs.pop("_from_auto", None)
+        codec_path = kwargs.pop("codec_path", None)
+        if codec_path is None:
+            raise ValueError("`codec_path` must be supplied to load the MossSpeech codec.")
+
+        device = kwargs.pop("device", None) or "cpu"
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        _ensure_dependencies()
+        audio_codec = MossSpeechCodec.from_pretrained(codec_path, trust_remote_code=True).to(device)
+        return cls(tokenizer=tokenizer, audio_codec=audio_codec)
+
+    def __call__(
+        self,
+        data: Union[Mapping[str, Any], Sequence[Sequence[Mapping[str, Any]]]],
+        output_modalities: Union[str, Sequence[str]],
+        **kwargs: Any,
+    ) -> BatchEncoding:
+        _ensure_dependencies()
+        if isinstance(data, Mapping):
+            samples: list[Sequence[Mapping[str, Any]]] = [data]  # type: ignore[list-item]
+        elif isinstance(data, Sequence):
+            if len(data) == 0:
+                raise ValueError("`data` must contain at least one sample.")
+            if all(isinstance(turn, Mapping) for turn in data):
+                samples = [data]  # type: ignore[list-item]
+            else:
+                samples = list(data)  # type: ignore[list-item]
+        else:
+            raise TypeError("`data` must be a conversation dictionary or a sequence of conversations.")
+
+        if isinstance(output_modalities, str):
+            output_modalities = [output_modalities] * len(samples)
+        elif len(output_modalities) != len(samples):
+            raise ValueError("`output_modalities` length must match number of samples.")
+
+        merged_kwargs = self._merge_kwargs(MossSpeechProcessorKwargs, **kwargs)
+        common_kwargs = merged_kwargs["common_kwargs"]
+        padding = common_kwargs.get("padding", True)
+        if not padding:
+            raise NotImplementedError("Only padded batches are currently supported.")
+        return_tensors = common_kwargs.get("return_tensors", "pt")
+
+        chat_samples = [
+            self.sample_processor.prepare_sample(conversation, modality)
+            for conversation, modality in zip(samples, output_modalities)
+        ]
+
+        pad_token_id = self.tokenizer.pad_token_id
+        if pad_token_id is None:
+            raise ValueError("Tokenizer must define `pad_token_id` for MossSpeech processing.")
+
+        batch_inputs = self.sample_processor.collate(
+            chat_samples,
+            pad_token_id=pad_token_id,
+            audio_pad_token_id=_AUDIO_PAD_TOKEN_ID,
+        )
+        payload = {key: value for key, value in asdict(batch_inputs).items() if value is not None}
+        return BatchEncoding(payload, tensor_type=return_tensors)
+
+    def decode(
+        self,
+        token_ids: "torch.Tensor",
+        output_modalities: Union[str, Sequence[str]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> list[MossSpeechResponse]:
+        _ensure_dependencies()
+        if token_ids.dim() != 3:
+            raise ValueError("`token_ids` must be shaped as (batch, sequence_length, channels).")
+
+        if isinstance(output_modalities, str):
+            output_modalities = [output_modalities] * token_ids.shape[0]
+        elif len(output_modalities) != token_ids.shape[0]:
+            raise ValueError("`output_modalities` length must equal the batch size.")
+
+        if token_ids.shape[0] != 1:
+            raise NotImplementedError("Batch decoding is not yet implemented for MossSpeech.")
+
+        responses: list[MossSpeechResponse] = []
+        for batch_index, modality in enumerate(output_modalities):
+            tokens = token_ids[batch_index].int().cpu()
+            if tokens.shape[1] != 2:
+                tokens = tokens.transpose(0, 1)
+            if tokens.shape[0] != 2:
+                raise ValueError("Decoded tensor must contain exactly two channels (text and audio).")
+
+            if modality == "audio":
+                prefix = torch.tensor([[_SOSP_TOKEN_ID], [_AUDIO_PAD_TOKEN_ID]], dtype=torch.long)
+                tokens = torch.cat([prefix, tokens], dim=1)
+
+            text_channel = tokens[0, :-1]
+            audio_channel = tokens[1, :-1]
+            decoded_text = (
+                self.tokenizer.decode(text_channel, skip_special_tokens=True)
+                .replace("<|empty|>", ".")
+                .replace("<|end_empty|>", ":")
+            )
+
+            sosp_indices = (text_channel == self.sosp_token_id).nonzero(as_tuple=True)[0]
+            eosp_indices = (audio_channel == self.eosp_token_id).nonzero(as_tuple=True)[0]
+
+            waveform: Optional["torch.Tensor"] = None
+            if len(sosp_indices) > 0:
+                start_idx = sosp_indices[0].item() + 1
+                stop_idx = eosp_indices[0].item() if len(eosp_indices) > 0 else text_channel.shape[0]
+                audio_tokens = tokens[:, start_idx:stop_idx]
+                flattened_audio_tokens = audio_tokens[1].reshape(-1).tolist()
+
+                continuation = "".join(f"<{token}>" for token in flattened_audio_tokens)
+                codec_tokens = [int(match) for match in re.findall(r"(\d+)>", continuation)]
+                codec_tensor = torch.tensor(codec_tokens, dtype=torch.long).reshape(1, 1, -1)
+
+                prompt_path = kwargs.get("decoder_audio_prompt_path")
+                if prompt_path is None:
+                    raise ValueError("`decoder_audio_prompt_path` must be provided to decode audio outputs.")
+                codec_output = self.audio_codec.decode(codec_tensor, prompt_speech=prompt_path)
+                waveform = codec_output["syn_wav_list"][0].reshape(1, -1).detach().cpu()
+
+            responses.append(
+                MossSpeechResponse(
+                    audio=waveform,
+                    generated_text=decoded_text,
+                    sampling_rate=24000 if waveform is not None else None,
+                )
+            )
+
+        return responses
+
+
+__all__ = [
+    "MossSpeechProcessor",
+    "MossSpeechProcessorKwargs",
+    "MossSpeechResponse",
+]

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|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "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:9c5ae00e602b8860cbd784ba82a8aa14e8feecec692e7076590d014d7b7fdafa
+size 11421896

tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "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": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "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": "<|video_pad|>",
+      "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": "<|fim_pad|>",
+      "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
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}