update to transformers v5: modular conformer, new class names, processor support

README.md

@@ -64,12 +64,12 @@ Measured RTFx of ~1820 on a single H100 GPU (batched inference, batch size 128).
 
 ### Installation
 We require `flash_attention_2` for inference, since this backend supports sequence packing and respects the `is_causal=False` flag. 
-Tested with `transformers==4.57.6` and `transformers==5.5.3`, using `torch==2.9.1`.
+Requires `transformers>=5.5.3` and `torch>=2.9.1`.
 
 ```shell
 # Fresh install (CUDA 12.8, Python 3.10+)
 pip install torch==2.9.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
-pip install transformers==4.57.6 accelerate==1.13.0 safetensors==0.7.0 huggingface-hub==0.36.2 tokenizers==0.22.2
+pip install transformers>=5.5.3 accelerate safetensors huggingface-hub tokenizers
 pip install soundfile
 pip install flash-attn==2.8.3 --no-build-isolation
 ```

__init__.py

@@ -1,3 +1,42 @@
-from .configuration_nle import NLEEncoderConfig, NLEProjectorConfig, NLEConfig
-from .modeling_nle import NLENARDecoder
-from .modeling_ctc import NLECTCEncoder
+# Copyright 2026 IBM and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .configuration_granite_speech_nar import GraniteSpeechNarConfig, GraniteSpeechNarEncoderConfig, GraniteSpeechNarProjectorConfig
+from .feature_extraction_granite_speech_nar import GraniteSpeechNarFeatureExtractor
+from .modeling_granite_speech_nar import (
+    GraniteSpeechNarCTCEncoder,
+    GraniteSpeechNarEncoderOutput,
+    GraniteSpeechNarProjector,
+    GraniteSpeechNarForASR,
+    GraniteSpeechNarLM,
+    GraniteSpeechNarModel,
+    GraniteSpeechNarOutput,
+    GraniteSpeechNarPreTrainedModel,
+)
+from .processing_granite_speech_nar import GraniteSpeechNarProcessor
+
+__all__ = [
+    "GraniteSpeechNarConfig",
+    "GraniteSpeechNarEncoderConfig",
+    "GraniteSpeechNarProjectorConfig",
+    "GraniteSpeechNarFeatureExtractor",
+    "GraniteSpeechNarCTCEncoder",
+    "GraniteSpeechNarEncoderOutput",
+    "GraniteSpeechNarProjector",
+    "GraniteSpeechNarForASR",
+    "GraniteSpeechNarLM",
+    "GraniteSpeechNarModel",
+    "GraniteSpeechNarOutput",
+    "GraniteSpeechNarPreTrainedModel",
+    "GraniteSpeechNarProcessor",
+]

config.json

@@ -1,482 +1,84 @@
 {
-  "architectures": [
-    "NLENARDecoder"
-  ],
-  "attn_implementation": "flash_attention_2",
   "auto_map": {
-    "AutoConfig": "configuration_nle.NLEConfig",
-    "AutoFeatureExtractor": "feature_extraction_nle.NLEFeatureExtractor",
-    "AutoModel": "modeling_nle.NLENARDecoder"
-  },
-  "ctc_tokenizer_config": {
-    "char2idx": {
-      " ": 32,
-      "!": 33,
-      "\"": 34,
-      "#": 35,
-      "$": 36,
-      "%": 37,
-      "&": 38,
-      "'": 39,
-      "(": 40,
-      ")": 41,
-      "*": 42,
-      "+": 43,
-      ",": 44,
-      "-": 45,
-      ".": 46,
-      "/": 47,
-      "0": 48,
-      "1": 49,
-      "2": 50,
-      "3": 51,
-      "4": 52,
-      "5": 53,
-      "6": 54,
-      "7": 55,
-      "8": 56,
-      "9": 57,
-      ":": 58,
-      ";": 59,
-      "<": 60,
-      "=": 61,
-      ">": 62,
-      "?": 63,
-      "@": 64,
-      "A": 65,
-      "B": 66,
-      "C": 67,
-      "D": 68,
-      "E": 69,
-      "F": 70,
-      "G": 71,
-      "H": 72,
-      "I": 73,
-      "J": 74,
-      "K": 75,
-      "L": 76,
-      "M": 77,
-      "N": 78,
-      "O": 79,
-      "P": 80,
-      "Q": 81,
-      "R": 82,
-      "S": 83,
-      "T": 84,
-      "U": 85,
-      "V": 86,
-      "W": 87,
-      "X": 88,
-      "Y": 89,
-      "Z": 90,
-      "[": 91,
-      "\\": 92,
-      "]": 93,
-      "^": 94,
-      "_": 95,
-      "`": 96,
-      "a": 97,
-      "b": 98,
-      "c": 99,
-      "d": 100,
-      "e": 101,
-      "f": 102,
-      "g": 103,
-      "h": 104,
-      "i": 105,
-      "j": 106,
-      "k": 107,
-      "l": 108,
-      "m": 109,
-      "n": 110,
-      "o": 111,
-      "p": 112,
-      "q": 113,
-      "r": 114,
-      "s": 115,
-      "t": 116,
-      "u": 117,
-      "v": 118,
-      "w": 119,
-      "x": 120,
-      "y": 121,
-      "z": 122,
-      "{": 123,
-      "|": 124,
-      "}": 125,
-      "~": 126,
-      "\u007f": 127,
-      "\u0080": 128,
-      "\u0081": 129,
-      "\u0082": 130,
-      "\u0083": 131,
-      "\u0084": 132,
-      "\u0085": 133,
-      "\u0086": 134,
-      "\u0087": 135,
-      "\u0088": 136,
-      "\u0089": 137,
-      "\u008a": 138,
-      "\u008b": 139,
-      "\u008c": 140,
-      "\u008d": 141,
-      "\u008e": 142,
-      "\u008f": 143,
-      "\u0090": 144,
-      "\u0091": 145,
-      "\u0092": 146,
-      "\u0093": 147,
-      "\u0094": 148,
-      "\u0095": 149,
-      "\u0096": 150,
-      "\u0097": 151,
-      "\u0098": 152,
-      "\u0099": 153,
-      "\u009a": 154,
-      "\u009b": 155,
-      "\u009c": 156,
-      "\u009d": 157,
-      "\u009e": 158,
-      "\u009f": 159,
-      "\u00a0": 160,
-      "\u00a1": 161,
-      "\u00a2": 162,
-      "\u00a3": 163,
-      "\u00a4": 164,
-      "\u00a5": 165,
-      "\u00a6": 166,
-      "\u00a7": 167,
-      "\u00a8": 168,
-      "\u00a9": 169,
-      "\u00aa": 170,
-      "\u00ab": 171,
-      "\u00ac": 172,
-      "\u00ad": 173,
-      "\u00ae": 174,
-      "\u00af": 175,
-      "\u00b0": 176,
-      "\u00b1": 177,
-      "\u00b2": 178,
-      "\u00b3": 179,
-      "\u00b4": 180,
-      "\u00b5": 181,
-      "\u00b6": 182,
-      "\u00b7": 183,
-      "\u00b8": 184,
-      "\u00b9": 185,
-      "\u00ba": 186,
-      "\u00bb": 187,
-      "\u00bc": 188,
-      "\u00bd": 189,
-      "\u00be": 190,
-      "\u00bf": 191,
-      "\u00c0": 192,
-      "\u00c1": 193,
-      "\u00c2": 194,
-      "\u00c3": 195,
-      "\u00c4": 196,
-      "\u00c5": 197,
-      "\u00c6": 198,
-      "\u00c7": 199,
-      "\u00c8": 200,
-      "\u00c9": 201,
-      "\u00ca": 202,
-      "\u00cb": 203,
-      "\u00cc": 204,
-      "\u00cd": 205,
-      "\u00ce": 206,
-      "\u00cf": 207,
-      "\u00d0": 208,
-      "\u00d1": 209,
-      "\u00d2": 210,
-      "\u00d3": 211,
-      "\u00d4": 212,
-      "\u00d5": 213,
-      "\u00d6": 214,
-      "\u00d7": 215,
-      "\u00d8": 216,
-      "\u00d9": 217,
-      "\u00da": 218,
-      "\u00db": 219,
-      "\u00dc": 220,
-      "\u00dd": 221,
-      "\u00de": 222,
-      "\u00df": 223,
-      "\u00e0": 224,
-      "\u00e1": 225,
-      "\u00e2": 226,
-      "\u00e3": 227,
-      "\u00e4": 228,
-      "\u00e5": 229,
-      "\u00e6": 230,
-      "\u00e7": 231,
-      "\u00e8": 232,
-      "\u00e9": 233,
-      "\u00ea": 234,
-      "\u00eb": 235,
-      "\u00ec": 236,
-      "\u00ed": 237,
-      "\u00ee": 238,
-      "\u00ef": 239,
-      "\u00f0": 240,
-      "\u00f1": 241,
-      "\u00f2": 242,
-      "\u00f3": 243,
-      "\u00f4": 244,
-      "\u00f5": 245,
-      "\u00f6": 246,
-      "\u00f7": 247,
-      "\u00f8": 248,
-      "\u00f9": 249,
-      "\u00fa": 250,
-      "\u00fb": 251,
-      "\u00fc": 252,
-      "\u00fd": 253,
-      "\u00fe": 254,
-      "\u00ff": 255,
-      "\u30a1": 256,
-      "\u30a2": 257,
-      "\u30a3": 258,
-      "\u30a4": 259,
-      "\u30a5": 260,
-      "\u30a6": 261,
-      "\u30a7": 262,
-      "\u30a8": 263,
-      "\u30a9": 264,
-      "\u30aa": 265,
-      "\u30ab": 266,
-      "\u30ac": 267,
-      "\u30ad": 268,
-      "\u30ae": 269,
-      "\u30af": 270,
-      "\u30b0": 271,
-      "\u30b1": 272,
-      "\u30b2": 273,
-      "\u30b3": 274,
-      "\u30b4": 275,
-      "\u30b5": 276,
-      "\u30b6": 277,
-      "\u30b7": 278,
-      "\u30b8": 279,
-      "\u30b9": 280,
-      "\u30ba": 281,
-      "\u30bb": 282,
-      "\u30bc": 283,
-      "\u30bd": 284,
-      "\u30be": 285,
-      "\u30bf": 286,
-      "\u30c0": 287,
-      "\u30c1": 288,
-      "\u30c2": 289,
-      "\u30c3": 290,
-      "\u30c4": 291,
-      "\u30c5": 292,
-      "\u30c6": 293,
-      "\u30c7": 294,
-      "\u30c8": 295,
-      "\u30c9": 296,
-      "\u30ca": 297,
-      "\u30cb": 298,
-      "\u30cc": 299,
-      "\u30cd": 300,
-      "\u30ce": 301,
-      "\u30cf": 302,
-      "\u30d0": 303,
-      "\u30d1": 304,
-      "\u30d2": 305,
-      "\u30d3": 306,
-      "\u30d4": 307,
-      "\u30d5": 308,
-      "\u30d6": 309,
-      "\u30d7": 310,
-      "\u30d8": 311,
-      "\u30d9": 312,
-      "\u30da": 313,
-      "\u30db": 314,
-      "\u30dc": 315,
-      "\u30dd": 316,
-      "\u30de": 317,
-      "\u30df": 318,
-      "\u30e0": 319,
-      "\u30e1": 320,
-      "\u30e2": 321,
-      "\u30e3": 322,
-      "\u30e4": 323,
-      "\u30e5": 324,
-      "\u30e6": 325,
-      "\u30e7": 326,
-      "\u30e8": 327,
-      "\u30e9": 328,
-      "\u30ea": 329,
-      "\u30eb": 330,
-      "\u30ec": 331,
-      "\u30ed": 332,
-      "\u30ee": 333,
-      "\u30ef": 334,
-      "\u30f0": 335,
-      "\u30f1": 336,
-      "\u30f2": 337,
-      "\u30f3": 338,
-      "\u30f4": 339,
-      "\u30f5": 340,
-      "\u30f6": 341,
-      "\u30f7": 342,
-      "\u30f8": 343,
-      "\u30f9": 344,
-      "\u30fa": 345,
-      "\u30fb": 346,
-      "\u30fc": 347
-    }
+    "AutoConfig": "configuration_granite_speech_nar.GraniteSpeechNarConfig",
+    "AutoFeatureExtractor": "feature_extraction_granite_speech_nar.GraniteSpeechNarFeatureExtractor",
+    "AutoProcessor": "processing_granite_speech_nar.GraniteSpeechNarProcessor",
+    "AutoModel": "modeling_granite_speech_nar.GraniteSpeechNarForASR"
   },
-  "dtype": "bfloat16",
+  "architectures": [
+    "GraniteSpeechNarForASR"
+  ],
+  "model_type": "granite_speech_nar",
+  "torch_dtype": "bfloat16",
   "encoder_config": {
-    "attn_type": "block",
-    "backbone": null,
-    "backbone_requires_grad": false,
-    "bpe_output_dim": 100353,
-    "bpe_pooling_window": 4,
-    "context_size": 200,
-    "conv_expansion_factor": 2,
-    "conv_kernel_size": 15,
-    "dim_head": 128,
-    "dropout": 0.1,
-    "feedforward_mult": 4,
-    "fmask_F": 6,
-    "fmask_m": 2,
-    "fmask_prob": 0.9,
-    "hidden_dim": 1024,
-    "initializer_range": 0.02,
+    "model_type": "granite_speech_nar_encoder",
     "input_dim": 160,
-    "loss_lambda": 0.2,
-    "max_pos_emb": 512,
-    "model.bpe_pooling_window": 4,
-    "model_type": "nle_encoder",
-    "num_heads": 8,
     "num_layers": 16,
-    "old_encoder_mask": true,
+    "hidden_dim": 1024,
+    "feedforward_mult": 4,
+    "num_heads": 8,
+    "dim_head": 128,
     "output_dim": 348,
+    "context_size": 200,
+    "max_pos_emb": 512,
+    "dropout": 0.1,
     "pred_dropout": 0.25,
+    "conv_kernel_size": 15,
+    "conv_expansion_factor": 2,
     "self_conditioning_layer": 8,
-    "tmask_T": 50,
-    "tmask_m": 2,
-    "tmask_m_relative_max": 0.02,
-    "tmask_prob": 0.9
+    "bpe_output_dim": 100353,
+    "bpe_pooling_window": 4,
+    "initializer_range": 0.02
   },
-  "encoder_layer_indices": [
-    4,
-    8,
-    12,
-    -1
-  ],
-  "initializer_range": 0.02,
-  "llm_config": {
-    "_name_or_path": "ibm-granite/granite-4.0-1b-base",
-    "add_cross_attention": false,
-    "architectures": [
-      "GraniteForCausalLM"
-    ],
-    "attention_bias": false,
-    "attention_dropout": 0.0,
-    "attention_multiplier": 0.0078125,
-    "bad_words_ids": null,
-    "begin_suppress_tokens": null,
-    "bos_token_id": 100257,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "dtype": "bfloat16",
-    "early_stopping": false,
-    "embedding_multiplier": 12,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 100257,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "hidden_act": "silu",
+  "projector_config": {
+    "model_type": "granite_speech_nar_projector",
+    "encoder_dim": 1024,
+    "llm_dim": 2048,
+    "downsample_rate": 5,
+    "num_encoder_layers": 4,
     "hidden_size": 2048,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "initializer_range": 0.1,
-    "intermediate_size": 4096,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "length_penalty": 1.0,
-    "logits_scaling": 8,
-    "max_length": 20,
-    "max_position_embeddings": 4096,
-    "min_length": 0,
-    "mlp_bias": false,
+    "num_heads": 32,
+    "num_layers": 2,
+    "dropout_prob": 0.1,
+    "block_size": 15,
+    "mlp_ratio": 2,
+    "layernorm_eps": 1e-06,
+    "attn_bias": true,
+    "mlp_bias": true
+  },
+  "text_config": {
     "model_type": "granite",
-    "no_repeat_ngram_size": 0,
-    "num_attention_heads": 16,
-    "num_beam_groups": 1,
-    "num_beams": 1,
+    "vocab_size": 100352,
+    "hidden_size": 2048,
+    "intermediate_size": 4096,
     "num_hidden_layers": 40,
+    "num_attention_heads": 16,
     "num_key_value_heads": 4,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": 100256,
-    "prefix": null,
-    "problem_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "residual_multiplier": 0.22,
-    "return_dict": true,
-    "return_dict_in_generate": false,
+    "hidden_act": "silu",
+    "max_position_embeddings": 4096,
+    "initializer_range": 0.1,
     "rms_norm_eps": 1e-05,
+    "use_cache": true,
+    "pad_token_id": 100256,
+    "bos_token_id": 100257,
+    "eos_token_id": 100257,
+    "tie_word_embeddings": true,
     "rope_parameters": {
       "rope_theta": 10000,
       "rope_type": "default"
     },
-    "rope_scaling": null,
-    "rope_theta": 10000.0,
-    "sep_token_id": null,
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torchscript": false,
-    "transformers_version": "4.57.6",
-    "typical_p": 1.0,
-    "use_bfloat16": false,
-    "use_cache": true,
-    "vocab_size": 100352
-  },
-  "llm_name": "ibm-granite/granite-4.0-1b-base",
-  "model_type": "nle",
-  "projector_config": {
-    "attn_bias": true,
-    "block_size": 15,
-    "downsample_rate": 5,
-    "dropout_prob": 0.1,
-    "encoder_dim": 1024,
-    "hidden_size": 2048,
-    "layernorm_eps": 1e-06,
-    "llm_dim": 2048,
-    "mlp_bias": true,
-    "mlp_ratio": 2,
-    "model_type": "nle_projector",
-    "num_encoder_layers": 4,
-    "num_heads": 32,
-    "num_layers": 2
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "embedding_multiplier": 12,
+    "logits_scaling": 8,
+    "residual_multiplier": 0.22,
+    "attention_multiplier": 0.0078125
   },
+  "encoder_layer_indices": [4, 8, 12, -1],
   "scale_projected_embeddings": true,
-  "transformers_version": "4.57.6"
+  "blank_token_id": 100257,
+  "min_edit_sequence_length": 8,
+  "ce_loss_lambda": 0.0,
+  "encoder_ctc_loss_lambda": 0.0
 }

configuration_granite_speech_nar.py

@@ -0,0 +1,214 @@
+# Copyright 2026 IBM and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Config classes for Granite Speech NAR (Non-Autoregressive ASR)."""
+
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+
+
+@strict
+class GraniteSpeechNarEncoderConfig(PreTrainedConfig):
+    r"""
+    Configuration for the conformer encoder component of GraniteSpeechNar.
+
+    feedforward_mult (`int`, *optional*, defaults to 4):
+        Multiplier for the feedforward layers; intermediate dim = `hidden_dim * feedforward_mult`.
+    output_dim (`int`, *optional*, defaults to 348):
+        Output dimension of the mid-layer CTC prediction head.
+    context_size (`int`, *optional*, defaults to 200):
+        Context size for block-wise conformer attention.
+    max_pos_emb (`int`, *optional*, defaults to 512):
+        Maximum relative positional embedding index (Shaw's relative positional encoding).
+    pred_dropout (`float`, *optional*, defaults to 0.25):
+        Dropout applied to encoder hidden states before prediction heads.
+    conv_expansion_factor (`int`, *optional*, defaults to 2):
+        Expansion factor for conformer convolution module.
+    self_conditioning_layer (`int`, *optional*):
+        Layer index at which self-conditioning (mid-layer CTC feedback) is applied.
+        Defaults to `num_layers // 2`.
+    bpe_output_dim (`int`, *optional*):
+        Vocabulary size for the BPE CTC head (shifted by +1 for blank). If None, BPE head is disabled.
+    bpe_pooling_window (`int`, *optional*, defaults to 4):
+        Window size for posterior-weighted pooling before the BPE CTC head.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechNarEncoderConfig
+
+    >>> configuration = GraniteSpeechNarEncoderConfig()
+    >>> print(configuration.hidden_dim)
+    1024
+    ```"""
+
+    model_type = "granite_speech_nar_encoder"
+    attribute_map = {
+        "hidden_size": "hidden_dim",
+        "num_hidden_layers": "num_layers",
+        "num_attention_heads": "num_heads",
+        "num_mel_bins": "input_dim",
+    }
+
+    input_dim: int = 160
+    num_layers: int = 16
+    hidden_dim: int = 1024
+    feedforward_mult: int = 4
+    num_heads: int = 8
+    dim_head: int | None = None
+    output_dim: int = 348
+    context_size: int = 200
+    max_pos_emb: int = 512
+    dropout: float = 0.1
+    pred_dropout: float = 0.25
+    conv_kernel_size: int = 15
+    conv_expansion_factor: int = 2
+    self_conditioning_layer: int | None = None
+    bpe_output_dim: int | None = None
+    bpe_pooling_window: int = 4
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        if self.dim_head is None:
+            self.dim_head = self.hidden_dim // self.num_heads
+        if self.self_conditioning_layer is None:
+            self.self_conditioning_layer = self.num_layers // 2
+
+
+@strict
+class GraniteSpeechNarProjectorConfig(PreTrainedConfig):
+    r"""
+    Configuration for the QFormer-based audio projector in GraniteSpeechNar.
+
+    encoder_dim (`int`, *optional*, defaults to 1024):
+        Hidden dimension of the encoder (per layer).
+    llm_dim (`int`, *optional*, defaults to 2048):
+        Hidden dimension of the language model.
+    downsample_rate (`int`, *optional*, defaults to 5):
+        Temporal downsampling rate within each window block.
+    num_encoder_layers (`int`, *optional*, defaults to 4):
+        Number of encoder layers concatenated as projector input.
+    block_size (`int`, *optional*, defaults to 15):
+        Window size for blocked cross-attention in the projector.
+    layernorm_eps (`float`, *optional*, defaults to 1e-6):
+        Epsilon for layer normalization.
+    attn_bias (`bool`, *optional*, defaults to `True`):
+        Whether to use bias in attention projections.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechNarProjectorConfig
+
+    >>> configuration = GraniteSpeechNarProjectorConfig()
+    >>> print(configuration.hidden_size)
+    2048
+    ```"""
+
+    model_type = "granite_speech_nar_projector"
+
+    encoder_dim: int = 1024
+    llm_dim: int = 2048
+    downsample_rate: int = 5
+    num_encoder_layers: int = 4
+    hidden_size: int = 2048
+    num_heads: int = 32
+    num_layers: int = 2
+    dropout_prob: float = 0.1
+    block_size: int = 15
+    mlp_ratio: int = 2
+    layernorm_eps: float = 1e-6
+    attn_bias: bool = True
+    mlp_bias: bool = True
+
+
+@strict
+class GraniteSpeechNarConfig(PreTrainedConfig):
+    r"""
+    Configuration for the GraniteSpeechNar non-autoregressive ASR model.
+
+    This model uses a conformer encoder with BPE CTC head, a QFormer-based projector,
+    and a bidirectional Granite LLM backbone for single-pass speech recognition.
+
+    projector_config (`GraniteSpeechNarProjectorConfig` or `dict`, *optional*):
+        Configuration for the QFormer-based audio projector.
+    tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+        Whether the LLM's input and output word embeddings should be tied.
+    encoder_layer_indices (`list[int]`, *optional*, defaults to `[4, 8, 12, -1]`):
+        Indices of encoder layers whose hidden states are concatenated as projector input.
+    scale_projected_embeddings (`bool`, *optional*, defaults to `True`):
+        Whether to divide projected audio embeddings by the LLM's embedding multiplier.
+    blank_token_id (`int`, *optional*):
+        Token ID used as the CTC blank symbol. Defaults to `text_config.eos_token_id`.
+    min_edit_sequence_length (`int`, *optional*, defaults to 8):
+        Minimum length of the edit sequence (CTC tokens + insertion slots) fed to the LLM.
+    ce_loss_lambda (`float`, *optional*, defaults to 0.0):
+        Weight for auxiliary cross-entropy loss on the LLM output.
+    encoder_ctc_loss_lambda (`float`, *optional*, defaults to 0.0):
+        Weight for auxiliary encoder BPE CTC loss.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechNarConfig, GraniteSpeechNarForASR
+
+    >>> configuration = GraniteSpeechNarConfig()
+    >>> model = GraniteSpeechNarForASR(configuration)
+    >>> print(configuration.model_type)
+    'granite_speech_nar'
+    ```"""
+
+    model_type = "granite_speech_nar"
+    sub_configs = {
+        "encoder_config": GraniteSpeechNarEncoderConfig,
+        "projector_config": GraniteSpeechNarProjectorConfig,
+        "text_config": "AutoConfig",
+    }
+
+    encoder_config: dict | PreTrainedConfig | None = None
+    projector_config: dict | PreTrainedConfig | None = None
+    text_config: dict | PreTrainedConfig | None = None
+    tie_word_embeddings: bool = True
+    encoder_layer_indices: list[int] | None = None
+    scale_projected_embeddings: bool = True
+    blank_token_id: int | None = None
+    min_edit_sequence_length: int = 8
+    ce_loss_lambda: float = 0.0
+    encoder_ctc_loss_lambda: float = 0.0
+
+    def __post_init__(self, **kwargs):
+        if isinstance(self.text_config, dict):
+            self.text_config["model_type"] = self.text_config.get("model_type", "granite")
+            self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+        elif self.text_config is None:
+            self.text_config = CONFIG_MAPPING["granite"]()
+
+        if not isinstance(self.encoder_config, GraniteSpeechNarEncoderConfig):
+            self.encoder_config = GraniteSpeechNarEncoderConfig(**(self.encoder_config or {}))
+
+        if not isinstance(self.projector_config, GraniteSpeechNarProjectorConfig):
+            self.projector_config = GraniteSpeechNarProjectorConfig(**(self.projector_config or {}))
+
+        if self.encoder_layer_indices is None:
+            self.encoder_layer_indices = [4, 8, 12, -1]
+
+        if self.blank_token_id is None:
+            self.blank_token_id = self.text_config.eos_token_id
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["GraniteSpeechNarEncoderConfig", "GraniteSpeechNarProjectorConfig", "GraniteSpeechNarConfig"]

configuration_nle.py

@@ -1,151 +0,0 @@
-from typing import List, Optional, Union, Any
-
-from transformers.configuration_utils import PretrainedConfig
-
-
-class NLEEncoderConfig(PretrainedConfig):
-    model_type = "nle_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,
-        max_pos_emb=512,
-        dropout=0.1,
-        pred_dropout=0.25,
-        conv_kernel_size=15,
-        conv_expansion_factor=2,
-        loss_lambda=0.2,
-        initializer_range=0.02,
-        self_conditioning_layer=None,
-        old_encoder_mask=True,
-        bpe_output_dim=None,
-        bpe_pooling_window=4,
-        **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.pred_dropout = pred_dropout
-        self.conv_kernel_size = conv_kernel_size
-        self.conv_expansion_factor = conv_expansion_factor
-        self.max_pos_emb = max_pos_emb
-        self.loss_lambda = loss_lambda
-        self.initializer_range = initializer_range
-        if self_conditioning_layer is None:
-            self_conditioning_layer = num_layers // 2
-        self.self_conditioning_layer = self_conditioning_layer
-        self.old_encoder_mask = old_encoder_mask
-        self.bpe_output_dim = bpe_output_dim
-        self.bpe_pooling_window = bpe_pooling_window
-
-
-class NLEProjectorConfig(PretrainedConfig):
-    """Config for the QFormer-based encoder-to-LLM projector."""
-    model_type = "nle_projector"
-
-    def __init__(
-        self,
-        encoder_dim: int = 1024,
-        llm_dim: int = 2048,
-        downsample_rate: int = 5,
-        num_encoder_layers: int = 1,
-        hidden_size: Optional[int] = None,
-        num_heads: Optional[int] = None,
-        num_layers: int = 1,
-        dropout_prob: float = 0.0,
-        block_size: int = 15,
-        mlp_ratio: int = 2,
-        layernorm_eps: float = 1e-6,
-        attn_bias: bool = True,
-        mlp_bias: bool = True,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.encoder_dim = encoder_dim
-        self.llm_dim = llm_dim
-        self.downsample_rate = downsample_rate
-        self.num_encoder_layers = num_encoder_layers
-        self.hidden_size = hidden_size if hidden_size is not None else encoder_dim
-        self.num_heads = num_heads if num_heads is not None else self.hidden_size // 64
-        self.num_layers = num_layers
-        self.dropout_prob = dropout_prob
-        self.block_size = block_size
-        self.mlp_ratio = mlp_ratio
-        self.layernorm_eps = layernorm_eps
-        self.attn_bias = attn_bias
-        self.mlp_bias = mlp_bias
-
-
-class NLEConfig(PretrainedConfig):
-    model_type = "nle"
-
-    def __init__(
-        self,
-        encoder_config: Union[NLEEncoderConfig, dict, None] = None,
-        projector_config: Union[NLEProjectorConfig, dict, None] = None,
-        llm_name: str = "ibm-granite/granite-3.3-2b-base",
-        llm_config: Optional[dict] = None,
-        attn_implementation: str = "flash_attention_2",
-        initializer_range: float = 0.02,
-        encoder_layer_indices: Optional[List[int]] = None,
-        scale_projected_embeddings: bool = False,
-        ctc_tokenizer_config: Optional[dict] = None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        if isinstance(encoder_config, dict):
-            self.encoder_config = NLEEncoderConfig(**encoder_config)
-        elif isinstance(encoder_config, NLEEncoderConfig):
-            self.encoder_config = encoder_config
-        elif encoder_config is None:
-            self.encoder_config = NLEEncoderConfig()
-        else:
-            raise TypeError("encoder_config must be NLEEncoderConfig or dict")
-
-        if isinstance(projector_config, dict):
-            self.projector_config = NLEProjectorConfig(**projector_config)
-        elif isinstance(projector_config, NLEProjectorConfig):
-            self.projector_config = projector_config
-        elif projector_config is None:
-            self.projector_config = NLEProjectorConfig()
-        else:
-            raise TypeError("projector_config must be NLEProjectorConfig or dict")
-
-        self.llm_name = llm_name
-        self.llm_config = llm_config
-        self.attn_implementation = attn_implementation
-        self.initializer_range = initializer_range
-        self.encoder_layer_indices = list(encoder_layer_indices) if encoder_layer_indices is not None else [-1]
-        self.scale_projected_embeddings = scale_projected_embeddings
-        self.ctc_tokenizer_config = ctc_tokenizer_config
-        self.auto_map = {
-            "AutoConfig": "configuration_nle.NLEConfig",
-            "AutoModel": "modeling_nle.NLENARDecoder",
-            "AutoFeatureExtractor": "feature_extraction_nle.NLEFeatureExtractor",
-        }
-
-    def to_dict(self):
-        d = super().to_dict()
-        d["encoder_config"] = self.encoder_config.to_dict()
-        d["projector_config"] = self.projector_config.to_dict()
-        if self.llm_config is not None:
-            d["llm_config"] = self.llm_config
-        return d
-
-
-__all__ = ["NLEEncoderConfig", "NLEProjectorConfig", "NLEConfig"]

feature_extraction_nle.py → feature_extraction_granite_speech_nar.py

@@ -1,32 +1,36 @@
-"""Feature extractor for NLE models.
+# Copyright 2026 IBM and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Feature extraction for Granite Speech NAR."""
 
-Handles mel spectrogram extraction, frame stacking, batching, and computing
-the correct attention_mask / x_sizes at encoder-frame rate.
-
-Usage:
-    from feature_extraction_nle import NLEFeatureExtractor
-
-    feature_extractor = NLEFeatureExtractor()
-    inputs = feature_extractor([waveform1, waveform2])
-    output = model.generate(**inputs)
-"""
+from transformers.feature_extraction_utils import FeatureExtractionMixin
+from transformers.tokenization_utils_base import AudioInput
+from transformers.utils import is_torch_available, is_torchaudio_available
+from transformers.utils.import_utils import requires_backends
 
-from typing import List, Optional, Union
 
-import torch
-import torchaudio
-from transformers.feature_extraction_utils import FeatureExtractionMixin
+if is_torch_available():
+    import torch
 
+if is_torchaudio_available():
+    import torchaudio
 
-class NLEFeatureExtractor(FeatureExtractionMixin):
-    """Prepares raw audio for the NLENARDecoder.
 
-    Pipeline: raw 16kHz audio -> MelSpectrogram -> log-mel normalize -> stack 2 frames.
-    Encoder frame rate = sample_rate / (hop_length * 2) = 50 fps, i.e. T_samples // 320.
+class GraniteSpeechNarFeatureExtractor(FeatureExtractionMixin):
+    """Extracts log-mel spectrogram features for GraniteSpeechNar.
 
-    Returns:
-      - input_features:  [B, T_enc, 160] stacked log-mel features
-      - attention_mask:   [B, T_enc] bool mask at encoder-frame rate
+    Produces stacked pairs of 80-band mel frames, yielding 160-dim features
+    at half the original frame rate.
     """
 
     model_input_names = ["input_features", "attention_mask"]
@@ -40,56 +44,50 @@ class NLEFeatureExtractor(FeatureExtractionMixin):
         n_mels: int = 80,
         **kwargs,
     ):
+        requires_backends(self, ["torch", "torchaudio"])
         super().__init__(**kwargs)
-        self.auto_map = {"AutoFeatureExtractor": "feature_extraction_nle.NLEFeatureExtractor"}
         self.sampling_rate = sampling_rate
         self.n_fft = n_fft
         self.win_length = win_length
         self.hop_length = hop_length
         self.n_mels = n_mels
         self.mel_filters = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sampling_rate, n_fft=n_fft, win_length=win_length,
-            hop_length=hop_length, n_mels=n_mels,
+            sample_rate=sampling_rate,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            n_mels=n_mels,
         )
 
-    @torch.no_grad()
-    def _extract_features(self, raw_audio: torch.Tensor) -> torch.Tensor:
-        """Convert raw waveform batch to stacked log-mel features.
-
-        Args:
-            raw_audio: [B, T] raw 16kHz waveform
-
-        Returns:
-            [B, T_enc, n_mels * 2] stacked log-mel features
-        """
-        melspec = self.mel_filters.to(raw_audio.device)
-        B, T = raw_audio.shape
-        # Ensure even number of mel frames for stacking
-        l = 2 * (T // (2 * self.hop_length))
-        mel = melspec(raw_audio.float())[..., :l]
-        logmel = mel.transpose(-1, -2).clamp_min_(1e-10).log10_()
-        mx = logmel.amax(dim=(-2, -1), keepdim=True)
-        logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1)
-        # Stack 2 consecutive frames
-        return logmel.reshape(B, -1, 2 * self.n_mels)
+    def to_dict(self):
+        output = super().to_dict()
+        output.pop("mel_filters", None)
+        return output
+
+    def to_json_string(self):
+        import json
+        dictionary = self.to_dict()
+        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
+
+    def __repr__(self):
+        return f"{self.__class__.__name__} {self.to_json_string()}"
+
+    def _extract_features(self, raw_audio: "torch.Tensor") -> "torch.Tensor":
+        with torch.no_grad():
+            mel_filters = self.mel_filters.to(raw_audio.device)
+            B, T = raw_audio.shape
+            l = 2 * (T // (2 * self.hop_length))
+            mel = mel_filters(raw_audio.float())[..., :l]
+            logmel = mel.transpose(-1, -2).clamp_min_(1e-10).log10_()
+            mx = logmel.amax(dim=(-2, -1), keepdim=True)
+            logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1)
+            return logmel.reshape(B, -1, 2 * self.n_mels)
 
     def __call__(
         self,
-        audios: Union[torch.Tensor, List[torch.Tensor]],
-        device: Optional[Union[str, torch.device]] = None,
+        audios: AudioInput,
+        device: str | None = None,
     ) -> dict:
-        """Prepare a batch of raw audio waveforms for the model.
-
-        Args:
-            audios: A single tensor [T] or [B, T], or a list of 1-D tensors
-                    (variable length). Expected 16 kHz float waveforms.
-            device: Target device for the output tensors.
-
-        Returns:
-            Dict with keys: input_features, attention_mask — ready to
-            unpack into model.generate(**inputs).
-        """
-        # Normalise to list of 1-D tensors
         if isinstance(audios, torch.Tensor):
             if audios.ndim == 1:
                 audios = [audios]
@@ -101,7 +99,6 @@ class NLEFeatureExtractor(FeatureExtractionMixin):
         raw_lengths = [a.shape[-1] for a in audios]
         encoder_frame_counts = [l // (2 * self.hop_length) for l in raw_lengths]
 
-        # Pad waveforms to same length
         raw_audio = torch.nn.utils.rnn.pad_sequence(
             [a.squeeze(0) if a.ndim > 1 else a for a in audios],
             batch_first=True,
@@ -110,10 +107,8 @@ class NLEFeatureExtractor(FeatureExtractionMixin):
         if device is not None:
             raw_audio = raw_audio.to(device)
 
-        # Extract mel features on the padded batch
         input_features = self._extract_features(raw_audio)
 
-        # Build attention_mask at encoder-frame rate
         max_enc_frames = input_features.shape[1]
         x_sizes = torch.tensor(encoder_frame_counts, dtype=torch.long)
         attention_mask = torch.arange(max_enc_frames).unsqueeze(0) < x_sizes.unsqueeze(1)
@@ -128,4 +123,4 @@ class NLEFeatureExtractor(FeatureExtractionMixin):
         }
 
 
-__all__ = ["NLEFeatureExtractor"]
+__all__ = ["GraniteSpeechNarFeatureExtractor"]

generation_config.json

@@ -0,0 +1,3 @@
+{
+  "transformers_version": "5.8.1"
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ae13622e43596054266cfd773dd2f411d608da268e036f9fd0c640adf9b66e1c
-size 4509380570
+oid sha256:668b9fb26c417a4a8b78e3b99cb2fa3c52335f2cdcaafe99a9ad76f19df04fbb
+size 4509384522

modeling_conformer.py

@@ -1,156 +0,0 @@
-
-import math, torch
-import torch.nn.functional as F
-from torch import nn
-from .configuration_nle import NLEEncoderConfig
-
-
-class NLEConformerFeedForward(nn.Module):
-    """Feedforward module for conformer encoder blocks."""
-
-    def __init__(self, config: NLEEncoderConfig):
-        super().__init__()
-        self.pre_norm = nn.LayerNorm(config.hidden_dim)
-        self.up_proj = nn.Linear(config.hidden_dim, config.hidden_dim * config.feedforward_mult)
-        self.silu = nn.SiLU()
-        self.dropout = nn.Dropout(config.dropout)
-        self.down_proj = nn.Linear(config.hidden_dim * config.feedforward_mult, config.hidden_dim)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.pre_norm(hidden_states)
-        hidden_states = self.up_proj(hidden_states)
-        hidden_states = self.dropout(self.silu(hidden_states))
-        hidden_states = self.down_proj(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        return hidden_states
-
-
-class NLEConformerAttention(nn.Module):
-    """Attention for conformer blocks using Shaw's relative positional embeddings.
-    See the following [paper](https://arxiv.org/pdf/1803.02155) for more details.
-    """
-
-    def __init__(self, config: NLEEncoderConfig):
-        super().__init__()
-        self.config = config
-        inner_dim = config.dim_head * config.num_heads
-        self.max_pos_emb = config.max_pos_emb
-        self.context_size = config.context_size
-        self.num_heads = config.num_heads
-        self.dim_head = config.dim_head
-        self.scale = self.dim_head**-0.5
-        self.pre_norm = nn.LayerNorm(config.hidden_dim)
-        self.to_q = nn.Linear(config.hidden_dim, inner_dim, bias=False)
-        self.to_kv = nn.Linear(config.hidden_dim, inner_dim * 2, bias=False)
-        self.to_out = nn.Linear(inner_dim, config.hidden_dim)
-        self.rel_pos_emb = nn.Embedding(2 * self.max_pos_emb + 1, self.dim_head)
-        self.dropout = nn.Dropout(config.dropout)
-
-        if self.context_size <= 0 or self.context_size > self.max_pos_emb:
-            raise ValueError("Context size is either less than 0 or exceeds the max_pos_emb")
-
-    def forward(self, hidden_states: torch.Tensor, 
-                attention_mask: torch.Tensor) -> torch.Tensor:
-
-        hidden_states = self.pre_norm(hidden_states)
-        bsz, num_features, _ = hidden_states.shape
-        num_blocks = math.ceil(num_features / self.context_size)
-        remainder = num_features % self.context_size
-        if self.config.old_encoder_mask:
-            attention_mask = torch.ones_like(attention_mask)
-        if remainder > 0:
-            # right padding to reach block size
-            hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, self.context_size - remainder))
-            attention_mask = torch.nn.functional.pad(attention_mask, (0, self.context_size - remainder))
-
-        query_states = self.to_q(hidden_states)
-        key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
-
-        query_states = query_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
-        key_states = key_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
-        value_states = value_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
-        seq = torch.arange(self.config.context_size, device=hidden_states.device)
-        dist = seq.view(-1, 1) - seq.view(1, -1) + self.config.max_pos_emb
-        rel_pos_emb = self.rel_pos_emb(dist).to(query_states.dtype)
-        pos_attn = torch.einsum('b m h c d, c r d -> b m h c r', query_states, rel_pos_emb) * self.scale
-        mask_value = -torch.finfo(pos_attn.dtype).max
-        expanded_attention_mask = attention_mask.reshape(bsz, num_blocks, 1, 1, -1)
-        pos_attn.masked_fill_(~expanded_attention_mask, mask_value)
-
-        with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
-            out = F.scaled_dot_product_attention(
-                query_states, key_states, value_states, attn_mask=pos_attn, scale=self.scale
-            )
-        out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
-        out = self.to_out(out[:, :num_features, :])
-        return self.dropout(out)
-
-
-class NLEConformerDepthWiseConv1d(nn.Module):
-    """Wrapper for padded 1D pointwise convolution."""
-
-    def __init__(self, chan_in: int, chan_out: int, kernel_size: int):
-        super().__init__()
-        # Padding for the 1D conv is symmetric or close (i.e., offset by one).
-        pad = kernel_size // 2
-        pad_offset = (kernel_size + 1) % 2
-        self.padding = (pad, pad - pad_offset)
-
-        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in, bias=False)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = F.pad(hidden_states, self.padding)
-        return self.conv(hidden_states)
-
-
-class NLEConformerConvModule(nn.Module):
-    """Conformer conv module consisting of several 1D/depthwise 1D convolutional layers."""
-
-    def __init__(self, config: NLEEncoderConfig):
-        super().__init__()
-        inner_dim = config.hidden_dim * config.conv_expansion_factor
-
-        self.norm = nn.LayerNorm(config.hidden_dim)
-        self.up_conv = nn.Conv1d(config.hidden_dim, inner_dim * 2, 1)
-        self.glu = nn.GLU(dim=1)
-        self.depth_conv = NLEConformerDepthWiseConv1d(
-            inner_dim,
-            inner_dim,
-            kernel_size=config.conv_kernel_size,
-        )
-        self.silu = nn.SiLU()
-        self.batch_norm = nn.BatchNorm1d(inner_dim)
-        self.down_conv = nn.Conv1d(inner_dim, config.hidden_dim, 1)
-        self.dropout = nn.Dropout(config.dropout)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.norm(hidden_states)
-        hidden_states = self.up_conv(hidden_states.permute(0, 2, 1))
-        hidden_states = self.glu(hidden_states)
-        hidden_states = self.depth_conv(hidden_states)
-        hidden_states = self.silu(self.batch_norm(hidden_states))
-        hidden_states = self.down_conv(hidden_states).permute(0, 2, 1)
-        hidden_states = self.dropout(hidden_states)
-        return hidden_states
-
-
-class NLEConformerBlock(nn.Module):
-    """Conformer block, consisting largely of linear layers, attention, and convolutional layers."""
-
-    def __init__(self, config: NLEEncoderConfig):
-        super().__init__()
-        self.ff1 = NLEConformerFeedForward(config)
-        self.attn = NLEConformerAttention(config)
-        self.conv = NLEConformerConvModule(config)
-        self.ff2 = NLEConformerFeedForward(config)
-        self.post_norm = nn.LayerNorm(config.hidden_dim)
-
-    def forward(self, hidden_states: torch.Tensor, 
-                attention_mask: torch.Tensor) -> torch.Tensor:
-        hidden_states = 0.5 * self.ff1(hidden_states) + hidden_states
-        hidden_states = self.attn(hidden_states, 
-                                  attention_mask=attention_mask) + hidden_states
-        hidden_states = self.conv(hidden_states) + hidden_states
-        hidden_states = 0.5 * self.ff2(hidden_states) + hidden_states
-        hidden_states = self.post_norm(hidden_states)
-        return hidden_states

modeling_ctc.py

@@ -1,135 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-from transformers.modeling_outputs import ModelOutput
-from transformers.modeling_utils import PreTrainedModel
-from .configuration_nle import NLEEncoderConfig
-from .modeling_conformer import NLEConformerBlock
-
-
-def posterior_weighted_pool(hidden, importance, window_size=4):
-    """
-    Downsample hidden states using importance-weighted pooling.
-
-    Args:
-        hidden: [B, T, D] hidden states
-        importance: [B, T] importance weights (e.g., 1 - blank_prob)
-        window_size: number of frames to pool together
-
-    Returns:
-        pooled: [B, T//window_size, D] downsampled hidden states
-    """
-    B, T, D = hidden.shape
-    pad_len = (window_size - T % window_size) % window_size
-    if pad_len > 0:
-        hidden = F.pad(hidden, (0, 0, 0, pad_len))
-        importance = F.pad(importance, (0, pad_len))
-
-    num_windows = hidden.shape[1] // window_size
-    hidden = hidden.view(B, num_windows, window_size, D)
-    importance = importance.view(B, num_windows, window_size)
-
-    weights = importance / (importance.sum(dim=-1, keepdim=True) + 1e-8)
-    pooled = (hidden * weights.unsqueeze(-1)).sum(dim=2)
-    return pooled
-
-
-@dataclass
-class NLEEncoderOutput(ModelOutput):
-    logits: Optional[torch.FloatTensor] = None
-    logits_bpe: Optional[torch.FloatTensor] = None
-    last_hidden_state: Optional[torch.FloatTensor] = None
-    all_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-class NLECTCEncoder(PreTrainedModel):
-    config_class = NLEEncoderConfig
-
-    def __init__(self, config: NLEEncoderConfig):
-        super().__init__(config)
-        self.config = config
-        self.input_linear = nn.Linear(config.input_dim, config.hidden_dim, bias=True)
-        self.layers = nn.ModuleList([NLEConformerBlock(config) for _ 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.out_bpe = None
-        if config.bpe_output_dim is not None:
-            self.out_bpe = nn.Linear(config.hidden_dim, config.bpe_output_dim, bias=True)
-
-        self.dropout = nn.Dropout(config.pred_dropout)
-
-        self.post_init()
-
-    def _init_weights(self, module: nn.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)
-        elif isinstance(module, (nn.LayerNorm, nn.BatchNorm1d)):
-            module.weight.data.fill_(1.0)
-            module.bias.data.zero_()
-
-    def forward(
-        self,
-        input_features: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,   # [B, T_enc] bool after stacking
-        output_hidden_states: Optional[bool] = None,
-    ) -> NLEEncoderOutput:
-
-        inputs_embeds = input_features
-        if attention_mask is None:
-            mask_shape = inputs_embeds.shape[:-1]
-            attention_mask = torch.ones(mask_shape, dtype=bool, device=inputs_embeds.device)
-
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-
-        hidden_states = self.input_linear(inputs_embeds.to(self.dtype))
-        all_hidden_states = (hidden_states,) if output_hidden_states else None
-        blank_prob_mid = None
-
-        for idx, layer in enumerate(self.layers, start=1):
-            hidden_states = layer(hidden_states, attention_mask=attention_mask)
-
-            if idx == self.config.self_conditioning_layer:
-                logits_mid_plain = self.out(self.dropout(hidden_states))
-                probs_mid = torch.softmax(logits_mid_plain.float(), dim=-1)
-                blank_prob_mid = probs_mid[:, :, 0]
-                hidden_states = hidden_states + self.out_mid(probs_mid.to(hidden_states.dtype))
-
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-        hidden_states = self.dropout(hidden_states)
-        logits = self.out(hidden_states)
-
-        logits_bpe = None
-        if self.out_bpe is not None and blank_prob_mid is not None:
-            importance = 1.0 - blank_prob_mid
-            pooled = posterior_weighted_pool(
-                hidden_states.float(), importance, window_size=self.config.bpe_pooling_window
-            ).to(hidden_states.dtype)
-            # Downsample attention_mask to pooled resolution
-            pool_window = self.config.bpe_pooling_window
-            T = attention_mask.shape[1]
-            pad_len = (pool_window - T % pool_window) % pool_window
-            bpe_mask = F.pad(attention_mask, (0, pad_len), value=False)[:, ::pool_window]
-            # Only compute BPE projection on valid frames (100K vocab is expensive)
-            # Returns flat [N_valid, bpe_output_dim] — caller must unpack using bpe_mask
-            valid_pooled = pooled[bpe_mask]  # [N_valid, D]
-            logits_bpe = self.out_bpe(valid_pooled)  # [N_valid, bpe_output_dim]
-
-        return NLEEncoderOutput(
-            logits=logits,
-            logits_bpe=logits_bpe,
-            last_hidden_state=hidden_states,
-            all_hidden_states=all_hidden_states
-        )
-

modeling_granite_speech_nar.py

@@ -0,0 +1,1149 @@
+# Copyright 2026 IBM and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from transformers.masking_utils import create_bidirectional_mask, find_packed_sequence_indices, packed_sequence_mask_function
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import ModelOutput, TransformersKwargs
+from transformers.utils.generic import maybe_autocast
+from .configuration_granite_speech_nar import (
+    GraniteSpeechNarConfig,
+    GraniteSpeechNarEncoderConfig,
+    GraniteSpeechNarProjectorConfig,
+)
+
+
+@dataclass
+class GraniteSpeechNarEncoderOutput(ModelOutput):
+    """Output of the GraniteSpeechNar encoder."""
+
+    loss: torch.Tensor | None = None
+    logits: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    all_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+class GraniteSpeechNarOutput(ModelOutput):
+    """Output of the GraniteSpeechNarForASR model.
+
+    Attributes:
+        loss: Combined CTC + auxiliary losses (only when labels provided).
+        preds: List of predicted token ID tensors per sample (after CTC collapse, inference only).
+        logits: List of per-sample logit tensors from the LLM head.
+        encoder_logits: Flat BPE CTC logits from the encoder.
+        encoder_preds: List of CTC-collapsed encoder predictions per sample.
+    """
+
+    loss: torch.Tensor | None = None
+    preds: list[torch.Tensor] | None = None
+    logits: list[torch.Tensor] | None = None
+    encoder_logits: torch.Tensor | None = None
+    encoder_preds: list[torch.Tensor] | None = None
+
+
+### Encoder - conformer is adapted from: https://github.com/lucidrains/conformer.git
+class GraniteSpeechNarConformerFeedForward(nn.Module):
+    """Feedforward module for conformer encoder blocks."""
+
+    def __init__(self, config: GraniteSpeechNarEncoderConfig):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(config.hidden_dim)
+        self.up_proj = nn.Linear(config.hidden_dim, config.hidden_dim * config.feedforward_mult)
+        self.silu = nn.SiLU()
+        self.dropout = nn.Dropout(config.dropout)
+        self.down_proj = nn.Linear(config.hidden_dim * config.feedforward_mult, config.hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states)
+        hidden_states = self.up_proj(hidden_states)
+        hidden_states = self.dropout(self.silu(hidden_states))
+        hidden_states = self.down_proj(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechNarConformerAttention(nn.Module):
+    """Attention for conformer blocks using Shaw's relative positional embeddings.
+    See the following [paper](https://huggingface.co/papers/1803.02155) for more details.
+    """
+
+    def __init__(self, config: GraniteSpeechNarEncoderConfig):
+        super().__init__()
+
+        inner_dim = config.dim_head * config.num_heads
+        self.max_pos_emb = config.max_pos_emb
+        self.context_size = config.context_size
+        self.num_heads = config.num_heads
+        self.dim_head = config.dim_head
+        self.scale = self.dim_head**-0.5
+        self.pre_norm = nn.LayerNorm(config.hidden_dim)
+        self.to_q = nn.Linear(config.hidden_dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(config.hidden_dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, config.hidden_dim)
+        self.rel_pos_emb = nn.Embedding(2 * self.max_pos_emb + 1, self.dim_head)
+        self.dropout = nn.Dropout(config.dropout)
+
+        if self.context_size <= 0 or self.context_size > self.max_pos_emb:
+            raise ValueError("Context size is either less than 0 or exceeds the max_pos_emb")
+
+    def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states)
+        bsz, num_features, _ = hidden_states.shape
+
+        num_blocks = math.ceil(num_features / self.context_size)
+        remainder = num_features % self.context_size
+        if remainder > 0:
+            # right padding to reach block size
+            hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, self.context_size - remainder))
+
+        query_states = self.to_q(hidden_states)
+        key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
+
+        query_states = query_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+        key_states = key_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+        value_states = value_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+
+        # shaw's relative positional embedding
+        rel_pos_emb = self.rel_pos_emb(attention_dists)
+        # alternative computation of `pos_attn` - for readability
+        # rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] + list(rel_pos_emb.shape))
+        # pos_attn = torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded, dim=-1) * self.scale
+        # einsum implementation of pos_attn - gives x30 speedup over the alternative
+        # TODO (@avihu111) find a fast alternative to einsum
+        pos_attn = torch.einsum("b m h c d, c r d -> b m h c r", query_states, rel_pos_emb) * self.scale
+
+        if remainder > 0:
+            # masked attention in the extended block
+            mask = torch.ones(self.context_size, self.context_size, dtype=bool, device=hidden_states.device)
+            mask[:remainder, :remainder] = 0
+            mask_value = -torch.finfo(pos_attn.dtype).max
+            pos_attn[:, -1, :].masked_fill_(mask, mask_value)
+
+        with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+            out = F.scaled_dot_product_attention(
+                query_states, key_states, value_states, attn_mask=pos_attn, scale=self.scale
+            )
+        out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
+        out = self.to_out(out[:, :num_features, :])
+        return self.dropout(out)
+
+
+class GraniteSpeechNarConformerDepthWiseConv1d(nn.Module):
+    """Wrapper for padded 1D pointwise convolution."""
+
+    def __init__(self, chan_in: int, chan_out: int, kernel_size: int):
+        super().__init__()
+        # Padding for the 1D conv is symmetric or close (i.e., offset by one).
+        pad = kernel_size // 2
+        pad_offset = (kernel_size + 1) % 2
+        self.padding = (pad, pad - pad_offset)
+
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.pad(hidden_states, self.padding)
+        return self.conv(hidden_states)
+
+
+class GraniteSpeechNarConformerConvModule(nn.Module):
+    """Conformer conv module consisting of several 1D/depthwise 1D convolutional layers."""
+
+    def __init__(self, config: GraniteSpeechNarEncoderConfig):
+        super().__init__()
+        inner_dim = config.hidden_dim * config.conv_expansion_factor
+
+        self.norm = nn.LayerNorm(config.hidden_dim)
+        self.up_conv = nn.Conv1d(config.hidden_dim, inner_dim * 2, 1)
+        self.glu = nn.GLU(dim=1)
+        self.depth_conv = GraniteSpeechNarConformerDepthWiseConv1d(
+            inner_dim,
+            inner_dim,
+            kernel_size=config.conv_kernel_size,
+        )
+        self.silu = nn.SiLU()
+        self.batch_norm = nn.BatchNorm1d(inner_dim)
+        self.down_conv = nn.Conv1d(inner_dim, config.hidden_dim, 1)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.up_conv(hidden_states.permute(0, 2, 1))
+        hidden_states = self.glu(hidden_states)
+        hidden_states = self.depth_conv(hidden_states)
+        hidden_states = self.silu(self.batch_norm(hidden_states))
+        hidden_states = self.down_conv(hidden_states).permute(0, 2, 1)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechNarConformerBlock(nn.Module):
+    """Conformer block, consisting largely of linear layers, attention, and convolutional layers."""
+
+    def __init__(self, config: GraniteSpeechNarEncoderConfig):
+        super().__init__()
+        self.ff1 = GraniteSpeechNarConformerFeedForward(config)
+        self.attn = GraniteSpeechNarConformerAttention(config)
+        self.conv = GraniteSpeechNarConformerConvModule(config)
+        self.ff2 = GraniteSpeechNarConformerFeedForward(config)
+        self.post_norm = nn.LayerNorm(config.hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+        hidden_states = 0.5 * self.ff1(hidden_states) + hidden_states
+        hidden_states = self.attn(hidden_states, attention_dists=attention_dists) + hidden_states
+        hidden_states = self.conv(hidden_states) + hidden_states
+        hidden_states = 0.5 * self.ff2(hidden_states) + hidden_states
+        hidden_states = self.post_norm(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechNarQFormerCrossAttention(nn.Module):
+    def __init__(self, config: GraniteSpeechNarProjectorConfig):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.head_dim = config.hidden_size // config.num_heads
+        self.hidden_size = config.hidden_size
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
+        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
+        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
+
+    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, query_len, _ = hidden_states.shape
+        encoder_len = encoder_hidden_states.shape[1]
+
+        query_states = (
+            self.q_proj(hidden_states).view(batch_size, query_len, self.num_heads, self.head_dim).transpose(1, 2)
+        )
+        key_states = (
+            self.k_proj(encoder_hidden_states)
+            .view(batch_size, encoder_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        value_states = (
+            self.v_proj(encoder_hidden_states)
+            .view(batch_size, encoder_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+
+        attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, is_causal=False)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, query_len, self.hidden_size)
+        return self.o_proj(attn_output)
+
+
+class GraniteSpeechNarQFormerMLP(nn.Module):
+    def __init__(self, config: GraniteSpeechNarProjectorConfig):
+        super().__init__()
+        mlp_hidden_size = int(config.hidden_size * config.mlp_ratio)
+        self.fc1 = nn.Linear(config.hidden_size, mlp_hidden_size, bias=config.mlp_bias)
+        self.act = nn.SiLU()
+        self.fc2 = nn.Linear(mlp_hidden_size, config.hidden_size, bias=config.mlp_bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.act(self.fc1(hidden_states)))
+
+
+class GraniteSpeechNarQFormerLayer(nn.Module):
+    def __init__(self, config: GraniteSpeechNarProjectorConfig):
+        super().__init__()
+        self.attn_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
+        self.cross_attention = GraniteSpeechNarQFormerCrossAttention(config)
+        self.mlp_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
+        self.mlp = GraniteSpeechNarQFormerMLP(config)
+
+    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.cross_attention(self.attn_norm(hidden_states), encoder_hidden_states)
+        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
+        return hidden_states
+
+
+class GraniteSpeechNarQFormer(nn.Module):
+    def __init__(self, config: GraniteSpeechNarProjectorConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([GraniteSpeechNarQFormerLayer(config) for _ in range(config.num_layers)])
+
+    def forward(self, query_embeds: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = query_embeds
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, encoder_hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechNarProjector(nn.Module):
+    """Windowed QFormer projector that maps multi-layer encoder features to LLM embedding space."""
+
+    def __init__(self, config: GraniteSpeechNarProjectorConfig):
+        super().__init__()
+        self.config = config
+        self.layer_norms = nn.ModuleList(
+            [nn.LayerNorm(config.encoder_dim, eps=config.layernorm_eps) for _ in range(config.num_encoder_layers)]
+        )
+        self.layer_projector = nn.Linear(config.encoder_dim * config.num_encoder_layers, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout_prob)
+        self.projector_act = nn.GELU()
+        self.qformer = GraniteSpeechNarQFormer(config)
+
+        query_length = config.block_size // config.downsample_rate
+        embed_std = config.hidden_size**-0.5
+        self.query = nn.Parameter(torch.randn(1, query_length, config.hidden_size) * embed_std)
+        self.window_positions = nn.Parameter(torch.randn(1, config.block_size, config.hidden_size) * embed_std)
+        self.out_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
+        self.out_linear = nn.Linear(config.hidden_size, config.llm_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, dim = hidden_states.size()
+
+        hidden_states = hidden_states.view(
+            batch_size, seq_len, self.config.num_encoder_layers, self.config.encoder_dim
+        )
+        normalized_layers = []
+        for i, layer_norm in enumerate(self.layer_norms):
+            normalized_layers.append(layer_norm(hidden_states[:, :, i]))
+        hidden_states = torch.cat(normalized_layers, dim=-1)
+
+        hidden_states = self.projector_act(self.layer_projector(hidden_states))
+
+        block_size = self.config.block_size
+        nblocks = seq_len // block_size
+        rest = seq_len % block_size
+        if rest > 0:
+            hidden_states = F.pad(hidden_states, (0, 0, 0, block_size - rest), "constant", 0)
+            nblocks += 1
+
+        hidden_states = hidden_states.view(batch_size * nblocks, block_size, self.config.hidden_size)
+        query_length = self.query.shape[1]
+        mean_pool = hidden_states.view(
+            batch_size * nblocks, query_length, self.config.downsample_rate, self.config.hidden_size
+        ).mean(dim=-2)
+
+        hidden_states = self.qformer(
+            query_embeds=self.dropout(self.query + mean_pool),
+            encoder_hidden_states=self.dropout(hidden_states + self.window_positions),
+        )
+
+        hidden_states = hidden_states.view(batch_size, nblocks * query_length, -1)
+        hidden_states = self.dropout(self.out_norm(hidden_states))
+        return self.out_linear(hidden_states)
+
+
+class GraniteSpeechNarPreTrainedModel(PreTrainedModel):
+    config_class = GraniteSpeechNarConfig
+    base_model_prefix = "encoder"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _no_split_modules = ["GraniteSpeechNarConformerBlock", "GraniteDecoderLayer"]
+    input_modalities = ("audio",)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class GraniteSpeechNarAttention(nn.Module):
+    """GraniteAttention with is_causal=False for bidirectional attention."""
+
+    is_causal = False
+
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = config.attention_multiplier
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GraniteSpeechNarRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        GraniteSpeechNarRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class GraniteSpeechNarMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class GraniteSpeechNarDecoderLayer(GradientCheckpointingLayer):
+    """GraniteDecoderLayer using bidirectional attention."""
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = GraniteSpeechNarAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = GraniteSpeechNarMLP(config)
+        self.input_layernorm = GraniteSpeechNarRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GraniteSpeechNarRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.residual_multiplier = config.residual_multiplier
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            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`).
+            past_key_values (`Cache`, *optional*): cached past key and value projection states
+            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        return hidden_states
+
+
+class GraniteSpeechNarRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: GraniteSpeechNarConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: GraniteSpeechNarConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class GraniteSpeechNarModel(GraniteSpeechNarPreTrainedModel):
+    """GraniteModel with bidirectional (non-causal) attention.
+
+    Uses GraniteSpeechNarDecoderLayer which sets is_causal=False,
+    and replaces create_causal_mask() with create_bidirectional_mask() so all
+    attention backends (SDPA, FA2, eager, flex) get a proper non-causal mask.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [GraniteSpeechNarDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = GraniteSpeechNarRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = GraniteSpeechNarRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.embedding_multiplier = config.embedding_multiplier
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        inputs_embeds = inputs_embeds * self.embedding_multiplier
+
+        if position_ids is None:
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
+
+        packed_seq_mask = find_packed_sequence_indices(position_ids)
+        and_mask_fn = packed_sequence_mask_function(packed_seq_mask) if packed_seq_mask is not None else None
+        bidirectional_mask = create_bidirectional_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            and_mask_function=and_mask_fn,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        # KV cache is not needed in a non-autoregressive model
+        kwargs["use_cache"] = False
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=bidirectional_mask,
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states)
+
+
+def _posterior_weighted_pool(hidden: torch.Tensor, importance: torch.Tensor, window_size: int = 4) -> torch.Tensor:
+    batch_size, seq_len, hidden_dim = hidden.shape
+    pad_len = (window_size - seq_len % window_size) % window_size
+    if pad_len > 0:
+        hidden = F.pad(hidden, (0, 0, 0, pad_len))
+        importance = F.pad(importance, (0, pad_len))
+    num_windows = hidden.shape[1] // window_size
+    hidden = hidden.view(batch_size, num_windows, window_size, hidden_dim)
+    importance = importance.view(batch_size, num_windows, window_size)
+    weights = importance / (importance.sum(dim=-1, keepdim=True) + 1e-8)
+    pooled = (hidden * weights.unsqueeze(-1)).sum(dim=2)
+    return pooled
+
+
+class GraniteSpeechNarCTCEncoder(GraniteSpeechNarPreTrainedModel):
+    """Conformer encoder with BPE CTC head and multi-layer output."""
+
+    config_class = GraniteSpeechNarEncoderConfig
+
+    def __init__(self, config: GraniteSpeechNarEncoderConfig):
+        super().__init__(config)
+        self.input_linear = nn.Linear(config.input_dim, config.hidden_dim, bias=True)
+        self.layers = nn.ModuleList([GraniteSpeechNarConformerBlock(config) for _ 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.out_bpe = None
+        if config.bpe_output_dim is not None:
+            self.out_bpe = nn.Linear(config.hidden_dim, config.bpe_output_dim, bias=True)
+        self.dropout = nn.Dropout(config.pred_dropout)
+        self.post_init()
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        labels: torch.Tensor | None = None,
+        label_lengths: torch.Tensor | None = None,
+        **kwargs,
+    ) -> GraniteSpeechNarEncoderOutput:
+        if attention_mask is None:
+            attention_mask = torch.ones(input_features.shape[:-1], dtype=torch.bool, device=input_features.device)
+
+        hidden_states = self.input_linear(input_features.to(self.dtype))
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        blank_probs = None
+
+        context_size = self.config.context_size
+        seq = torch.arange(context_size, device=hidden_states.device)
+        relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
+        attention_dists = torch.clamp(relpos_dist, -context_size, context_size) + self.config.max_pos_emb
+
+        for layer_idx, layer in enumerate(self.layers, start=1):
+            hidden_states = layer(hidden_states, attention_dists=attention_dists)
+
+            if layer_idx == self.config.self_conditioning_layer:
+                mid_logits = self.out(self.dropout(hidden_states))
+                mid_probs = torch.softmax(mid_logits.float(), dim=-1)
+                blank_probs = mid_probs[:, :, 0]
+                hidden_states = hidden_states + self.out_mid(mid_probs.to(hidden_states.dtype))
+
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+        hidden_states = self.dropout(hidden_states)
+
+        logits = None
+        loss = None
+        if self.out_bpe is not None and blank_probs is not None:
+            pool_window = self.config.bpe_pooling_window
+            importance = 1.0 - blank_probs
+            pooled = _posterior_weighted_pool(hidden_states.float(), importance, window_size=pool_window).to(
+                hidden_states.dtype
+            )
+            encoder_lengths = attention_mask.sum(dim=1)
+            lengths = -(encoder_lengths // -pool_window)
+            lengths_list = lengths.tolist()
+            logits = self.out_bpe(torch.cat([pooled[i, :length] for i, length in enumerate(lengths_list)]))
+
+            if labels is not None:
+                logits_padded = logits.new_zeros(len(lengths_list), max(lengths_list), logits.shape[-1])
+                offset = 0
+                for i, length in enumerate(lengths_list):
+                    logits_padded[i, :length] = logits[offset : offset + length]
+                    offset += length
+
+                log_probs = torch.log_softmax(logits_padded.float(), dim=-1)
+                loss = (
+                    F.ctc_loss(
+                        log_probs.transpose(0, 1),
+                        labels + 1,
+                        lengths,
+                        label_lengths,
+                        blank=0,
+                        reduction="sum",
+                        zero_infinity=True,
+                    )
+                    / lengths.sum()
+                )
+
+        return GraniteSpeechNarEncoderOutput(
+            loss=loss,
+            logits=logits,
+            last_hidden_state=hidden_states,
+            all_hidden_states=all_hidden_states,
+        )
+
+
+class GraniteSpeechNarLM(GraniteSpeechNarPreTrainedModel, GenerationMixin):
+    """GraniteForCausalLM with a bidirectional (non-causal) backbone."""
+
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GraniteSpeechNarModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, GraniteSpeechNarLM
+
+        >>> model = GraniteSpeechNarLM.from_pretrained("meta-granite_speech_nar/GraniteSpeechNar-2-7b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-granite_speech_nar/GraniteSpeechNar-2-7b-hf")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = logits / self.config.logits_scaling  # main diff with Llama
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class GraniteSpeechNarForASR(GraniteSpeechNarPreTrainedModel):
+    def __init__(self, config: GraniteSpeechNarConfig):
+        super().__init__(config)
+
+        self.encoder = GraniteSpeechNarCTCEncoder(config.encoder_config)
+        self.projector = GraniteSpeechNarProjector(config.projector_config)
+
+        text_config = config.text_config
+        if hasattr(config, "_attn_implementation"):
+            text_config._attn_implementation = config._attn_implementation
+        self.language_model = GraniteSpeechNarLM._from_config(text_config)
+
+        self.post_init()
+
+    def _ctc_collapse_decode(
+        self,
+        bpe_logits_flat: torch.Tensor,
+        bpe_lengths: list[int],
+    ) -> list[torch.Tensor]:
+        """GPU CTC greedy decode: argmax -> unique_consecutive -> remove blank -> shift."""
+        preds_flat = bpe_logits_flat.argmax(dim=-1)
+        per_sample = preds_flat.split(bpe_lengths)
+        return [(collapsed := torch.unique_consecutive(seq))[collapsed != 0] - 1 for seq in per_sample]
+
+    def _add_insertion_slots(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """Insert blank tokens between each CTC token as editing slots for the LLM."""
+        blank_id = self.config.blank_token_id
+        n = token_ids.numel()
+        total_len = max(2 * n + 1, self.config.min_edit_sequence_length)
+        idx = torch.arange(n, device=token_ids.device)
+        out_idx = 2 * idx + 1
+        out = torch.full((total_len,), fill_value=blank_id, dtype=token_ids.dtype, device=token_ids.device)
+        out[out_idx] = token_ids
+        return out
+
+    def _build_flat_inputs(
+        self,
+        ctc_token_ids: list[torch.Tensor],
+        audio_embeds: torch.Tensor,
+        audio_lengths: list[int],
+    ) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
+        """Build flat (pad-free) LLM input: [audio_0, text_0, audio_1, text_1, ...]"""
+        embed_tokens = self.language_model.model.embed_tokens
+
+        embeds_list = []
+        position_ids_list = []
+        text_lengths = []
+
+        for i, audio_len in enumerate(audio_lengths):
+            audio_emb = audio_embeds[i, :audio_len]
+            text_ids_with_slots = self._add_insertion_slots(ctc_token_ids[i])
+            text_emb = embed_tokens(text_ids_with_slots)
+            sample_embeds = torch.cat([audio_emb, text_emb], dim=0)
+            embeds_list.append(sample_embeds)
+            position_ids_list.append(torch.arange(sample_embeds.shape[0], device=audio_embeds.device))
+            text_lengths.append(text_ids_with_slots.shape[0])
+
+        flat_embeds = torch.cat(embeds_list, dim=0).unsqueeze(0)
+        flat_position_ids = torch.cat(position_ids_list, dim=0).unsqueeze(0)
+        return flat_embeds, flat_position_ids, text_lengths
+
+    def forward(
+        self,
+        *,
+        input_features: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        label_lengths: torch.Tensor | None = None,
+        output_encoder_logits: bool = False,
+        **kwargs,
+    ) -> GraniteSpeechNarOutput:
+        r"""
+        Args:
+            input_features (`torch.Tensor` of shape `(batch_size, seq_len, input_dim)`):
+                Mel spectrogram features.
+            attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`, *optional*):
+                Encoder attention mask (1 for valid frames, 0 for padding).
+            labels (`torch.Tensor` of shape `(batch_size, max_label_len)`, *optional*):
+                Ground truth LLM token IDs for training.
+            label_lengths (`torch.Tensor` of shape `(batch_size,)`, *optional*):
+                Number of valid tokens per sample in `labels`.
+            output_encoder_logits (`bool`, *optional*, defaults to `False`):
+                Whether to return encoder BPE logits. When False, the large logits
+                tensor is freed early to reduce peak memory.
+
+        Returns:
+            [`GraniteSpeechNarOutput`]
+        """
+        encoder_labels = labels if self.config.encoder_ctc_loss_lambda else None
+        enc_out = self.encoder(
+            input_features=input_features,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            labels=encoder_labels,
+            label_lengths=label_lengths if encoder_labels is not None else None,
+        )
+
+        if attention_mask is None:
+            attention_mask = torch.ones(input_features.shape[:-1], dtype=torch.bool, device=input_features.device)
+
+        encoder_lengths = attention_mask.sum(dim=1)
+
+        pool_window = self.encoder.config.bpe_pooling_window
+        bpe_lengths = (-(encoder_lengths // -pool_window)).tolist()
+        ctc_token_ids = self._ctc_collapse_decode(enc_out.logits, bpe_lengths)
+
+        multilayer_features = torch.cat(
+            [enc_out.all_hidden_states[idx] for idx in self.config.encoder_layer_indices], dim=-1
+        )
+
+        encoder_loss = enc_out.loss
+        encoder_logits = enc_out.logits if output_encoder_logits else None
+        del enc_out
+
+        audio_embeds = self.projector(multilayer_features)
+        del multilayer_features
+        if self.config.scale_projected_embeddings:
+            embedding_multiplier = getattr(self.config.text_config, "embedding_multiplier", 1.0)
+            audio_embeds = audio_embeds / embedding_multiplier
+        audio_embeds = audio_embeds.to(self.language_model.model.embed_tokens.weight.dtype)
+
+        audio_lengths = (encoder_lengths // self.projector.config.downsample_rate).cpu().tolist()
+
+        flat_embeds, flat_position_ids, text_lengths = self._build_flat_inputs(
+            ctc_token_ids, audio_embeds, audio_lengths
+        )
+
+        llm_out = self.language_model(
+            inputs_embeds=flat_embeds,
+            position_ids=flat_position_ids,
+        )
+        all_logits = llm_out.logits.squeeze(0)
+
+        segment_lengths = [l for a, t in zip(audio_lengths, text_lengths) for l in (a, t)]
+        text_logits = torch.cat(list(all_logits.split(segment_lengths)[1::2]))
+        logits_per_sample = list(text_logits.split(text_lengths))
+
+        loss = None
+        if labels is not None:
+            log_probs = torch.log_softmax(text_logits.float(), dim=-1)
+
+            log_probs_padded = log_probs.new_zeros(len(text_lengths), max(text_lengths), log_probs.shape[-1])
+            offset = 0
+            for i, tl in enumerate(text_lengths):
+                log_probs_padded[i, :tl] = log_probs[offset : offset + tl]
+                offset += tl
+
+            input_lengths = torch.tensor(text_lengths, device=text_logits.device)
+
+            loss = (
+                F.ctc_loss(
+                    log_probs_padded.transpose(0, 1),
+                    labels,
+                    input_lengths,
+                    label_lengths,
+                    blank=self.config.blank_token_id,
+                    reduction="sum",
+                    zero_infinity=True,
+                )
+                / input_lengths.sum()
+            )
+
+            if self.config.ce_loss_lambda > 0.0:
+                ce_targets = torch.cat([self._add_insertion_slots(ids) for ids in ctc_token_ids])
+                ce_loss = F.cross_entropy(
+                    text_logits,
+                    ce_targets.long(),
+                    reduction="mean",
+                    ignore_index=-100,
+                )
+                loss = loss + self.config.ce_loss_lambda * ce_loss
+
+            if encoder_loss is not None:
+                loss = loss + self.config.encoder_ctc_loss_lambda * encoder_loss
+
+        return GraniteSpeechNarOutput(
+            loss=loss,
+            logits=logits_per_sample,
+            encoder_logits=encoder_logits,
+            encoder_preds=ctc_token_ids,
+        )
+
+    @torch.inference_mode()
+    def transcribe(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_encoder_logits: bool = False,
+    ) -> GraniteSpeechNarOutput:
+        """Single-pass non-autoregressive inference: forward + CTC collapse on LLM output.
+
+        Returns token ID tensors in `preds`. Use `GraniteSpeechNarProcessor.batch_decode()`
+        to convert to strings.
+        """
+        output = self.forward(
+            input_features=input_features,
+            attention_mask=attention_mask,
+            output_encoder_logits=output_encoder_logits,
+        )
+
+        blank_id = self.config.blank_token_id
+        preds = []
+        for sample_logits in output.logits:
+            pred = torch.unique_consecutive(sample_logits.argmax(-1))
+            pred = pred[pred != blank_id]
+            preds.append(pred)
+
+        return GraniteSpeechNarOutput(
+            preds=preds,
+            logits=output.logits,
+            encoder_logits=output.encoder_logits,
+            encoder_preds=output.encoder_preds,
+        )
+
+
+__all__ = [
+    "GraniteSpeechNarModel",
+    "GraniteSpeechNarCTCEncoder",
+    "GraniteSpeechNarForASR",
+    "GraniteSpeechNarLM",
+    "GraniteSpeechNarPreTrainedModel",
+]

modeling_nle.py

@@ -1,267 +0,0 @@
-from dataclasses import dataclass
-from pathlib import Path
-from typing import List, Optional
-import shutil
-
-import torch
-from transformers import (
-    PreTrainedModel,
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    AutoConfig,
-)
-from transformers.utils import ModelOutput
-
-from .modeling_ctc import NLECTCEncoder
-from .modeling_projector import EncoderProjectorQFormer
-from .configuration_nle import NLEConfig
-from .tokenizer import Tokenizer
-from .modeling_conformer import NLEConformerBlock  # noqa: F401 — needed so HF auto-import finds this module
-
-
-@dataclass
-class NLENARDecoderOutput(ModelOutput):
-    loss: Optional[torch.Tensor] = None
-    text_preds: Optional[List[str]] = None
-    text_ctc_preds: Optional[List[str]] = None
-    editing_logits: Optional[torch.Tensor] = None
-    encoder_logits: Optional[torch.Tensor] = None
-
-
-class NLENARDecoder(PreTrainedModel):
-    config_class = NLEConfig
-
-    def __init__(self, config: NLEConfig):
-        super().__init__(config)
-
-        self.encoder = NLECTCEncoder(config.encoder_config)
-
-        if config.ctc_tokenizer_config is not None:
-            self.ctc_tokenizer = Tokenizer(**config.ctc_tokenizer_config)
-        else:
-            self.ctc_tokenizer = Tokenizer()
-
-        if config.llm_config is not None:
-            llm_cfg = AutoConfig.for_model(**config.llm_config)
-            if config.attn_implementation is not None:
-                llm_cfg._attn_implementation = config.attn_implementation
-            self.llm = AutoModelForCausalLM.from_config(llm_cfg)
-        else:
-            llm_kwargs = {"device_map": "cpu", "torch_dtype": torch.bfloat16}
-            if config.attn_implementation is not None:
-                llm_kwargs["attn_implementation"] = config.attn_implementation
-            self.llm = AutoModelForCausalLM.from_pretrained(config.llm_name, **llm_kwargs)
-
-        for layer in self.llm.model.layers:
-            layer.self_attn.is_causal = False
-
-        self.llm_tokenizer = AutoTokenizer.from_pretrained(config.llm_name)
-        self.projector = EncoderProjectorQFormer(config.projector_config)
-
-        self.post_init()
-
-    def save_pretrained(self, save_directory, **kwargs):
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
-
-        self.config.llm_config = self.llm.config.to_dict()
-        if self.config.ctc_tokenizer_config is None and self.ctc_tokenizer is not None:
-            self.config.ctc_tokenizer_config = {"char2idx": self.ctc_tokenizer.char2idx}
-
-        self.llm_tokenizer.save_pretrained(save_directory)
-        super().save_pretrained(save_directory, **kwargs)
-
-        src_dir = Path(__file__).parent
-        for py_file in src_dir.glob("*.py"):
-            shutil.copy2(py_file, save_directory / py_file.name)
-
-    def add_insertion_slots(self, x: torch.Tensor) -> torch.Tensor:
-        """Inserts pad_id (EOS) tokens between each CTC token."""
-        pad_id = self.llm.config.eos_token_id
-        n = x.numel()
-        total_len = max(2 * n + 1, 8)
-        idx = torch.arange(n, device=x.device)
-        out_idx = 2 * idx + 1
-        out = torch.full((total_len,), fill_value=pad_id, dtype=x.dtype, device=x.device)
-        out[out_idx] = x
-        return out
-
-    def _decode_encoder_greedy(
-        self,
-        encoder_logits: torch.Tensor,
-        attention_mask: torch.Tensor
-    ) -> List[str]:
-        ctc_preds = torch.where(attention_mask, encoder_logits.argmax(dim=-1), 0).cpu().numpy()
-        text_ctc_preds = [self.ctc_tokenizer.decode(pred).strip() for pred in ctc_preds]
-        text_ctc_preds = [x if x != "" else " " for x in text_ctc_preds]
-        return text_ctc_preds
-
-    def _decode_bpe_ctc_greedy(
-        self,
-        bpe_logits_flat: torch.Tensor,
-        bpe_lengths: List[int],
-    ) -> List[str]:
-        """
-        Decode BPE CTC logits to text using greedy decoding.
-
-        BPE label mapping: index 0 = CTC blank, index i = llm_token_id (i-1).
-
-        Args:
-            bpe_logits_flat: [N_valid, V_bpe] raw logits (flat, valid positions only)
-            bpe_lengths: per-sample lengths in pooled resolution
-        """
-        preds_flat = bpe_logits_flat.argmax(dim=-1)  # [N_valid]
-
-        text_ctc_preds = []
-        offset = 0
-        for length in bpe_lengths:
-            seq = preds_flat[offset:offset + length]
-            offset += length
-            collapsed = torch.unique_consecutive(seq)
-            collapsed = collapsed[collapsed != 0] - 1  # remove blank, shift to LLM token IDs
-            if collapsed.numel() > 0:
-                text = self.llm_tokenizer.decode(collapsed.tolist(), skip_special_tokens=True)
-            else:
-                text = " "
-            text = text.strip().lower() or " "
-            text_ctc_preds.append(text)
-        return text_ctc_preds
-
-    def _build_flat_llm_inputs(
-        self,
-        text_ctc_preds: List[str],
-        projected_lengths: List[int],
-        encoder_embs: torch.Tensor,
-    ):
-        """
-        Build a flat (pad-free) LLM input sequence and position IDs.
-
-        Layout per sample: [audio_embeds_i, text_embeds_i]
-        All samples are concatenated into a single flat sequence.
-
-        Returns:
-            flat_embeds: [1, N_total, D] — ready for LLM forward
-            flat_position_ids: [1, N_total] — per-sample cumulative positions
-            text_lengths: List[int] — length of each text segment (for slicing output)
-        """
-        projected = self.projector(encoder_embs)
-        if self.config.scale_projected_embeddings and hasattr(self.llm.config, "embedding_multiplier"):
-            projected = projected / self.llm.config.embedding_multiplier
-        projected = projected.to(self.llm.model.embed_tokens.weight.dtype)
-
-        pred_text_llm_tokens = self.llm_tokenizer(text_ctc_preds)
-        text_ids_with_slots = [
-            self.add_insertion_slots(torch.tensor(x, device=encoder_embs.device))
-            for x in pred_text_llm_tokens.input_ids
-        ]
-
-        embed_tokens = self.llm.model.embed_tokens
-        embeds_list = []
-        position_ids_list = []
-        text_lengths = []
-
-        for i, audio_len in enumerate(projected_lengths):
-            audio_embs = projected[i, :audio_len]  # [audio_len, D]
-            text_embs = embed_tokens(text_ids_with_slots[i])  # [text_len, D]
-            sample_embeds = torch.cat([audio_embs, text_embs], dim=0)
-            embeds_list.append(sample_embeds)
-            position_ids_list.append(torch.arange(sample_embeds.shape[0], device=encoder_embs.device))
-            text_lengths.append(text_ids_with_slots[i].shape[0])
-
-        flat_embeds = torch.cat(embeds_list, dim=0).unsqueeze(0)  # [1, N_total, D]
-        flat_position_ids = torch.cat(position_ids_list, dim=0).unsqueeze(0)  # [1, N_total]
-
-        return flat_embeds, flat_position_ids, text_lengths
-
-    def forward(
-        self,
-        *,
-        input_features: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> NLENARDecoderOutput:
-
-        need_hidden_states = self.config.encoder_layer_indices != [-1]
-        enc_out = self.encoder(
-            input_features=input_features,
-            attention_mask=attention_mask,
-            output_hidden_states=need_hidden_states,
-        )
-
-        encoder_logits = enc_out.logits
-        encoder_logits_bpe = enc_out.logits_bpe
-
-        if enc_out.all_hidden_states is not None and len(self.config.encoder_layer_indices) > 0:
-            selected_list = [enc_out.all_hidden_states[idx] for idx in self.config.encoder_layer_indices]
-            encoder_embs = torch.cat(selected_list, dim=-1)
-        else:
-            encoder_embs = enc_out.last_hidden_state
-        enc_out = None
-
-        if attention_mask is None:
-            attention_mask = torch.ones_like(encoder_logits[..., 0], dtype=torch.bool)
-
-        encoder_lengths = attention_mask.sum(dim=1)
-        projected_lengths = (encoder_lengths // self.config.projector_config.downsample_rate).cpu().tolist()
-
-        if encoder_logits_bpe is not None:
-            pool_window = self.encoder.config.bpe_pooling_window
-            bpe_lengths = (-(encoder_lengths // -pool_window)).tolist()  # ceil division
-            text_ctc_preds = self._decode_bpe_ctc_greedy(encoder_logits_bpe, bpe_lengths)
-            encoder_logits_bpe = None
-        else:
-            text_ctc_preds = self._decode_encoder_greedy(encoder_logits, attention_mask)
-
-        assert self.config.attn_implementation == "flash_attention_2", (
-            "NLENARDecoder requires flash_attention_2. Other backends produce "
-            "causal attention even though is_causal=False"
-        )
-
-        flat_embeds, flat_position_ids, text_lengths = self._build_flat_llm_inputs(
-            text_ctc_preds, projected_lengths, encoder_embs
-        )
-        encoder_embs = None
-
-        # Run LLM backbone (without lm_head) on the full flat sequence
-        llm_hidden = self.llm.model(
-            inputs_embeds=flat_embeds,
-            position_ids=flat_position_ids,
-            use_cache=False,
-        ).last_hidden_state.squeeze(0)  # [N_total, D]
-        flat_embeds = None
-
-        # Slice text hidden states, apply lm_head once on the concatenation
-        # Flat layout: [audio_0, text_0, audio_1, text_1, ...]
-        text_hidden_list = []
-        offset = 0
-        for i in range(len(projected_lengths)):
-            offset += projected_lengths[i]
-            text_hidden_list.append(llm_hidden[offset:offset + text_lengths[i]])
-            offset += text_lengths[i]
-        llm_hidden = None
-
-        text_logits = self.llm.lm_head(torch.cat(text_hidden_list))  # [N_text_total, V]
-        editing_logits_list = text_logits.split(text_lengths)
-
-        return NLENARDecoderOutput(
-            editing_logits=editing_logits_list,
-            encoder_logits=encoder_logits,
-            text_ctc_preds=text_ctc_preds,
-        )
-
-    @torch.inference_mode()
-    def generate(self, input_features, attention_mask):
-        """Single-pass inference: forward + argmax decoding."""
-        output = self.forward(input_features=input_features, attention_mask=attention_mask)
-
-        eos_id = self.llm.config.eos_token_id
-        text_preds = []
-        for logits in output.editing_logits:
-            cur_pred = torch.unique_consecutive(logits.argmax(-1))
-            cur_pred = cur_pred[cur_pred != eos_id]
-            text_preds.append(self.llm_tokenizer.decode(cur_pred, skip_special_tokens=True))
-
-        return NLENARDecoderOutput(
-            text_preds=text_preds,
-            text_ctc_preds=output.text_ctc_preds,
-            encoder_logits=output.encoder_logits,
-        )

modeling_projector.py

@@ -1,147 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .configuration_nle import NLEProjectorConfig
-
-
-class QFormerCrossAttention(nn.Module):
-    def __init__(self, config: NLEProjectorConfig):
-        super().__init__()
-        self.num_heads = config.num_heads
-        self.head_dim = config.hidden_size // config.num_heads
-        self.hidden_size = config.hidden_size
-
-        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
-        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
-        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
-        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attn_bias)
-
-    def forward(self, hidden_states, encoder_hidden_states):
-        batch_size, query_len, _ = hidden_states.shape
-        encoder_len = encoder_hidden_states.shape[1]
-
-        query_states = self.q_proj(hidden_states).view(
-            batch_size, query_len, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        key_states = self.k_proj(encoder_hidden_states).view(
-            batch_size, encoder_len, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        value_states = self.v_proj(encoder_hidden_states).view(
-            batch_size, encoder_len, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        with torch.nn.attention.sdpa_kernel([
-            torch.nn.attention.SDPBackend.FLASH_ATTENTION,
-            torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
-            torch.nn.attention.SDPBackend.MATH,
-        ]):
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                query_states, key_states, value_states, is_causal=False,
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, query_len, self.hidden_size)
-        return self.o_proj(attn_output)
-
-
-class QFormerMLP(nn.Module):
-    def __init__(self, config: NLEProjectorConfig):
-        super().__init__()
-        mlp_hidden_size = int(config.hidden_size * config.mlp_ratio)
-        self.fc1 = nn.Linear(config.hidden_size, mlp_hidden_size, bias=config.mlp_bias)
-        self.act = nn.SiLU()
-        self.fc2 = nn.Linear(mlp_hidden_size, config.hidden_size, bias=config.mlp_bias)
-
-    def forward(self, hidden_states):
-        return self.fc2(self.act(self.fc1(hidden_states)))
-
-
-class QFormerLayer(nn.Module):
-    def __init__(self, config: NLEProjectorConfig):
-        super().__init__()
-        self.attn_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
-        self.cross_attention = QFormerCrossAttention(config)
-        self.mlp_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
-        self.mlp = QFormerMLP(config)
-
-    def forward(self, hidden_states, encoder_hidden_states):
-        hidden_states = hidden_states + self.cross_attention(
-            self.attn_norm(hidden_states), encoder_hidden_states
-        )
-        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
-        return hidden_states
-
-
-class SimplifiedQFormer(nn.Module):
-    def __init__(self, config: NLEProjectorConfig):
-        super().__init__()
-        self.layers = nn.ModuleList([
-            QFormerLayer(config) for _ in range(config.num_layers)
-        ])
-
-    def forward(self, query_embeds, encoder_hidden_states):
-        hidden_states = query_embeds
-        for layer in self.layers:
-            hidden_states = layer(hidden_states, encoder_hidden_states)
-        return hidden_states
-
-
-class EncoderProjectorQFormer(nn.Module):
-    def __init__(self, config: NLEProjectorConfig):
-        super().__init__()
-        self.config = config
-
-        self.layer_norms = nn.ModuleList([
-            nn.LayerNorm(config.encoder_dim, eps=config.layernorm_eps)
-            for _ in range(config.num_encoder_layers)
-        ])
-
-        self.layer_projector = nn.Linear(
-            config.encoder_dim * config.num_encoder_layers, config.hidden_size
-        )
-        self.dropout = nn.Dropout(config.dropout_prob)
-        self.projector_act = nn.GELU()
-
-        self.qformer = SimplifiedQFormer(config)
-
-        query_length = config.block_size // config.downsample_rate
-        embed_std = config.hidden_size ** -0.5
-        self.query = nn.Parameter(
-            torch.randn(1, query_length, config.hidden_size) * embed_std
-        )
-        self.window_positions = nn.Parameter(
-            torch.randn(1, config.block_size, config.hidden_size) * embed_std
-        )
-        self.out_norm = nn.LayerNorm(config.hidden_size, eps=config.layernorm_eps)
-        self.out_linear = nn.Linear(config.hidden_size, config.llm_dim)
-
-    def forward(self, x):
-        batch_size, seq_len, dim = x.size()
-
-        x = x.view(batch_size, seq_len, self.config.num_encoder_layers, self.config.encoder_dim)
-        normalized_layers = []
-        for i, layer_norm in enumerate(self.layer_norms):
-            normalized_layers.append(layer_norm(x[:, :, i]))
-        x = torch.cat(normalized_layers, dim=-1)
-
-        x = self.projector_act(self.layer_projector(x))
-
-        block_size = self.config.block_size
-        nblocks = seq_len // block_size
-        rest = seq_len % block_size
-        if rest > 0:
-            x = nn.functional.pad(x, (0, 0, 0, block_size - rest), 'constant', 0)
-            nblocks += 1
-
-        x = x.view(batch_size * nblocks, block_size, self.config.hidden_size)
-        query_length = self.query.shape[1]
-        mean_pool = x.view(
-            batch_size * nblocks, query_length, self.config.downsample_rate, self.config.hidden_size
-        ).mean(dim=-2)
-
-        query_output = self.qformer(
-            query_embeds=self.dropout(self.query + mean_pool),
-            encoder_hidden_states=self.dropout(x + self.window_positions),
-        )
-
-        query_output = query_output.view(batch_size, nblocks * query_length, -1)
-        query_output = self.dropout(self.out_norm(query_output))
-        return self.out_linear(query_output)

preprocessor_config.json

@@ -1,8 +1,8 @@
 {
   "auto_map": {
-    "AutoFeatureExtractor": "feature_extraction_nle.NLEFeatureExtractor"
+    "AutoFeatureExtractor": "feature_extraction_granite_speech_nar.GraniteSpeechNarFeatureExtractor"
   },
-  "feature_extractor_type": "NLEFeatureExtractor",
+  "feature_extractor_type": "GraniteSpeechNarFeatureExtractor",
   "hop_length": 160,
   "n_fft": 512,
   "n_mels": 80,

processing_granite_speech_nar.py

@@ -0,0 +1,49 @@
+# Copyright 2026 IBM and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Processor for Granite Speech NAR."""
+
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import AudioInput
+from transformers.utils import is_torch_available
+from .feature_extraction_granite_speech_nar import GraniteSpeechNarFeatureExtractor
+
+
+if is_torch_available():
+    import torch
+
+
+class GraniteSpeechNarProcessor(ProcessorMixin):
+    """Processor combining audio feature extraction and tokenizer for GraniteSpeechNar."""
+
+    feature_extractor_class = "AutoFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, feature_extractor: GraniteSpeechNarFeatureExtractor, tokenizer=None, **kwargs):
+        super().__init__(feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs)
+
+    def __call__(
+        self,
+        audios: AudioInput,
+        device: str | None = None,
+        **kwargs,
+    ) -> dict:
+        return self.feature_extractor(audios, device=device)
+
+    def batch_decode(self, token_ids_list: list["torch.Tensor"], **kwargs) -> list[str]:
+        if self.tokenizer is None:
+            raise ValueError("Tokenizer not set. Pass tokenizer to GraniteSpeechNarProcessor.")
+        return [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in token_ids_list]
+
+
+__all__ = ["GraniteSpeechNarProcessor"]

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "processor_class": "GraniteSpeechNarProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_granite_speech_nar.GraniteSpeechNarProcessor"
+  }
+}

tokenizer.py

@@ -1,30 +0,0 @@
-"""Minimal CTC tokenizer for Granite Speech."""
-
-import numpy as np
-
-
-class Tokenizer:
-    """
-    CTC tokenizer with char2idx mapping. Index 0 is always blank.
-    Default vocab: latin256_kana92 (348 tokens).
-    """
-
-    def __init__(self, char2idx=None, **kwargs):
-        if char2idx is None:
-            # Default: latin256_kana92
-            char2idx = {chr(n): n for n in range(32, 256)}
-            char2idx |= {chr(0x30A1 + n): 256 + n for n in range(92)}
-
-        # char2idx values may be strings after JSON roundtrip
-        self.char2idx = {k: int(v) for k, v in char2idx.items()}
-        self.idx2char = {v: k for k, v in self.char2idx.items()}
-        self.vocab_size = len(self.char2idx) + 1
-
-    def encode(self, text: str) -> np.ndarray:
-        return np.array([self.char2idx[c] for c in text if c in self.char2idx], dtype=np.int64)
-
-    def decode(self, tokens: np.ndarray) -> str:
-        """Decode CTC output: unique_consecutive + remove blanks."""
-        pred = tokens[np.insert(tokens[1:] != tokens[:-1], 0, True)]
-        pred = pred[pred != 0]
-        return "".join([self.idx2char[idx] for idx in pred.tolist()])