Upload 16 files

adapter_config.json

@@ -0,0 +1,29 @@
+{
+    "alpha_pattern": {},
+    "auto_mapping": null,
+    "base_model_name_or_path": "ibm-granite/granite-speech-3.2-8b",
+    "bias": "none",
+    "fan_in_fan_out": false,
+    "inference_mode": true,
+    "init_lora_weights": true,
+    "layer_replication": null,
+    "layers_pattern": null,
+    "layers_to_transform": null,
+    "loftq_config": {},
+    "lora_alpha": 32,
+    "lora_dropout": 0.0,
+    "megatron_config": null,
+    "megatron_core": "megatron.core",
+    "modules_to_save": null,
+    "peft_type": "LORA",
+    "r": 64,
+    "rank_pattern": {},
+    "revision": null,
+    "target_modules": [
+        "q_proj",
+        "v_proj"
+    ],
+    "task_type": "CAUSAL_LM",
+    "use_dora": false,
+    "use_rslora": false
+}

adapter_model.safetensors

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

added_tokens.json

@@ -0,0 +1,6 @@
+{
+  "<|audio|>": 49155,
+  "<|end_of_role|>": 49153,
+  "<|start_of_role|>": 49152,
+  "<|tool_call|>": 49154
+}

configuration_granite_speech.py

@@ -0,0 +1,126 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.auto import CONFIG_MAPPING, AutoConfig
+
+
+class GraniteSpeechEncoderConfig(PretrainedConfig):
+    model_type = "granite_speech_encoder"
+
+    def __init__(
+        self,
+        input_dim=160,
+        num_layers=10,
+        hidden_dim=1024,
+        feedforward_mult=4,
+        num_heads=8,
+        dim_head=128,
+        output_dim=42,
+        context_size=200,
+        dropout=0.1,
+        conv_kernel_size=15,
+        conv_expansion_factor=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.feedforward_mult = feedforward_mult
+        self.num_heads = num_heads
+        self.dim_head = dim_head
+        self.output_dim = output_dim
+        self.context_size = context_size
+        self.dropout = dropout
+        self.conv_kernel_size = conv_kernel_size
+        self.conv_expansion_factor = conv_expansion_factor
+
+
+## adapted from transformers.models.blip.configuration_blip_2.Blip2VisionConfig
+class GraniteSpeechProjectorConfig(PretrainedConfig):
+    model_type = "granite_speech_qformer"
+
+    def __init__(
+        self,
+        llm_dim=4096,
+        downsample_rate=5,
+        window_size=15,
+        hidden_size=1024,
+        num_attention_heads=16,
+        intermediate_size=4096,
+        num_hidden_layers=2,
+        encoder_hidden_size=1024,
+        cross_attention_frequency=1,
+        max_position_embeddings=2048,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        position_embedding_type="absolute",
+        use_qformer_text_input=False,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.cross_attention_frequency = cross_attention_frequency
+        self.encoder_hidden_size = encoder_hidden_size
+        self.use_qformer_text_input = use_qformer_text_input
+        self.downsample_rate = downsample_rate
+        self.window_size = window_size
+        self.llm_dim = llm_dim
+
+
+class GraniteSpeechConfig(PretrainedConfig):
+    model_type = "granite_speech"
+    sub_configs = {
+        "text_config": AutoConfig,
+        "encoder_config": GraniteSpeechEncoderConfig,
+        "projector_config": GraniteSpeechProjectorConfig,
+    }
+
+    def __init__(
+        self,
+        encoder_config=None,
+        text_config=None,
+        projector_config=None,
+        audio_token_index=49155,
+        initializer_range=0.02,
+        has_lora_adapter=True,
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "granite"
+            text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+        elif text_config is None:
+            text_config = CONFIG_MAPPING["granite"]()
+
+        if isinstance(projector_config, dict):
+            # TODO - In the future, we should make this generic.
+            projector_config = GraniteSpeechProjectorConfig(**projector_config)
+        elif projector_config is None:
+            projector_config = GraniteSpeechProjectorConfig()
+
+        if not isinstance(encoder_config, GraniteSpeechEncoderConfig):
+            encoder_config = {} if encoder_config is None else encoder_config
+            encoder_config = GraniteSpeechEncoderConfig(**encoder_config)
+
+        self.text_config = text_config
+        self.encoder_config = encoder_config
+        self.projector_config = projector_config
+        self.audio_token_index = audio_token_index
+        self.initializer_range = initializer_range
+        self.has_lora_adapter = has_lora_adapter
+        super().__init__(**kwargs)
+
+
+__all__ = ["GraniteSpeechEncoderConfig", "GraniteSpeechProjectorConfig", "GraniteSpeechConfig"]

feature_extraction_granite_speech.py

@@ -0,0 +1,118 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Feature extractor class for Speech Granite
+"""
+
+import math
+from typing import List, Optional
+
+from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
+from transformers.utils import is_torch_available, is_torchaudio_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+if is_torch_available():
+    import torch
+
+if is_torchaudio_available():
+    import torchaudio
+
+
+class GraniteSpeechFeatureExtractor(FeatureExtractionMixin):
+    model_input_names = ["input_features"]
+
+    def __init__(
+        self,
+        sampling_rate=16000,
+        n_fft=512,
+        win_length=400,
+        hop_length=160,
+        n_mels=80,
+        projector_window_size=15,
+        projector_downsample_rate=5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.melspec_kwargs = {
+            "sample_rate": sampling_rate,
+            "n_fft": n_fft,
+            "win_length": win_length,
+            "hop_length": hop_length,
+            "n_mels": n_mels,
+        }
+        # HACK - for now, lazily initialize the mel spectrogram transform;
+        # the feature extractor mixin explodes otherwise because
+        # it tries to log the feature extractor, and the melspectrogram
+        # transform isn't json serializable...
+        self.melspec = None
+        self.projector_window_size = projector_window_size
+        self.projector_downsample_rate = projector_downsample_rate
+
+    def _ensure_melspec_transform_is_initialized(self):
+        if self.melspec is None:
+            self.melspec = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs)
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        device: Optional[str] = "cpu",
+    ) -> BatchFeature:
+        # TODO there is probably a better way to do both of these things...
+        self._ensure_melspec_transform_is_initialized()
+        if device is not None:
+            melspec = self.melspec.to(device)
+            x = x.to(device)
+        else:
+            melspec = self.melspec
+
+        B, _ = x.shape
+        with torch.no_grad():
+            mel = melspec(x.float())
+            logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_()
+            mx = logmel.amax(dim=(-2, -1), keepdim=True)
+            logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1)
+            if logmel.shape[1] % 2 == 1:
+                logmel = logmel[:, :-1]  # remove last frame if odd
+            x = logmel.reshape(B, -1, 2 * logmel.shape[-1])  # stacking and skipping by 2
+
+        if x.device != "cpu":
+            return x.detach().cpu()
+        return x
+
+    def _get_num_audio_features(self, audio_lengths: List[int]) -> List[int]:
+        """
+        Gets the (variable length) variable length number of features
+        (i.e., projector output) for the sequences being considered.
+        """
+        hop_length = self.melspec_kwargs["hop_length"]
+        effective_window_size = self.projector_window_size // self.projector_downsample_rate
+
+        projector_lengths = []
+        for raw_length in audio_lengths:
+            # mel sequence length computation
+            mel_length = raw_length // hop_length + 1
+            # encoder frame takes two mel features
+            encoder_length = mel_length // 2
+            nblocks = math.ceil(encoder_length / self.projector_window_size)
+            # projector output length
+            projector_length = nblocks * effective_window_size
+            projector_lengths.append(projector_length)
+
+        return projector_lengths
+
+
+__all__ = ["GraniteSpeechFeatureExtractor"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 0,
+  "transformers_version": "4.50.0.dev0"
+}

model.safetensors.index.json

@@ -0,0 +1,762 @@
+{
+  "metadata": {
+    "total_size": 16966988452
+  },
+  "weight_map": {
+    "encoder.out.bias": "model-00009-of-00009.safetensors",
+    "encoder.out.weight": "model-00009-of-00009.safetensors",
+    "encoder.out_mid.bias": "model-00009-of-00009.safetensors",
+    "encoder.out_mid.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.1.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.10.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.2.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.3.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.4.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.5.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.6.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.7.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.8.post_norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.fn.rel_pos_emb.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.fn.to_kv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.fn.to_out.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.fn.to_out.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.fn.to_q.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.attn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.2.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.2.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.4.conv.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.5.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.5.num_batches_tracked": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.5.running_mean": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.5.running_var": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.5.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.7.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.conv.net.7.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff1.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.fn.net.0.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.fn.net.0.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.fn.net.3.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.fn.net.3.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.ff2.fn.norm.weight": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.post_norm.bias": "model-00009-of-00009.safetensors",
+    "encoder.rnn_tr.9.post_norm.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.embed_tokens.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.mlp.down_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.mlp.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.mlp.up_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.self_attn.k_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.self_attn.o_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.self_attn.q_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.0.self_attn.v_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.mlp.down_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.mlp.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.mlp.up_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.self_attn.k_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.self_attn.o_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.self_attn.q_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.1.self_attn.v_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.10.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.mlp.down_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.mlp.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.mlp.up_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.self_attn.k_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.self_attn.o_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.self_attn.q_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.10.self_attn.v_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.mlp.down_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.mlp.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.mlp.up_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.self_attn.k_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.self_attn.o_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.self_attn.q_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.11.self_attn.v_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.mlp.down_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.mlp.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.mlp.up_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.self_attn.k_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.self_attn.o_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.self_attn.q_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.12.self_attn.v_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.mlp.down_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.mlp.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.mlp.up_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.self_attn.k_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.self_attn.o_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.self_attn.q_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.13.self_attn.v_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.14.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.mlp.down_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.mlp.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.mlp.up_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.self_attn.k_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.self_attn.o_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.self_attn.q_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.14.self_attn.v_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.mlp.down_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.mlp.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.mlp.up_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.self_attn.k_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.self_attn.o_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.self_attn.q_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.15.self_attn.v_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.mlp.down_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.mlp.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.mlp.up_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.self_attn.k_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.self_attn.o_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.self_attn.q_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.16.self_attn.v_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.mlp.down_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.mlp.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.mlp.up_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.self_attn.k_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.self_attn.o_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.self_attn.q_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.17.self_attn.v_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.input_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.mlp.down_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.mlp.gate_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.mlp.up_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.post_attention_layernorm.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.self_attn.k_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.self_attn.o_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.self_attn.q_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.18.self_attn.v_proj.weight": "model-00004-of-00009.safetensors",
+    "language_model.model.layers.19.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.mlp.down_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.mlp.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.mlp.up_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.self_attn.k_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.self_attn.o_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.self_attn.q_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.19.self_attn.v_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.2.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.mlp.down_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.mlp.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.mlp.up_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.self_attn.k_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.self_attn.o_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.self_attn.q_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.2.self_attn.v_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.20.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.mlp.down_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.mlp.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.mlp.up_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.self_attn.k_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.self_attn.o_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.self_attn.q_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.20.self_attn.v_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.mlp.down_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.mlp.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.mlp.up_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.self_attn.k_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.self_attn.o_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.self_attn.q_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.21.self_attn.v_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.mlp.down_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.mlp.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.mlp.up_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.self_attn.k_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.self_attn.o_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.self_attn.q_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.22.self_attn.v_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.input_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.mlp.down_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.mlp.gate_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.mlp.up_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.post_attention_layernorm.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.self_attn.k_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.self_attn.o_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.self_attn.q_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.23.self_attn.v_proj.weight": "model-00005-of-00009.safetensors",
+    "language_model.model.layers.24.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.mlp.down_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.mlp.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.mlp.up_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.self_attn.k_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.self_attn.o_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.self_attn.q_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.24.self_attn.v_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.mlp.down_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.mlp.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.mlp.up_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.self_attn.k_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.self_attn.o_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.self_attn.q_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.25.self_attn.v_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.mlp.down_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.mlp.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.mlp.up_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.self_attn.k_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.self_attn.o_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.self_attn.q_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.26.self_attn.v_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.mlp.down_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.mlp.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.mlp.up_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.self_attn.k_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.self_attn.o_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.self_attn.q_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.27.self_attn.v_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.input_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.mlp.down_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.mlp.gate_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.mlp.up_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.post_attention_layernorm.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.self_attn.k_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.self_attn.o_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.self_attn.q_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.28.self_attn.v_proj.weight": "model-00006-of-00009.safetensors",
+    "language_model.model.layers.29.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.mlp.down_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.mlp.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.mlp.up_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.self_attn.k_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.self_attn.o_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.self_attn.q_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.29.self_attn.v_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.3.input_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.mlp.down_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.mlp.gate_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.mlp.up_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.post_attention_layernorm.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.self_attn.k_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.self_attn.o_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.self_attn.q_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.3.self_attn.v_proj.weight": "model-00001-of-00009.safetensors",
+    "language_model.model.layers.30.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.mlp.down_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.mlp.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.mlp.up_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.self_attn.k_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.self_attn.o_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.self_attn.q_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.30.self_attn.v_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.mlp.down_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.mlp.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.mlp.up_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.self_attn.k_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.self_attn.o_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.self_attn.q_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.31.self_attn.v_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.mlp.down_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.mlp.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.mlp.up_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.self_attn.k_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.self_attn.o_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.self_attn.q_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.32.self_attn.v_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.input_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.mlp.down_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.mlp.gate_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.mlp.up_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.post_attention_layernorm.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.self_attn.k_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.self_attn.o_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.self_attn.q_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.33.self_attn.v_proj.weight": "model-00007-of-00009.safetensors",
+    "language_model.model.layers.34.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.mlp.down_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.mlp.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.mlp.up_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.self_attn.k_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.self_attn.o_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.self_attn.q_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.34.self_attn.v_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.mlp.down_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.mlp.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.mlp.up_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.self_attn.k_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.self_attn.o_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.self_attn.q_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.35.self_attn.v_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.mlp.down_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.mlp.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.mlp.up_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.self_attn.k_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.self_attn.o_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.self_attn.q_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.36.self_attn.v_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.mlp.down_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.mlp.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.mlp.up_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.self_attn.k_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.self_attn.o_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.self_attn.q_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.37.self_attn.v_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.input_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.mlp.down_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.mlp.gate_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.mlp.up_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.post_attention_layernorm.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.self_attn.k_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.self_attn.o_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.self_attn.q_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.38.self_attn.v_proj.weight": "model-00008-of-00009.safetensors",
+    "language_model.model.layers.39.input_layernorm.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.mlp.down_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.mlp.gate_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.mlp.up_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.post_attention_layernorm.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.self_attn.k_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.self_attn.o_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.self_attn.q_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.39.self_attn.v_proj.weight": "model-00009-of-00009.safetensors",
+    "language_model.model.layers.4.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.mlp.down_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.mlp.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.mlp.up_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.self_attn.k_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.self_attn.o_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.self_attn.q_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.4.self_attn.v_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.mlp.down_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.mlp.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.mlp.up_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.self_attn.k_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.self_attn.o_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.self_attn.q_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.5.self_attn.v_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.mlp.down_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.mlp.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.mlp.up_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.self_attn.k_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.self_attn.o_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.self_attn.q_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.6.self_attn.v_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.mlp.down_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.mlp.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.mlp.up_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.self_attn.k_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.self_attn.o_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.self_attn.q_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.7.self_attn.v_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.input_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.mlp.down_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.mlp.gate_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.mlp.up_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.post_attention_layernorm.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.self_attn.k_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.self_attn.o_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.self_attn.q_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.8.self_attn.v_proj.weight": "model-00002-of-00009.safetensors",
+    "language_model.model.layers.9.input_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.mlp.down_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.mlp.gate_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.mlp.up_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.post_attention_layernorm.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.self_attn.k_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.self_attn.o_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.self_attn.q_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.layers.9.self_attn.v_proj.weight": "model-00003-of-00009.safetensors",
+    "language_model.model.norm.weight": "model-00009-of-00009.safetensors",
+    "projector.linear.bias": "model-00009-of-00009.safetensors",
+    "projector.linear.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.key.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.key.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.query.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.query.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.value.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.attention.value.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.output.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.output.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.output.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.attention.output.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.key.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.key.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.query.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.query.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.value.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.attention.value.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.output.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.output.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.output.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.crossattention.output.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.intermediate_query.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.intermediate_query.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.output_query.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.output_query.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.output_query.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.0.output_query.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.key.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.key.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.query.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.query.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.value.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.attention.value.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.output.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.output.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.output.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.attention.output.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.key.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.key.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.query.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.query.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.value.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.attention.value.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.output.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.output.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.output.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.crossattention.output.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.intermediate_query.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.intermediate_query.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.output_query.LayerNorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.output_query.LayerNorm.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.output_query.dense.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.encoder.layer.1.output_query.dense.weight": "model-00009-of-00009.safetensors",
+    "projector.qformer.layernorm.bias": "model-00009-of-00009.safetensors",
+    "projector.qformer.layernorm.weight": "model-00009-of-00009.safetensors",
+    "projector.query": "model-00009-of-00009.safetensors"
+  }
+}

modeling_granite_speech.py

@@ -0,0 +1,1393 @@
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import einsum, nn
+
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    ModelOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.auto import AutoModelForCausalLM
+from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
+from transformers.utils import (
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_peft_available,
+    logging,
+    replace_return_docstrings,
+)
+
+from .configuration_granite_speech import (
+    GraniteSpeechConfig,
+    GraniteSpeechEncoderConfig,
+    GraniteSpeechProjectorConfig,
+)
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "GraniteSpeechConfig"
+
+
+@dataclass
+class GraniteSpeechCausalLMOutputWithPast(ModelOutput):
+    """
+    Base class for LlavaNext causal language model (or autoregressive) outputs.
+
+    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 (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[List[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+### Projector
+# Currently, we copy the Qformer code directly to avoid depending on Blip2;
+# it would be better to create the model from config, similar to the LLM,
+# but to do this, we will need to register the QFormer model into an automodel,
+# which will should involve pulling it out into its own dir so that it is accessible
+# under transformers.models.X.
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerMultiHeadAttention with Blip2->GraniteSpeech
+class GraniteSpeechQFormerMultiHeadAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
+                % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
+        self.save_attention = False
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+
+    def get_attention_map(self):
+        return self.attention_map
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        mixed_query_layer = self.query(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            seq_length = hidden_states.size()[1]
+            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
+            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
+            distance = position_ids_l - position_ids_r
+            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
+            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        if is_cross_attention and self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->GraniteSpeechQFormer
+class GraniteSpeechQFormerSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerAttention with Blip2->GraniteSpeech
+class GraniteSpeechQFormerAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False):
+        super().__init__()
+        self.attention = GraniteSpeechQFormerMultiHeadAttention(config, is_cross_attention)
+        self.output = GraniteSpeechQFormerSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.attention.query = prune_linear_layer(self.attention.query, index)
+        self.attention.key = prune_linear_layer(self.attention.key, index)
+        self.attention.value = prune_linear_layer(self.attention.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
+        self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        self_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->GraniteSpeechQFormer
+class GraniteSpeechQFormerIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->GraniteSpeechQFormer
+class GraniteSpeechQFormerOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerLayer with Blip2->GraniteSpeech
+class GraniteSpeechQFormerLayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = GraniteSpeechQFormerAttention(config)
+
+        self.layer_idx = layer_idx
+
+        if layer_idx % config.cross_attention_frequency == 0:
+            self.crossattention = GraniteSpeechQFormerAttention(config, is_cross_attention=True)
+            self.has_cross_attention = True
+        else:
+            self.has_cross_attention = False
+
+        if config.use_qformer_text_input:
+            self.intermediate = GraniteSpeechQFormerIntermediate(config)
+            self.output = GraniteSpeechQFormerOutput(config)
+
+        self.intermediate_query = GraniteSpeechQFormerIntermediate(config)
+        self.output_query = GraniteSpeechQFormerOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        query_length=0,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+        outputs = self_attention_outputs[1:-1]
+
+        present_key_value = self_attention_outputs[-1]
+
+        if query_length > 0:
+            query_attention_output = attention_output[:, :query_length, :]
+
+            if self.has_cross_attention:
+                if encoder_hidden_states is None:
+                    raise ValueError("encoder_hidden_states must be given for cross-attention layers")
+                cross_attention_outputs = self.crossattention(
+                    query_attention_output,
+                    attention_mask,
+                    head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions=output_attentions,
+                )
+                query_attention_output = cross_attention_outputs[0]
+                # add cross attentions if we output attention weights
+                outputs = outputs + cross_attention_outputs[1:-1]
+
+            layer_output = apply_chunking_to_forward(
+                self.feed_forward_chunk_query,
+                self.chunk_size_feed_forward,
+                self.seq_len_dim,
+                query_attention_output,
+            )
+
+            if attention_output.shape[1] > query_length:
+                layer_output_text = apply_chunking_to_forward(
+                    self.feed_forward_chunk,
+                    self.chunk_size_feed_forward,
+                    self.seq_len_dim,
+                    attention_output[:, query_length:, :],
+                )
+                layer_output = torch.cat([layer_output, layer_output_text], dim=1)
+        else:
+            layer_output = apply_chunking_to_forward(
+                self.feed_forward_chunk,
+                self.chunk_size_feed_forward,
+                self.seq_len_dim,
+                attention_output,
+            )
+        outputs = (layer_output,) + outputs
+
+        outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+    def feed_forward_chunk_query(self, attention_output):
+        intermediate_output = self.intermediate_query(attention_output)
+        layer_output = self.output_query(intermediate_output, attention_output)
+        return layer_output
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerEncoder with Blip2->GraniteSpeech
+class GraniteSpeechQFormerEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList(
+            [GraniteSpeechQFormerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        query_length=0,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions else None
+
+        next_decoder_cache = () if use_cache else None
+
+        for i in range(self.config.num_hidden_layers):
+            layer_module = self.layer[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False) and self.training:
+                if use_cache:
+                    logger.warning(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer_module.__call__,
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    query_length,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                if layer_module.has_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class GraniteSpeechEncoderProjectorPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = GraniteSpeechProjectorConfig
+    base_model_prefix = "qformer"
+    supports_gradient_checkpointing = True
+
+    _no_split_modules = [
+        "GraniteSpeechQFormerMultiHeadAttention",
+        "T5Block",
+        "OPTDecoderLayer",
+    ]
+    _skip_keys_device_placement = "past_key_values"
+    _keep_in_fp32_modules = ["query_tokens"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        factor = self.config.initializer_range
+        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Embedding) or isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=factor)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class GraniteSpeechQFormerModel(GraniteSpeechEncoderProjectorPreTrainedModel):
+    """
+    Querying Transformer (Q-Former), used in GraniteSpeech.
+    """
+
+    def __init__(self, config: GraniteSpeechProjectorConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.encoder = GraniteSpeechQFormerEncoder(config)
+
+        self.post_init()
+
+    # Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerModel.get_input_embeddings
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    # Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerModel.set_input_embeddings
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    # Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerModel._prune_heads
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    # Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerModel.get_extended_attention_mask
+    def get_extended_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int],
+        device: torch.device,
+        has_query: bool = False,
+    ) -> torch.Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (`Tuple[int]`):
+                The shape of the input to the model.
+            device (`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+
+    # Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerModel.forward
+    def forward(
+        self,
+        query_embeds: torch.FloatTensor,
+        query_length: Optional[int] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of:
+            shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and
+            value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are
+            used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key
+            value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape
+            `(batch_size, sequence_length)`.
+        use_cache (`bool`, `optional`):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # past_key_values_length
+        past_key_values_length = (
+            past_key_values[0][0].shape[2] - self.config.query_length if past_key_values is not None else 0
+        )
+
+        query_length = (
+            query_length if query_length is not None else query_embeds.shape[1] if query_embeds is not None else 0
+        )
+
+        embedding_output = self.layernorm(query_embeds)
+        embedding_output = self.dropout(embedding_output)
+
+        input_shape = embedding_output.size()[:-1]
+        batch_size, seq_length = input_shape
+        device = embedding_output.device
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if isinstance(encoder_hidden_states, list):
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
+            else:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+            if isinstance(encoder_attention_mask, list):
+                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+            else:
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            query_length=query_length,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = sequence_output[:, 0, :]
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+# TODO (alex) - refactor GraniteSpeechQformer to be available under
+# transformers.models.X, delete all of the code above, and
+# create the model through AutoModel.
+
+
+class GraniteSpeechEncoderProjectorQFormer(nn.Module):
+    def __init__(self, config: GraniteSpeechProjectorConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.ds_rate = config.downsample_rate
+        self.window_size = config.window_size
+        self.num_queries = self.window_size // self.ds_rate
+        self.query = nn.Parameter(torch.zeros(1, self.num_queries, config.hidden_size))
+        self.query.data.normal_(mean=0.0, std=1.0)
+        # NOTE: It would be better to create this from config, similar to the LLM.
+        # To do this, we need to register the QFormer model into an automodel, which
+        # will require pulling it out into its own dir so that it's accessible under
+        # transformers.models.X
+        self.qformer = GraniteSpeechQFormerModel(config)
+        self.linear = nn.Linear(config.hidden_size, config.llm_dim)
+
+    def forward(self, x, atts):
+        batch_size, seq_len, dim = x.size()
+        nblocks = math.ceil(seq_len / self.window_size)
+        pad = nblocks * self.window_size - seq_len
+        x = nn.functional.pad(x, (0, 0, 0, pad), "constant", 0)
+        x = x.view(batch_size * nblocks, self.window_size, dim)
+
+        query_output = self.qformer(
+            query_embeds=self.query.data,
+            encoder_hidden_states=x,
+            encoder_attention_mask=atts,
+            return_dict=True,
+        )
+        query_proj = self.linear(
+            query_output.last_hidden_state.view(batch_size, nblocks * self.window_size // self.ds_rate, -1)
+        )
+        return query_proj
+
+
+### Encoder
+class GraniteSpeechCTCModel(nn.Module):
+    def __init__(self, config: GraniteSpeechEncoderConfig):
+        super(GraniteSpeechCTCModel, self).__init__()
+
+        self.rnn_tr = nn.ModuleList(
+            [nn.Linear(config.input_dim, config.hidden_dim, bias=True)]
+            + [
+                GraniteSpeechConformerBlock(
+                    dim=config.hidden_dim,
+                    dim_head=config.dim_head,
+                    heads=config.num_heads,
+                    ff_mult=config.feedforward_mult,
+                    conv_expansion_factor=config.conv_expansion_factor,
+                    conv_kernel_size=config.conv_kernel_size,
+                    context_size=config.context_size,  # attention context size
+                    attn_dropout=config.dropout,
+                    ff_dropout=config.dropout,
+                    conv_dropout=config.dropout,
+                )
+                for layer_idx in range(config.num_layers)
+            ]
+        )
+
+        self.out = nn.Linear(config.hidden_dim, config.output_dim, bias=True)
+        self.out_mid = nn.Linear(config.output_dim, config.hidden_dim, bias=True)
+        self.context_size = config.context_size
+        self.input_dim = config.input_dim
+        self.num_layers = config.num_layers
+        self.hidden_dim = config.hidden_dim
+        self.output_dim = config.output_dim
+
+    def forward(self, x: torch.Tensor):
+        x = self.rnn_tr[0](x)
+        for idx, layer in enumerate(self.rnn_tr[1:], start=1):
+            x = layer(x, self.context_size)
+            if idx == self.num_layers // 2:
+                x_mid = x.clone()
+                x_mid = self.out(x_mid)
+                x += self.out_mid(nn.Softmax(dim=-1)(x_mid))
+        return x
+
+
+# NOTE: Conformer adapated from: https://github.com/lucidrains/conformer.git
+class GraniteSpeechConformerPermute(nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+
+    def forward(self, x):
+        x = x.permute(self.dims)
+        return x
+
+
+class GraniteSpeechConformerDepthWiseConv1d(nn.Module):
+    def __init__(self, chan_in, chan_out, kernel_size, padding):
+        super().__init__()
+        self.padding = padding
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in, bias=False)
+
+    def forward(self, x):
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+class GraniteSpeechConformerScale(nn.Module):
+    def __init__(self, scale, fn):
+        super().__init__()
+        self.fn = fn
+        self.scale = scale
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+class GraniteSpeechConformerPreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+
+class GraniteSpeechConformerPreNormAttn(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, context_size, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, context_size, **kwargs)
+
+
+class GraniteSpeechConformerAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        context_size=200,
+        max_pos_emb=512,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.dim_head = dim_head
+        self.scale = dim_head**-0.5
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim)
+
+        self.max_pos_emb = max_pos_emb
+        self.rel_pos_emb = nn.Embedding(2 * max_pos_emb + 1, dim_head)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context_size):
+        device, h, max_pos_emb = x.device, self.heads, self.max_pos_emb
+        bs, n, d = x.shape
+        assert context_size > 0 and context_size <= max_pos_emb
+
+        nb = math.ceil(n / context_size)
+        nr = n % context_size
+        if nr > 0:
+            # right padding to reach block size
+            x = torch.nn.functional.pad(x, (0, 0, 0, context_size - nr))
+
+        q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim=-1))
+        q, k, v = [t.reshape(bs, nb, context_size, h, -1).transpose(2, 3) for t in (q, k, v)]
+
+        dots = einsum("b m h i d, b m h j d -> b m h i j", q, k) * self.scale
+
+        # shaw's relative positional embedding
+        seq = torch.arange(context_size, device=device)
+        dist = seq.view(-1, 1) - seq.view(1, -1)
+        dist = torch.clamp(dist, -context_size, context_size) + max_pos_emb
+        rel_pos_emb = self.rel_pos_emb(dist).to(q)
+        pos_attn = einsum("b m h c d, c r d -> b m h c r", q, rel_pos_emb) * self.scale
+        dots = dots + pos_attn
+
+        if nr > 0:
+            # masked attention in the extended block
+            mask = torch.ones(context_size, context_size, dtype=bool, device=device)
+            mask[:nr, :nr] = 0
+            mask_value = -torch.finfo(dots.dtype).max
+            dots[:, -1, :].masked_fill_(mask, mask_value)
+
+        attn = dots.softmax(dim=-1)
+
+        out = einsum("b m h i j, b m h j d -> b m h i d", attn, v)
+        out = out.transpose(2, 3).reshape(bs, x.shape[1], -1)
+        out = self.to_out(out[:, :n, :])
+        return self.dropout(out)
+
+
+class GraniteSpeechConformerFeedForward(nn.Module):
+    def __init__(self, dim, mult=4, dropout=0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult), nn.SiLU(), nn.Dropout(dropout), nn.Linear(dim * mult, dim), nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GraniteSpeechConformerConvModule(nn.Module):
+    def __init__(self, dim, causal=False, expansion_factor=2, kernel_size=31, dropout=0.0):
+        super().__init__()
+
+        inner_dim = dim * expansion_factor
+        padding = self.calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)
+
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            GraniteSpeechConformerPermute(dims=(0, 2, 1)),
+            nn.Conv1d(dim, inner_dim * 2, 1),
+            nn.GLU(dim=1),
+            GraniteSpeechConformerDepthWiseConv1d(inner_dim, inner_dim, kernel_size=kernel_size, padding=padding),
+            nn.BatchNorm1d(inner_dim) if not causal else nn.Identity(),
+            nn.SiLU(),
+            nn.Conv1d(inner_dim, dim, 1),
+            GraniteSpeechConformerPermute(dims=(0, 2, 1)),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+    @staticmethod
+    def calc_same_padding(kernel_size: int):
+        pad = kernel_size // 2
+        return (pad, pad - (kernel_size + 1) % 2)
+
+
+class GraniteSpeechConformerBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=64,
+        heads=8,
+        ff_mult=2,
+        conv_expansion_factor=2,
+        conv_kernel_size=31,
+        context_size=-1,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        conv_dropout=0.0,
+    ):
+        super().__init__()
+        self.ff1 = GraniteSpeechConformerFeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+        self.attn = GraniteSpeechConformerAttention(
+            dim=dim,
+            dim_head=dim_head,
+            heads=heads,
+            dropout=attn_dropout,
+            context_size=context_size,
+        )
+        self.conv = GraniteSpeechConformerConvModule(
+            dim=dim,
+            causal=False,
+            expansion_factor=conv_expansion_factor,
+            kernel_size=conv_kernel_size,
+            dropout=conv_dropout,
+        )
+        self.ff2 = GraniteSpeechConformerFeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+
+        self.attn = GraniteSpeechConformerPreNormAttn(dim, self.attn)
+        self.ff1 = GraniteSpeechConformerScale(0.5, GraniteSpeechConformerPreNorm(dim, self.ff1))
+        self.ff2 = GraniteSpeechConformerScale(0.5, GraniteSpeechConformerPreNorm(dim, self.ff2))
+
+        self.post_norm = nn.LayerNorm(dim)
+
+    def forward(self, x, context_size):
+        x = self.ff1(x) + x
+        x = self.attn(x, context_size) + x
+        x = self.conv(x) + x
+        x = self.ff2(x) + x
+        x = self.post_norm(x)
+        return x
+
+
+GRANITE_SPEECH_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config (`GraniteSpeechConfig`):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Granite Speech Model outputting raw hidden-states without any specific head on top.",
+    GRANITE_SPEECH_START_DOCSTRING,
+)
+class GraniteSpeechPreTrainedModel(PreTrainedModel):
+    config_class = GraniteSpeechConfig
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+GRANITE_SPEECH_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        input_features (`torch.FloatTensor` of shape `(batch_size, audio seq len, mel feat dim)):
+            The tensors corresponding to the input audios. input features can be obtained using
+            [`AutoFeatureExtractor`]. See [`GraniteSpeechFeatureExtractor.__call__`] for details.
+            [`GraniteSpeechProcessor`] uses [`GraniteSpeechFeatureExtractor`] for processing audio.
+        input_mask (`torch.Tensor`, *optional*)
+            Mask for extracted audio features that should should be ignored when creating the merged
+            multimodal representation (i.e., due to padding).
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
+        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)`) and 2 additional tensors of shape
+            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
+            this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
+            the complete sequence length.
+"""
+
+
+@add_start_docstrings(
+    """The Granite Speech model, which consists of an audio encoder, projector, and language model.""",
+    GRANITE_SPEECH_START_DOCSTRING,
+)
+class GraniteSpeechForConditionalGeneration(GraniteSpeechPreTrainedModel, GenerationMixin):
+    def __init__(self, config: GraniteSpeechConfig):
+        super().__init__(config)
+        # NOTE: It doesn't matter when we initialize from config, but we should be careful
+        # to make sure this does not pick up the adapter_config if in the future we use
+        # from_pretrained or something similar, since that should be set by the composite
+        # model; don't need to consider it twice
+        self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+        if self.language_model._tied_weights_keys is not None:
+            self._tied_weights_keys = [f"language_model.{k}" for k in self.language_model._tied_weights_keys]
+
+        self.encoder = GraniteSpeechCTCModel(config.encoder_config)
+        self.projector = GraniteSpeechEncoderProjectorQFormer(config.projector_config)
+
+        if config.has_lora_adapter and not is_peft_available():
+            logger.warning(
+                "Config indicates that a lora adapter should be present, but "
+                "peft is not installed; this will cause the model to perform "
+                "incorrectly when audio inputs are provided. Please install "
+                "peft and reload the model!"
+            )
+
+        self.post_init()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def get_audio_features(self, input_features):
+        encoder_embeds = self.encoder(input_features)
+        projected_embeds = self.projector(encoder_embeds, None)
+        return projected_embeds
+
+    @add_start_docstrings_to_model_forward(GRANITE_SPEECH_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=GraniteSpeechCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        input_features: torch.FloatTensor = None,
+        input_features_mask: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **lm_kwargs,
+    ) -> Union[Tuple[torch.Tensor], GraniteSpeechCausalLMOutputWithPast]:
+        r"""
+            labels (`torch.LongTensor` of shape `(batch_size, 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]`.
+
+            logits_to_keep (`int` or `torch.Tensor`, *optional*):
+                If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+                `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+                token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+                If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+                This is useful when using packed tensor format (single dimension for batch and sequence length).
+
+        Returns:
+
+        Example:
+
+        TODO - add example for usage.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if input_features is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_features and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if inputs_embeds is None:
+            # Get the base embeddings; set all audio tokens to 0 index
+            # to avoid out of vocabulary issues with the LLM embedding.
+            # Audio features will be masked into is_audio_idx indices later.
+            is_audio_idx = input_ids == self.config.audio_token_index
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[is_audio_idx] = 0
+            inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+        if input_features is not None:
+            if input_features.dtype != self.dtype:
+                logger.warning(f"input features are casted to {self.dtype}")
+                input_features = input_features.to(self.dtype)
+            # Get the audio features from the encoder / projector
+            audio_features = self.get_audio_features(input_features)
+
+            # Merge the audio features into the LLM embeddings
+            inputs_embeds = self.get_merged_audio_embeddings(
+                input_ids=input_ids, audio_features=audio_features, input_features_mask=input_features_mask
+            )
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            **lm_kwargs,
+        )
+        logits = outputs[0]
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+                shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return GraniteSpeechCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        input_features=None,
+        attention_mask=None,
+        cache_position=None,
+        logits_to_keep=None,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward audio inputs to the model
+
+        model_inputs = self.language_model.prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+
+        # If we're in cached decoding stage, input_features should be None because
+        # input ids do not contain special audio token anymore Otherwise we need
+        # input feature values to be passed to the model
+        if cache_position[0] == 0:
+            model_inputs["input_features"] = input_features
+        return model_inputs
+
+    def get_merged_audio_embeddings(self, input_ids, audio_features, input_features_mask):
+        """
+        Adds the audio token to the model's LLM vocabulary so that we can pass it
+        through the tokenizer; it's assumed that the embeddings corresponding to the
+        <|audio|> token will be clobbered with speech features.
+
+        TODO - This needs to be adapted to handle batches of variable length sequences
+        and potentially labels.
+        """
+        is_audio_index = input_ids == self.config.audio_token_index
+        llm_input_ids = torch.where(is_audio_index, 0, input_ids)
+        inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids)  # [bsz, # features, hidden size]
+
+        # Mask the audio features into the text embeddings
+        special_audio_mask = is_audio_index.unsqueeze(-1)
+        audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)[input_features_mask]
+        inputs_embeds = inputs_embeds.masked_scatter(
+            special_audio_mask,
+            audio_features,
+        )
+        return inputs_embeds
+
+    def generate(self, *args, **kwargs):
+        """This model is expected to have a lora adapater, which is only
+        enabled when considering audio inputs. As such, we override generate
+        to conditionally enable / disable the lora adapter based on whether
+        or not any input features were provided.
+        """
+        input_features = kwargs.pop("input_features", None)
+        if is_peft_available and self._hf_peft_config_loaded:
+            if input_features is not None:
+                self.enable_adapters()
+            else:
+                self.disable_adapters()
+        return super().generate(*args, input_features=input_features, **kwargs)
+
+    def save_pretrained(self, *args, **kwargs):
+        # overwrite save_pretrained to first save the adapter if we have one
+        # NOTE - this will use the base model path we are exporting in the lora
+        # adapter, which may not necessarily be the best behavior, but for now
+        # we keep this for portability, since using the local dir causes problems
+        # if the model is loaded from outside of the current working dir.
+        if is_peft_available and self._hf_peft_config_loaded:
+            super().save_pretrained(*args, **kwargs)
+        # Then save the base model afterwards
+        self._hf_peft_config_loaded = False
+        super().save_pretrained(*args, **kwargs)
+
+
+__all__ = [
+    "GraniteSpeechForConditionalGeneration",
+    "GraniteSpeechPreTrainedModel",
+    "GraniteSpeechEncoderProjectorPreTrainedModel",
+    "GraniteSpeechQFormerModel",
+]

preprocessor_config.json

@@ -0,0 +1 @@
+{}

processing_granite_speech.py

@@ -0,0 +1,158 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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 Speech Granite.
+"""
+
+from collections.abc import Sequence
+from typing import List, Union
+
+import numpy as np
+import torch
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PreTokenizedInput, TextInput
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+# 🚨🚨🚨 HACK 🚨🚨🚨
+# This is needed to avoid custom registration issues for now,
+# since we have a custom subclass for the feature extractor as well.
+import transformers
+from .feature_extraction_granite_speech import GraniteSpeechFeatureExtractor
+transformers.GraniteSpeechFeatureExtractor = GraniteSpeechFeatureExtractor
+# The above code is the only change in the modeling code from the following
+# commit on Alex's fork: 397e03a4d76c5f3d8a651e47ade9f27c635e1617
+
+class GraniteSpeechProcessor(ProcessorMixin):
+    attributes = ["feature_extractor", "tokenizer"]
+    valid_kwargs = ["audio_token"]
+
+    feature_extractor_class = "GraniteSpeechFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        feature_extractor,
+        tokenizer,
+        audio_token="<|audio|>",
+    ):
+        self.audio_token = tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
+        super().__init__(feature_extractor, tokenizer)
+
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
+        audios: Union[torch.Tensor, List[torch.Tensor]] = None,
+        device: str = "cpu",
+        **kwargs,
+    ) -> BatchFeature:
+        speech_inputs = {}
+        text_inputs = {}
+
+        text = self._get_validated_text(text)
+        expected_num_audios = sum(t.count(self.audio_token) for t in text)
+
+        if audios is not None:
+            audios, audio_lengths = self._get_validated_audios(audios)
+            if any(text.count(self.audio_token) != 1 for text in text):
+                raise ValueError("Only one audio sample is currently supported per input")
+            if len(audio_lengths) != expected_num_audios:
+                raise ValueError("Text/Audio mismatch. The number of audios and audio tokens do not match")
+
+            # Calculate Mel features & the number of placeholders we will need
+            speech_inputs["input_features"] = self.feature_extractor(
+                audios,
+                device=device,
+            )
+            num_audio_features = self.feature_extractor._get_num_audio_features(audio_lengths)
+            speech_inputs["input_features_mask"] = torch.arange(max(num_audio_features)).view(1, -1) <= torch.tensor(
+                num_audio_features
+            ).view(-1, 1)
+
+            # duplicate the audio placeholders to match the feature dims
+            text = self._expand_audio_placeholders(text, num_audio_features)
+        else:
+            assert expected_num_audios == 0, "No audio is provided, expecting no audio tokens"
+
+        text_inputs = self.tokenizer(text, padding=True, **kwargs)
+        return BatchFeature(data={**text_inputs, **speech_inputs})
+
+    def _expand_audio_placeholders(self, text: list[str], num_audio_features: List[int]):
+        """
+        Expands audio placeholders in the formatted text to match the number of
+        features of the corresponding embeddings; we can use the resulting text
+        to conveniently mask the audio features into the text embeddings.
+        """
+        prompt_strings = []
+        num_replaced = 0
+        for sample in text:
+            while self.audio_token in sample:
+                sample = sample.replace(
+                    self.audio_token,
+                    "<placeholder>" * num_audio_features[num_replaced],
+                    1,
+                )
+                num_replaced += 1
+            prompt_strings.append(sample)
+
+        prompt_strings = [sample.replace("<placeholder>", self.audio_token) for sample in prompt_strings]
+        return prompt_strings
+
+    ##### Validation
+    def _get_validated_text(self, text: Union[str, list]) -> List[str]:
+        if isinstance(text, str):
+            return [text]
+        elif isinstance(text, list) and isinstance(text[0], str):
+            return text
+        raise TypeError("Invalid text provided! Text should be a string or list of strings.")
+
+    def _get_validated_audios(self, audios):
+        # Coerce to PyTorch tensors if we have numpy arrays, since
+        # currently we have a dependency on torch/torchaudio anyway
+        if isinstance(audios, np.ndarray):
+            audios = torch.from_numpy(audios)
+        elif isinstance(audios, Sequence) and isinstance(audios[0], np.ndarray):
+            audios = [torch.from_numpy(arr) for arr in audios]
+
+        if isinstance(audios, torch.Tensor):
+            if audios.ndim == 1:
+                audios = audios.unsqueeze(0)
+            if not torch.is_floating_point(audios):
+                raise ValueError("Invalid audio provided. Audio should be a floating point between 0 and 1")
+
+            if audios.shape[0] > 1:
+                logger.warning("Audio samples are already collated; assuming they all have the same length")
+            lengths = [audios.shape[-1]] * audios.shape[0]
+            return audios, lengths
+
+        elif isinstance(audios, Sequence) and isinstance(audios[0], torch.Tensor):
+            if not torch.is_floating_point(audios[0]):
+                raise ValueError("Invalid audio provided. Audio should be a floating point between 0 and 1")
+            lengths = [audio.shape[-1] for audio in audios]
+            padding = [max(lengths) - length for length in lengths]
+            # ensure all audios have a batch dimension:
+            audios = [audio.view(1, -1) for audio in audios]
+            padded = [torch.nn.functional.pad(audio, (0, pad)) for audio, pad in zip(audios, padding)]
+            audios = torch.cat(padded, dim=0)
+            return audios, lengths
+
+        raise TypeError("Invalid audio provided. Audio should be a one or more torch tensors or numpy arrays")
+
+
+__all__ = ["GraniteSpeechProcessor"]

special_tokens_map.json

@@ -0,0 +1,35 @@
+{
+  "additional_special_tokens": [
+    "<|start_of_role|>",
+    "<|end_of_role|>",
+    "<|tool_call|>"
+  ],
+  "bos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<fim_prefix>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<fim_middle>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<fim_suffix>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<fim_pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<filename>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<gh_stars>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<issue_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<issue_comment>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<issue_closed>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<jupyter_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<jupyter_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<jupyter_code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<jupyter_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<empty_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<commit_before>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<commit_msg>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "17": {
+      "content": "<commit_after>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "18": {
+      "content": "<reponame>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49152": {
+      "content": "<|start_of_role|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49153": {
+      "content": "<|end_of_role|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49154": {
+      "content": "<|tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49155": {
+      "content": "<|audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|start_of_role|>",
+    "<|end_of_role|>",
+    "<|tool_call|>"
+  ],
+  "bos_token": "<|end_of_text|>",
+  "chat_template": "{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set system_message = \"Knowledge Cutoff Date: April 2024.\nToday's Date: \" + strftime_now('%B %d, %Y') + \".\nYou are Granite, developed by IBM.\" %}\n    {%- if tools and documents %}\n        {%- set system_message = system_message + \" You are a helpful AI assistant with access to the following tools. When a tool is required to answer the user's query, respond with <|tool_call|> followed by a JSON list of tools used. If a tool does not exist in the provided list of tools, notify the user that you do not have the ability to fulfill the request.\n\nWrite the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data.\" %}\n    {%- elif tools %}\n        {%- set system_message = system_message + \" You are a helpful AI assistant with access to the following tools. When a tool is required to answer the user's query, respond with <|tool_call|> followed by a JSON list of tools used. If a tool does not exist in the provided list of tools, notify the user that you do not have the ability to fulfill the request.\" %}\n    {%- elif documents %}\n        {%- set system_message = system_message + \" Write the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data.\" %}\n    {%- elif thinking %}\n    {%- set system_message = system_message + \" You are a helpful AI assistant.\nRespond to every user query in a comprehensive and detailed way. You can write down your thoughts and reasoning process before responding. In the thought process, engage in a comprehensive cycle of analysis, summarization, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. In the response section, based on various attempts, explorations, and reflections from the thoughts section, systematically present the final solution that you deem correct. The response should summarize the thought process. Write your thoughts after 'Here is my thought process:' and write your response after 'Here is my response:' for each user query.\" %}\n    {%- else %}\n        {%- set system_message = system_message + \" You are a helpful AI assistant.\" %}    \n    {%- endif %}\n    {%- if 'citations' in controls and documents %}\n        {%- set system_message = system_message + '\n\nIn your response, use the symbols <co> and </co> to indicate when a fact comes from a document in the search result, e.g <co>0</co> for a fact from document 0. Afterwards, list all the citations with their corresponding documents in an ordered list.' %}\n    {%- endif %}\n    {%- if 'hallucinations' in controls and documents %}\n        {%- set system_message = system_message + '\n\nFinally, after the response is written, include a numbered list of sentences from the response that are potentially hallucinated and not based in the documents.' %}\n    {%- endif %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n{{- '<|start_of_role|>system<|end_of_role|>' + system_message + '<|end_of_text|>\n' }}\n{%- if tools %}\n    {{- '<|start_of_role|>tools<|end_of_role|>' }}\n    {{- tools | tojson(indent=4) }}\n    {{- '<|end_of_text|>\n' }}\n{%- endif %}\n{%- if documents %}\n    {{- '<|start_of_role|>documents<|end_of_role|>' }}\n    {%- for document in documents %}\n        {{- 'Document ' + loop.index0 | string + '\n' }}\n        {{- document['text'] }}\n        {%- if not loop.last %}\n            {{- '\n\n'}}\n        {%- endif%}\n    {%- endfor %}\n    {{- '<|end_of_text|>\n' }}\n{%- endif %}\n{%- for message in loop_messages %}\n    {{- '<|start_of_role|>' + message['role'] + '<|end_of_role|>' + message['content'] + '<|end_of_text|>\n' }}\n    {%- if loop.last and add_generation_prompt %}\n        {{- '<|start_of_role|>assistant' }}\n            {%- if controls %}\n                {{- ' ' + controls | tojson()}}\n            {%- endif %}\n        {{- '<|end_of_role|>' }}\n    {%- endif %}\n{%- endfor %}",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 9223372036854775807,
+  "pad_token": "<|end_of_text|>",
+  "padding_side": "left",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|end_of_text|>",
+  "vocab_size": 49152
+}