Promote tiny-audio-granite-qwen-500-3 (Granite NC encoder + fine-tuned Qwen3-0.6B)

Copy of mazesmazes/tiny-audio-granite-qwen-500-3@989df838 (step 16159). config.json: pretrained_model_path -> mazesmazes/tiny-audio, model_dtype -> bfloat16. Previous model tagged granite-qwen-frozen-4.

.gitattributes

@@ -1,3 +1,36 @@
-*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-tokenizer_config.json -filter -diff -merge text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,5 +1,5 @@
 ---
-license: mit
+license: cc-by-nc-sa-4.0
 language:
   - en
 datasets:
@@ -8,12 +8,13 @@ datasets:
   - fixie-ai/common_voice_17_0
   - speechcolab/gigaspeech
   - kensho/spgispeech
-  - facebook/voxpopuli
   - edinburghcstr/ami
   - sanchit-gandhi/tedlium-data
+  - edinburghcstr/edacc
+  - mazesmazes/notsofar-train
 base_model:
-  - ibm-granite/granite-speech-5.0-470m-turboctc
-  - Qwen/Qwen3.5-2B
+  - ibm-granite/granite-speech-5.0-470m-turboctc-nc
+  - Qwen/Qwen3-0.6B
 pipeline_tag: automatic-speech-recognition
 tags:
   - asr
@@ -25,171 +26,48 @@ tags:
   - speaker-diarization
   - qwen
   - granite-speech
-  - lora
 library_name: transformers
 metrics:
   - wer
-model-index:
-  - name: tiny-audio
-    results:
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "LibriSpeech test-clean, subset of up to 1,000 samples"
-          type: openslr/librispeech_asr
-          config: clean
-          split: test
-        metrics:
-          - type: wer
-            value: 1.84
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "LibriSpeech test-other, subset of up to 1,000 samples"
-          type: openslr/librispeech_asr
-          config: other
-          split: test
-        metrics:
-          - type: wer
-            value: 6.38
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "SPGISpeech, subset of up to 1,000 samples"
-          type: kensho/spgispeech
-          config: test
-          split: test
-        metrics:
-          - type: wer
-            value: 2.29
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "TED-LIUM, subset of up to 1,000 samples"
-          type: sanchit-gandhi/tedlium-data
-          config: default
-          split: test
-        metrics:
-          - type: wer
-            value: 3.71
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "LoquaciousSet (held out), subset of up to 1,000 samples"
-          type: speechbrain/LoquaciousSet
-          config: small
-          split: test
-        metrics:
-          - type: wer
-            value: 6.2
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "VoxPopuli, subset of up to 1,000 samples"
-          type: facebook/voxpopuli
-          config: en
-          split: test
-        metrics:
-          - type: wer
-            value: 7.11
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "Common Voice 17, subset of up to 1,000 samples"
-          type: fixie-ai/common_voice_17_0
-          config: en
-          split: test
-        metrics:
-          - type: wer
-            value: 7.18
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "AMI (IHM), subset of up to 1,000 samples"
-          type: edinburghcstr/ami
-          config: ihm
-          split: test
-        metrics:
-          - type: wer
-            value: 8.99
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "GigaSpeech (dev), subset of up to 1,000 samples"
-          type: fixie-ai/gigaspeech
-          config: dev
-          split: dev
-        metrics:
-          - type: wer
-            value: 9.06
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "Earnings22 (held out), subset of up to 1,000 samples"
-          type: sanchit-gandhi/earnings22_robust_split
-          config: default
-          split: test
-        metrics:
-          - type: wer
-            value: 10.58
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "People's Speech, subset of up to 1,000 samples"
-          type: fixie-ai/peoples_speech
-          config: clean
-          split: test
-        metrics:
-          - type: wer
-            value: 17.59
-            name: WER (normalized, subset of up to 1,000 samples)
-      - task:
-          type: automatic-speech-recognition
-          name: Automatic Speech Recognition
-        dataset:
-          name: "AMI (SDM), subset of up to 1,000 samples"
-          type: edinburghcstr/ami
-          config: sdm
-          split: test
-        metrics:
-          - type: wer
-            value: 23.53
-            name: WER (normalized, subset of up to 1,000 samples)
 ---
 
 # Tiny Audio
 
 **A tiny, hackable, open-source speech LLM.**
 
-It transcribes English with punctuation, capitalization, word timestamps, and speaker labels: 1.8%
-WER on LibriSpeech test-clean and 7.4% across 12 benchmarks (11,822 samples pooled). Built and
-trained with [Tiny Audio](https://github.com/alexkroman/tiny-audio).
+It transcribes English with punctuation, capitalization, word timestamps, and speaker labels. Built
+and trained with [Tiny Audio](https://github.com/alexkroman/tiny-audio).
 
 **[Try the live demo](https://huggingface.co/spaces/mazesmazes/tiny-audio)** ·
 **[Train your own](https://github.com/alexkroman/tiny-audio)** ·
 **[Free 3.5-hour course](https://github.com/alexkroman/tiny-audio/blob/main/docs/course/0-course-overview.md)**
 
+## Model Description
+
+Tiny Audio connects a frozen, pretrained speech encoder to a pretrained LLM with a small trainable
+projector:
+
+```text
+Audio (16 kHz) → speech encoder (frozen) → MLP projector (trained) → LLM decoder → Text
+```
+
+The published model pairs the non-causal Granite Speech 5.0 470M encoder (kept frozen) with the
+Qwen3-0.6B decoder. The MLP projector maps encoder frames into the decoder's embedding space; it is
+the only part trained from scratch, and the whole decoder is fine-tuned alongside it: ~604M trained
+parameters (596M decoder, 8.4M projector). The encoder, projector, and decoder are each swappable
+from config in the [training repo](https://github.com/alexkroman/tiny-audio).
+
+## Intended Use
+
+Tiny Audio is for transcribing English speech — meetings, calls, lectures, and other recorded audio
+— when you want the finished transcript itself: punctuation, capitalization, and formatted numbers
+included, plus word-level timestamps and speaker labels when you need them. It is also built to be
+read and retrained: the training code, data mix, and recipe are open, so you can adapt it to your
+own audio.
+
+Out of scope: non-English transcription, and audio far from the training distribution — far-field,
+noisy, or heavily overlapping speech — where word error rate rises (see Benchmarks and Limitations).
+
 ## Quick Start
 
 ```bash
@@ -211,36 +89,13 @@ the model writes punctuation, capitalization, and numbers itself.
 
 ## Benchmarks
 
-Word error rate (%, lower is better) on 11,822 samples (up to 1,000 per dataset), scored with the
+Benchmark results for this release have not been published yet. To score it yourself with the
 [Tiny Audio eval harness](https://github.com/alexkroman/tiny-audio/tree/main/scripts/eval)
-(`ta eval`) after text normalization.
-
-| Dataset                     |      WER |
-| --------------------------- | -------: |
-| LibriSpeech test-clean      |     1.84 |
-| SPGISpeech                  |     2.29 |
-| TED-LIUM                    |     3.71 |
-| LoquaciousSet †             |     6.20 |
-| LibriSpeech test-other      |     6.38 |
-| VoxPopuli                   |     7.11 |
-| Common Voice                |     7.18 |
-| AMI (IHM)                   |     8.99 |
-| GigaSpeech                  |     9.06 |
-| Earnings22 †                |    10.58 |
-| People's Speech             |    17.59 |
-| AMI (SDM)                   |    23.53 |
-| **Mean (12 sets)**          | **8.71** |
-| **Pooled (11,822 samples)** | **7.42** |
-
-† Held out: no data from this source was used in training.
-
-To reproduce a row, or score AssemblyAI, Deepgram, ElevenLabs, or Apple's on-device recognizer on
-the same samples:
+(`ta eval`):
 
 ```bash
 git clone https://github.com/alexkroman/tiny-audio.git && cd tiny-audio && poetry install
 poetry run ta eval -m mazesmazes/tiny-audio -d loquacious -n 100
-ASSEMBLYAI_API_KEY=... poetry run ta eval -m assemblyai -d loquacious -n 100
 ```
 
 ## More Than Plain Text
@@ -265,9 +120,9 @@ for r in pipe(["audio1.wav", "audio2.wav", "audio3.wav"], batch_size=4):
 
 `return_timestamps=True` times every word with
 [Qwen3-ForcedAligner-0.6B](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B-hf) and returns them
-under a `words` key. Audio of any length works: the pipeline transcribes it in 8-18 s chunks cut at
-pauses (the model trained on clips up to 19 s), aligns each chunk against its own transcript in
-batches, and places every word on the recording's timeline.
+under a `words` key. Audio of any length works: the pipeline transcribes it in chunks of up to 28 s
+cut at pauses (the model trained on clips up to 30 s), aligns each chunk against its own transcript
+in batches, and places every word on the recording's timeline.
 
 ```python
 result = pipe("audio.wav", return_timestamps=True)
@@ -323,12 +178,12 @@ support), the transcript is still returned and the error is reported under
 ## Train Your Own
 
 Everything that produced this model is open: the code, the data mix, and the recipe. Swap the
-encoder, the LLM, or the projector from config and train on your own data. A smoke test runs on a
-laptop in about five minutes:
+encoder, the LLM, or the projector from config and train on your own data. The recipe is the repo's
+default config:
 
 ```bash
 git clone https://github.com/alexkroman/tiny-audio.git && cd tiny-audio && poetry install
-poetry run python scripts/train.py +experiments=mps_smoke
+poetry run python scripts/train.py
 ```
 
 The
@@ -336,39 +191,48 @@ The
 walks through how the pieces fit, training a model, evaluating it, and publishing it with a demo
 like [this one](https://huggingface.co/spaces/mazesmazes/tiny-audio).
 
+## Serving and Streaming
+
+For more than one-off calls, the [training repo](https://github.com/alexkroman/tiny-audio) ships a
+batched HTTP server (`ta serve`) with an async jobs API for long recordings, and the model class
+exposes token-by-token streaming output via `ASRModel.generate_streaming`. See the
+[repo README](https://github.com/alexkroman/tiny-audio#as-an-http-api-on-runpod) for setup and the
+request format.
+
 ## Model Specifications
 
 | Specification    | Value                                                              |
 | ---------------- | ------------------------------------------------------------------ |
 | Input            | Audio (16kHz mono)                                                 |
 | Output           | Punctuated, capitalized text with formatted numbers                |
-| Max Audio Length | ~19 s per plain call; any length with timestamps or speakers       |
-| Vocabulary       | Qwen3.5 tokenizer                                                  |
+| Max Audio Length | Any length; audio over 28 s is transcribed in chunks cut at pauses |
+| Vocabulary       | Qwen3 tokenizer                                                    |
 | Languages        | English only                                                       |
 | Generation       | Greedy decoding (num_beams=1, do_sample=False), max 256 new tokens |
+| Optional output  | Word-level timestamps; speaker-labelled words and speaker segments |
+| Weights          | ~2.1 GB in bf16 (encoder + projector + decoder)                    |
 
 ## Training Details
 
-|                   |                                                                                                                                                                 |
-| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **Data**          | LibriHeavy (medium), People's Speech (clean), Common Voice 17, GigaSpeech (M), SPGISpeech (M), VoxPopuli, AMI (IHM + SDM), TED-LIUM: one pass, ~3.8M utterances |
-| **Hardware**      | Single NVIDIA H100 80GB                                                                                                                                         |
-| **Steps**         | 58,861                                                                                                                                                          |
-| **Batch Size**    | 64                                                                                                                                                              |
-| **Optimizer**     | AdamW (fused), cosine schedule, 1,000 warmup steps, no weight decay                                                                                             |
-| **Learning Rate** | 1e-3 (projector), 1e-4 (LoRA)                                                                                                                                   |
-| **Precision**     | bf16 (projector held in fp32)                                                                                                                                   |
+|                   |                                                                                                                                                     |
+| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------- |
+| **Data**          | LibriHeavy (medium), People's Speech (clean), Common Voice 17, GigaSpeech (M), SPGISpeech (M), AMI (IHM + SDM), TED-LIUM, NOTSOFAR, EdAcc: one pass |
+| **Steps**         | 16,159                                                                                                                                              |
+| **Batch Size**    | 304                                                                                                                                                 |
+| **Optimizer**     | AdamW (fused), cosine schedule, 1,000 warmup steps, no weight decay                                                                                 |
+| **Learning Rate** | 1e-3 (projector), 2e-5 (decoder)                                                                                                                    |
+| **Precision**     | bf16 mixed precision (weights held in fp32)                                                                                                         |
 
-The training recipe is
-[`configs/experiments/granite_qwen_frozen.yaml`](https://github.com/alexkroman/tiny-audio/blob/main/configs/experiments/granite_qwen_frozen.yaml).
+The training recipe is the default config,
+[`configs/config.yaml`](https://github.com/alexkroman/tiny-audio/blob/main/configs/config.yaml) with
+[`configs/training/production.yaml`](https://github.com/alexkroman/tiny-audio/blob/main/configs/training/production.yaml).
 
 ## Limitations
 
 - **English only**: Not trained on other languages.
 - **Sample rate**: Expects 16kHz audio (other rates are resampled automatically).
-- **Audio length**: A plain `pipe(audio)` call decodes the clip in one pass and works best up to
-  about 19 seconds (the training length). For longer audio pass `return_timestamps=True` or
-  `return_speakers=True`, which transcribe in 8-18 s chunks automatically.
+- **Audio length**: Audio longer than 28 seconds is cut into chunks at pauses and transcribed chunk
+  by chunk, so context does not carry across chunk boundaries.
 - **Speaker diarization**: At most 8 speakers per recording. The count is detected automatically;
   `num_speakers` and `max_speakers` can cap it, but `min_speakers` is not supported (passing it
   raises a `ValueError`). Needs transformers `main` until the next release.
@@ -382,22 +246,22 @@ The training recipe is
 | File                                                               | Description                                              |
 | ------------------------------------------------------------------ | -------------------------------------------------------- |
 | `config.json`                                                      | Model configuration                                      |
-| `model.safetensors`                                                | Projector weights (~50MB)                                |
-| `adapter_config.json` / `adapter_model.safetensors`                | LoRA adapters for the decoder (~270MB)                   |
+| `model.safetensors`                                                | Projector and fine-tuned decoder weights (~2.4 GB, fp32) |
 | `preprocessor_config.json`                                         | Audio preprocessing config                               |
 | `tokenizer.json` / `tokenizer_config.json` / `chat_template.jinja` | Tokenizer                                                |
 | `asr_*.py`, `projectors.py`, `alignment.py`, `diarization.py`      | Custom model code (loaded with `trust_remote_code=True`) |
 
-Only the projector and LoRA weights are stored here. The encoder (Granite Speech) and decoder
-(Qwen3.5-2B) are downloaded from their own Hugging Face repos.
+The encoder (Granite Speech) is downloaded from its own Hugging Face repo. `config.json` loads the
+weights in bf16, as the eval harness runs them.
 
-The previous GLM-ASR + Qwen3-0.6B model is still available at revision `glm-asr-qwen3-0.6b`:
+Earlier models are still available by revision: `granite-qwen-frozen-4` (Granite Speech 470M +
+Qwen3.5-2B with LoRA, MIT) and `glm-asr-qwen3-0.6b` (GLM-ASR + Qwen3-0.6B):
 
 ```python
 pipe = pipeline(
     "automatic-speech-recognition",
     model="mazesmazes/tiny-audio",
-    revision="glm-asr-qwen3-0.6b",
+    revision="granite-qwen-frozen-4",
     trust_remote_code=True,
 )
 ```
@@ -407,10 +271,10 @@ pipe = pipeline(
 If you use this model, please cite:
 
 ```bibtex
-@misc{tinyaudio2024,
+@misc{tinyaudio2025,
   author = {Alex Kroman},
   title = {Tiny Audio: Minimal ASR Training},
-  year = {2024},
+  year = {2025},
   publisher = {GitHub},
   url = {https://github.com/alexkroman/tiny-audio}
 }
@@ -418,15 +282,18 @@ If you use this model, please cite:
 
 ## Acknowledgments
 
-- [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc) for the
+- [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc-nc) for the
   audio encoder
-- [Qwen3.5](https://huggingface.co/Qwen/Qwen3.5-2B) for the language model
+- [Qwen3](https://huggingface.co/Qwen/Qwen3-0.6B) for the language model
 - [Qwen3-ForcedAligner](https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B-hf) for word timestamps
 - [Nemotron-3-Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization) for speaker
   diarization
-- The LibriHeavy, People's Speech, Common Voice, GigaSpeech, SPGISpeech, VoxPopuli, AMI, and
-  TED-LIUM teams for training data
+- The LibriHeavy, People's Speech, Common Voice, GigaSpeech, SPGISpeech, AMI, TED-LIUM, NOTSOFAR,
+  and EdAcc teams for training data
 
 ## License
 
-MIT
+CC BY-NC-SA 4.0 (non-commercial, share-alike), inherited from the
+[Granite Speech 5.0 470M non-causal encoder](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc-nc)
+this model is built on. For commercial use, see the MIT-licensed `granite-qwen-frozen-4` revision
+above. The Tiny Audio training code is MIT.

adapter_config.json

@@ -1,55 +0,0 @@
-{
-  "alora_invocation_tokens": null,
-  "alpha_pattern": {},
-  "arrow_config": null,
-  "auto_mapping": null,
-  "base_model_name_or_path": "mazesmazes/tiny-audio",
-  "bias": "none",
-  "corda_config": null,
-  "ensure_weight_tying": false,
-  "eva_config": null,
-  "exclude_modules": null,
-  "fan_in_fan_out": false,
-  "inference_mode": true,
-  "init_lora_weights": true,
-  "layer_replication": null,
-  "layers_pattern": null,
-  "layers_to_transform": null,
-  "loftq_config": {},
-  "lora_alpha": 64,
-  "lora_bias": false,
-  "lora_dropout": 0.0,
-  "lora_ga_config": null,
-  "megatron_config": null,
-  "megatron_core": "megatron.core",
-  "modules_to_save": null,
-  "monteclora_config": null,
-  "peft_type": "LORA",
-  "peft_version": "0.20.0",
-  "qalora_group_size": 16,
-  "r": 64,
-  "rank_pattern": {},
-  "revision": null,
-  "target_modules": [
-    "in_proj_qkv",
-    "in_proj_z",
-    "q_proj",
-    "o_proj",
-    "down_proj",
-    "in_proj_a",
-    "out_proj",
-    "in_proj_b",
-    "gate_proj",
-    "k_proj",
-    "v_proj",
-    "up_proj"
-  ],
-  "target_parameters": null,
-  "task_type": "CAUSAL_LM",
-  "trainable_token_indices": null,
-  "use_bdlora": null,
-  "use_dora": false,
-  "use_qalora": false,
-  "use_rslora": false,
-  "velora_config": null
-}

adapter_model.safetensors

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

alignment.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import itertools
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, TypedDict, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -93,7 +93,7 @@ class QwenForcedAligner:
                 cls.MODEL_ID, dtype=torch.bfloat16, attn_implementation=attn
             )
             # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a method.
-            model.to(device)  # pyright: ignore[reportArgumentType]
+            cast(torch.nn.Module, model).to(device)
             model.eval()
             cls._model = model
             cls._processor = Qwen3ASRProcessor.from_pretrained(cls.MODEL_ID)

asr_attention.py

@@ -12,14 +12,23 @@ import logging
 from typing import Any
 
 import torch
-from transformers import AutoConfig
+from transformers import AutoConfig, PretrainedConfig
 
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
 
 # FlashAttention's kernels are compiled for head dimensions up to 256.
 FLASH_ATTENTION_MAX_HEAD_DIM = 256
 
 
+def active_attn_implementation(config: PretrainedConfig) -> object:
+    """The attention kernel a loaded model runs (e.g. "sdpa"), read from its config.
+
+    transformers exposes it only as the private `_attn_implementation` field,
+    so this is the one place that reads it.
+    """
+    return config._attn_implementation  # noqa: SLF001  # pyright: ignore[reportPrivateUsage]
+
+
 def _max_attention_head_dim(text_config: object) -> int | None:
     """Largest attention head dim across layers, or None if undeterminable.
 
@@ -35,8 +44,9 @@ def _max_attention_head_dim(text_config: object) -> int | None:
     else:
         try:
             dims = [getattr(text_config, "head_dim", None)]
-        except Exception:
-            # Raised by the heterogeneity guard; treated as "unknown".
+        except RuntimeError:
+            # transformers' AmbiguousGlobalPerLayerAttributeError (a RuntimeError)
+            # from the heterogeneity guard; treated as "unknown".
             dims = []
     usable = [d for d in dims if isinstance(d, int) and d > 0]
     return max(usable) if usable else None
@@ -72,6 +82,7 @@ def _has_sliding_window_attention(model_id: str) -> bool:
             "Could not read the config for %s to check for sliding-window "
             "attention; assuming it has some and using eager attention on MPS.",
             model_id,
+            exc_info=True,
         )
         return True
 
@@ -148,9 +159,31 @@ def resolve_attn_implementation(requested: str | None, model_id: str | None = No
             "falling back to sdpa."
         )
         return "sdpa"
+    if not _flash_attention_runs():
+        return "sdpa"
     return requested
 
 
+@functools.cache
+def _flash_attention_runs() -> bool:
+    """Whether flash-attn's kernels actually run, forward and backward, on this GPU.
+
+    Importing proves only that the package is installed: its kernels are
+    compiled per GPU architecture, so a build without this card's (e.g. the
+    Blackwell RTX PRO 6000, sm_120) imports cleanly and fails at the first
+    attention call, mid-load or mid-step. Checked once per process.
+    """
+    try:
+        flash_attn = importlib.import_module("flash_attn")
+        q = torch.randn(1, 8, 2, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        flash_attn.flash_attn_func(q, q, q, causal=True).sum().backward()
+    # AssertionError: torch's own "not compiled with CUDA enabled".
+    except (RuntimeError, ImportError, AttributeError, AssertionError) as exc:
+        logger.warning("flash_attn is installed but its kernels fail here (%s); using sdpa", exc)
+        return False
+    return True
+
+
 def resolve_decoder_attn_implementation(requested: str | None, model_id: str) -> str | None:
     """`resolve_attn_implementation` for the decoder, plus FlashAttention's head-dim limit.
 

asr_config.py

@@ -24,50 +24,75 @@ DEFAULT_ENCODER_CONV_LAYERS: list[ConvLayerSpec] = [(1, 3, 1), (1, 3, 2)]
 # Full path: 100 Hz mel -> 2x FE stacking -> 4x encoder -> 12.5 Hz.
 GRANITE_ENCODER_CONV_LAYERS: list[ConvLayerSpec] = [(0, 2, 2), (0, 2, 2)]
 
+# The decoding knobs ASRConfig owns, with their defaults. ASRModel copies
+# exactly these onto its GenerationConfig, so a new knob is one entry here.
+# 12 is chosen to be inert on real speech: it blocks only exact 12-gram
+# repeats, and natural English -- including genuine stutters, list
+# recitation and "very very very" -- effectively never repeats a
+# 12-token span verbatim. Measured rate of runaway loops is 1 in 10,630
+# rows, so this should fire almost never; it bounds the tail rather
+# than changing normal output.
+GENERATION_DEFAULTS: dict[str, Any] = {
+    "max_new_tokens": 128,
+    "use_cache": True,
+    "no_repeat_ngram_size": 12,
+}
+
 # Encoders whose `input_features` are (batch, time, feature_dim) rather than
 # Whisper/GLM-ASR's (batch, n_mels, mel_len). Conformer-family checkpoints
 # (Granite Speech 5.0, Parakeet, Nemotron) are all time-major.
 _TIME_MAJOR_ENCODER_MARKERS = ("granite-speech", "parakeet", "nemotron")
 
-# Decoders that already ship a native audio-placeholder token. Reusing it is
-# strictly better than adding "<audio>" and resizing, above all when the
-# decoder is frozen: masked_scatter overwrites the placeholder's *input*
-# embedding, but Gemma4 also looks the token up in `embed_tokens_per_layer` to
-# build its per-layer embeddings (PLE), and that row is NOT overwritten. A
-# freshly added row would hand every audio position a random, permanently
-# untrainable PLE vector. Gemma's own "<|audio|>" row is pretrained for
-# exactly this placeholder role.
-_NATIVE_AUDIO_TOKENS = {"gemma-4": "<|audio|>"}
-
-
-def native_audio_token(text_model_id: str | None) -> str | None:
-    """Return the decoder's built-in audio placeholder token, if it has one."""
-    lowered = (text_model_id or "").lower()
-    return next((tok for key, tok in _NATIVE_AUDIO_TOKENS.items() if key in lowered), None)
-
 
 def is_time_major_encoder(audio_model_id: str | None) -> bool:
     """Whether `audio_model_id` expects (batch, time, feature) input_features."""
     return any(m in (audio_model_id or "").lower() for m in _TIME_MAJOR_ENCODER_MARKERS)
 
 
+# Qwen3-ASR's encoder (tiny_audio/qwen3_asr_encoder.py) runs its convolutions on
+# independent 1 s chunks of 100 mel frames, three (pad 1, kernel 3, stride 2)
+# layers each, so a full chunk yields 13 frames: transformers'
+# `_get_feat_extract_output_lengths` for that model, restated here.
+QWEN3_ASR_CHUNK_FRAMES = 100
+QWEN3_ASR_CONV_LAYERS: list[ConvLayerSpec] = [(1, 3, 2), (1, 3, 2), (1, 3, 2)]
+
+
+def is_qwen3_asr_encoder(audio_model_id: str | None) -> bool:
+    """Whether `audio_model_id` names a Qwen3-ASR checkpoint (its encoder is used)."""
+    return "qwen3-asr" in (audio_model_id or "").lower()
+
+
 @overload
 def compute_encoder_output_length(
-    mel_length: int, conv_layers: Sequence[ConvLayerSpec] | None = None
+    mel_length: int,
+    conv_layers: Sequence[ConvLayerSpec] | None = None,
+    chunk_frames: int | None = None,
 ) -> int: ...
 @overload
 def compute_encoder_output_length(
-    mel_length: "torch.Tensor", conv_layers: Sequence[ConvLayerSpec] | None = None
+    mel_length: "torch.Tensor",
+    conv_layers: Sequence[ConvLayerSpec] | None = None,
+    chunk_frames: int | None = None,
 ) -> "torch.Tensor": ...
 def compute_encoder_output_length(
-    mel_length: "int | torch.Tensor", conv_layers: Sequence[ConvLayerSpec] | None = None
+    mel_length: "int | torch.Tensor",
+    conv_layers: Sequence[ConvLayerSpec] | None = None,
+    chunk_frames: int | None = None,
 ) -> "int | torch.Tensor":
     """Apply encoder conv layer formulas to compute output length.
 
     Works with both Python ints and torch tensors of mel lengths; the formula
     `(L + 2*p - (k-1) - 1) // s + 1` per layer is identical for both.
+
+    With `chunk_frames`, the encoder convolves independent chunks of that many
+    mel frames (Qwen3-ASR): every full chunk contributes the same count and the
+    remainder its own, and an empty remainder contributes zero.
     """
     layers = conv_layers if conv_layers is not None else DEFAULT_ENCODER_CONV_LAYERS
+    if chunk_frames:
+        per_chunk = compute_encoder_output_length(chunk_frames, layers)
+        rest = compute_encoder_output_length(mel_length % chunk_frames, layers)
+        return (mel_length // chunk_frames) * per_chunk + rest
     length = mel_length
     for padding, kernel_size, stride in layers:
         length = (length + 2 * padding - (kernel_size - 1) - 1) // stride + 1
@@ -114,14 +139,19 @@ class ASRConfig(transformers.PretrainedConfig):
         # Encoder conv layers: list of (padding, kernel_size, stride) tuples
         # Default is Whisper/GLM-ASR structure: conv1(k=3,s=1,p=1) + conv2(k=3,s=2,p=1)
         encoder_conv_layers: list[ConvLayerSpec] | None = None,
+        # Mel frames per independently convolved encoder chunk (Qwen3-ASR's
+        # 100); None picks it from `audio_model_id`, 0 means no chunking.
+        encoder_chunk_frames: int | None = None,
+        # Pack labelled batches padding-free when the decoder runs FA2
+        # (asr_packing.py). Ignored under SDPA/eager, where it would be wrong.
+        pack_sequences: bool = True,
+        # Longest chunk the pipeline cuts long audio into. Must stay under the
+        # clip length the model trained on: a model trained on <=19 s skipped
+        # and paraphrased 25-30 s chunks (E22 8.00 -> 9.41, 2026-10-09). 18 s
+        # matches that 19 s training cap and is what older checkpoints, whose
+        # config.json lacks this key, keep using.
+        inference_chunk_seconds: float = 18.0,
         audio_sample_rate: int = 16000,
-        # Whether the encoder takes `input_features` as (batch, time, feature)
-        # instead of Whisper/GLM-ASR's (batch, n_mels, mel_len). Only
-        # SpecAugment unpacks all three dims, so this exists to keep
-        # `_mask_input_features` masking the time axis rather than the feature
-        # axis. Left as None it is auto-detected from `audio_model_id`, so a
-        # Granite/Parakeet swap can't silently mask the wrong axis.
-        audio_features_time_major: bool | None = None,
         # Whether to forward the mel padding mask into the audio encoder.
         # This matters a lot for Granite: it uses block attention over fixed
         # 128-frame blocks, so with `padding="longest"` batches the pad frames
@@ -135,16 +165,13 @@ class ASRConfig(transformers.PretrainedConfig):
         # incomparable to those baselines. Flip it explicitly to test.
         encoder_attention_mask: bool | None = None,
         # Placeholder token whose embeddings get replaced by projector output.
-        # Defaults to the decoder's native audio token when it has one (Gemma 4),
-        # otherwise "<audio>", which is added to the tokenizer and requires an
+        # Defaults to "<audio>", which is added to the tokenizer and requires an
         # embedding resize.
         audio_token: str | None = None,
         # dtype for the projector alone. The fp32-master-weights argument only
         # applies to parameters an optimizer actually updates, so pinning the
         # whole stack to float32 to protect a 10M-param projector wastes 2
-        # bytes on every frozen parameter -- 10.4 GiB on a frozen Gemma 4 E2B
-        # (5.12B stored params; "E2B" counts *active* params, and the
-        # per-layer-embedding table alone is 2.35B). Set this to float32 while
+        # bytes on every frozen parameter. Set this to float32 while
         # model_dtype stays bfloat16 to get master-weight precision where it
         # matters at frozen-model memory cost. Defaults to model_dtype, so
         # existing recipes are unchanged.
@@ -164,41 +191,12 @@ class ASRConfig(transformers.PretrainedConfig):
         encoder_dtype: str | None = None,
         projector_pool_stride: int = 4,
         projector_hidden_dim: int | None = None,
-        projector_type: str = "mlp",
-        # Label smoothing applied inside the LM's loss function (not HF Trainer's
-        # LabelSmoother). Train-only — ASRModel.forward zeros it on eval. Routing
-        # smoothing through the loss_function flows through liger's fused linear
-        # CE when apply_liger_kernel_to_qwen3() is active, avoiding the
-        # (B,T,V) fp32 log_softmax materialization that the HF LabelSmoother
-        # path requires (~15GB at B=50/V=152k on Qwen3-0.6B).
-        label_smoothing: float = 0.0,
         # LoRA configuration (for Stage 2 fine-tuning)
         use_lora: bool = False,
         lora_rank: int = 8,  # SALMONN default
         lora_alpha: int = 32,  # SALMONN default (scaling factor 4.0)
         lora_dropout: float = 0.0,
         lora_target_modules: list[str] | None = None,  # Default: all linear layers
-        # Per-module rank/alpha overrides, keyed by the module's leaf name and
-        # matched by PEFT as `(.*\.)?<key>$` against the full module path.
-        #
-        # These exist because `all-linear` is blind to a matrix's shape. On
-        # Qwen3.5 the gated-DeltaNet gates `in_proj_a` / `in_proj_b` are
-        # (16, 2048) -- one scalar per value head -- so a rank-64 adapter on
-        # them is capped at rank 16 by the output dimension and spends
-        # 64*2048 + 16*64 = 132K parameters to express an update that a full
-        # fine-tune of the same matrix would carry in 32K. Measured on the
-        # step-20000 granite_qwen_top4 checkpoint: 4.76M LoRA params across the
-        # 36 gate matrices, 4.03x what full fine-tuning them costs, and the
-        # realised update's effective rank was 6.3 / 7.3 against the 64
-        # allocated.
-        #
-        # PEFT looks up rank_pattern and alpha_pattern INDEPENDENTLY, each
-        # falling back to `r` / `lora_alpha`. Overriding rank alone therefore
-        # changes the scale: alpha/r goes 128/64 = 2.0 to 128/16 = 8.0. Always
-        # set both, and keep the ratio equal to the global one unless the
-        # scale change is the point.
-        lora_rank_pattern: dict[str, int] | None = None,
-        lora_alpha_pattern: dict[str, int] | None = None,
         # Seconds of silence prepended to every clip at INFERENCE. Measured on
         # Peoples (n=500, paired bootstrap): 20.51% -> 19.28% WER, delta -1.22
         # CI [-1.83, -0.64], with utterances losing a leading reference word
@@ -241,23 +239,6 @@ class ASRConfig(transformers.PretrainedConfig):
         # those two carry 33.57M params for 0.27% of the selection's squared
         # gradient -- see unfreeze_encoder_top_layers for the measurement.
         encoder_trainable_post_projections: bool = True,
-        # SpecAugment on mel input (training-only), parameters match
-        # transformers' WhisperConfig / Wav2Vec2 conventions. Most relevant
-        # when the encoder is trainable (`freeze_audio_encoder=False`) —
-        # without augmentation the encoder sees identical mel inputs on
-        # every visit and overfits fast. Standard for ASR encoder fine-
-        # tuning (Whisper, Conformer, wav2vec2 all use it). Applied to
-        # log-mel input where zero is in-distribution (silence);
-        # structurally different from the prior encoder-output ZM which
-        # was removed because zero was OOD for the encoder's emission
-        # distribution. Uses `_compute_mask_indices` from
-        # transformers.models.whisper.modeling_whisper — the same helper
-        # Whisper itself uses, vectorized over the batch and torch.compile
-        # compatible. Default values match Whisper's defaults.
-        apply_spec_augment: bool = False,
-        mask_time_prob: float = 0.05,
-        mask_time_length: int = 10,
-        mask_time_min_masks: int = 2,
         max_new_tokens: int | None = None,
         use_cache: bool | None = None,
         no_repeat_ngram_size: int | None = None,
@@ -270,7 +251,6 @@ class ASRConfig(transformers.PretrainedConfig):
             text_model_id: HuggingFace model ID for text decoder (Qwen)
             attn_implementation: Attention implementation ("flash_attention_2", "sdpa", "eager")
             model_dtype: Model dtype ("bfloat16", "float16", "float32")
-            projector_type: Projector architecture ("mlp")
             use_lora: Enable LoRA adapters for Stage 2 fine-tuning
         """
         # Set default generation parameters (greedy decoding only).
@@ -285,55 +265,42 @@ class ASRConfig(transformers.PretrainedConfig):
         # `_truncate_repetitions` postprocess missed it because its regex is
         # anchored to end-of-string and the loop was truncated mid-phrase.
         #
-        # 12 is chosen to be inert on real speech: it blocks only exact 12-gram
-        # repeats, and natural English -- including genuine stutters, list
-        # recitation and "very very very" -- effectively never repeats a
-        # 12-token span verbatim. Measured rate of runaway loops is 1 in 10,630
-        # rows, so this should fire almost never; it bounds the tail rather
-        # than changing normal output.
-        generation_defaults = {
-            "max_new_tokens": 128,
-            "use_cache": True,
-            "no_repeat_ngram_size": 12,
-        }
-
         self.audio_model_id = audio_model_id
         self.text_model_id = text_model_id
         self.attn_implementation = attn_implementation
         self.model_dtype = model_dtype
         self.encoder_dim = encoder_dim
         self.llm_dim = llm_dim
-        self.encoder_conv_layers = encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
-        self.audio_features_time_major = (
-            is_time_major_encoder(audio_model_id)
-            if audio_features_time_major is None
-            else audio_features_time_major
+        self.encoder_conv_layers: list[ConvLayerSpec] = (
+            encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
+        )
+        self.pack_sequences: bool = pack_sequences
+        self.inference_chunk_seconds: float = inference_chunk_seconds
+        self.encoder_chunk_frames: int = (
+            (QWEN3_ASR_CHUNK_FRAMES if is_qwen3_asr_encoder(audio_model_id) else 0)
+            if encoder_chunk_frames is None
+            else encoder_chunk_frames
         )
-        self.encoder_attention_mask = (
-            is_time_major_encoder(audio_model_id)
+        # Qwen3-ASR's encoder takes the mel mask to drop padding itself.
+        self.encoder_attention_mask: bool = (
+            (is_time_major_encoder(audio_model_id) or is_qwen3_asr_encoder(audio_model_id))
             if encoder_attention_mask is None
             else encoder_attention_mask
         )
-        self.audio_token = audio_token or native_audio_token(text_model_id) or "<audio>"
-        self.projector_dtype = projector_dtype or model_dtype
-        self.encoder_dtype = encoder_dtype or model_dtype
+        self.audio_token: str = audio_token or "<audio>"
+        self.projector_dtype: str = projector_dtype or model_dtype
+        self.encoder_dtype: str = encoder_dtype or model_dtype
         self.audio_sample_rate = audio_sample_rate
         self.projector_pool_stride = projector_pool_stride
         self.projector_hidden_dim = projector_hidden_dim
-        self.projector_type = projector_type
         self.transcribe_prompt = transcribe_prompt
-        self.label_smoothing = label_smoothing
         # LoRA configuration
         self.use_lora = use_lora
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
         self.lora_dropout = lora_dropout
-        # Plain dicts, never None: PEFT's LoraConfig defaults them to {} and
-        # `get_pattern_key` iterates the keys unconditionally.
-        self.lora_rank_pattern = dict(lora_rank_pattern or {})
-        self.lora_alpha_pattern = dict(lora_alpha_pattern or {})
-        self.inference_lead_in_seconds = float(inference_lead_in_seconds)
-        self.lora_target_modules = lora_target_modules or [
+        self.inference_lead_in_seconds: float = float(inference_lead_in_seconds)
+        self.lora_target_modules: list[str] = lora_target_modules or [
             "q_proj",
             "k_proj",
             "v_proj",
@@ -348,17 +315,13 @@ class ASRConfig(transformers.PretrainedConfig):
         self.freeze_audio_encoder = freeze_audio_encoder
         self.encoder_trainable_top_layers = encoder_trainable_top_layers
         self.encoder_trainable_post_projections = encoder_trainable_post_projections
-        self.apply_spec_augment = apply_spec_augment
-        self.mask_time_prob = mask_time_prob
-        self.mask_time_length = mask_time_length
-        self.mask_time_min_masks = mask_time_min_masks
 
         explicit_generation_args = {
             "max_new_tokens": max_new_tokens,
             "use_cache": use_cache,
             "no_repeat_ngram_size": no_repeat_ngram_size,
         }
-        for key, default in generation_defaults.items():
+        for key, default in GENERATION_DEFAULTS.items():
             value = explicit_generation_args[key]
             setattr(self, key, value if value is not None else default)
 
@@ -395,17 +358,17 @@ class ASRConfig(transformers.PretrainedConfig):
 
         # Point encoder to audio_config so pipeline uses correct feature extractor
         # The pipeline looks for config.encoder._name_or_path for feature extractor
-        self.encoder = self.audio_config
+        self.encoder: transformers.PretrainedConfig | dict[str, Any] = self.audio_config
 
         # The Hub infers the model's task from the first model class listed, so
         # AutoModelForSpeechSeq2Seq must precede AutoModel (else: feature-extraction).
-        self.auto_map = {
+        self.auto_map: dict[str, str] = {
             "AutoConfig": "asr_config.ASRConfig",
             "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
             "AutoModel": "asr_modeling.ASRModel",
             "AutoProcessor": "asr_processing.ASRProcessor",
         }
-        self.custom_pipelines = {
+        self.custom_pipelines: dict[str, dict[str, Any]] = {
             "automatic-speech-recognition": {
                 "impl": "asr_pipeline.ASRPipeline",
                 "pt": ["AutoModelForSpeechSeq2Seq"],
@@ -413,7 +376,7 @@ class ASRConfig(transformers.PretrainedConfig):
                 "type": "audio",
             }
         }
-        self.architectures = ["ASRModel"]
+        self.architectures: list[str] | None = ["ASRModel"]
         self.pipeline_tag = "automatic-speech-recognition"
 
 

asr_layers.py

@@ -1,108 +1,10 @@
-"""Module-surgery helpers for ASRModel: MPS-safe embeddings and encoder unfreezing.
+"""Module-surgery helpers for ASRModel: encoder unfreezing.
 
-Two self-contained groups of `nn.Module` utilities that ASRModel applies to its
-loaded sub-models:
-
-- `mps_unsafe_parameters` / `ChunkedEmbedding` / `chunk_oversized_embeddings`
-  work around PyTorch's 32-bit indexing limit in its Metal kernels.
-- `find_encoder_layer_stack` / `unfreeze_encoder_top_layers` locate an
-  encoder's transformer-block list and switch its top blocks back to trainable.
+`find_encoder_layer_stack` / `unfreeze_encoder_top_layers` locate an encoder's
+transformer-block list and switch its top blocks back to trainable.
 """
 
-import math
-
-import torch
-import torch.nn as nn
-from torch.nn import functional
-
-# PyTorch's Metal kernels index tensor storage with 32-bit offsets, so a gather
-# into a tensor holding more than INT32_MAX elements wraps around and silently
-# returns the wrong rows -- no error, no warning, just corrupt values.
-MPS_MAX_TENSOR_ELEMENTS = 2**31 - 1
-
-
-def mps_unsafe_parameters(model: nn.Module) -> list[tuple[str, int]]:
-    """Parameters too large for MPS to index correctly, as (name, numel) pairs.
-
-    Gemma 4 E2B trips this on `embed_tokens_per_layer`: its per-layer embedding
-    table is 262144 x (35*256) = 2,348,810,240 elements, 201M past INT32_MAX.
-    On MPS the lookup returns garbage, which Gemma then adds into all 35 layers
-    at every position. The failure is entirely silent -- the model loads, runs,
-    and emits fluent, confident, wrong transcripts. Measured on this stack:
-    teacher-forced loss 0.14 on CPU versus 3.80 on MPS from identical weights.
-
-    Returns an empty list for models that are safe, so callers can treat a
-    non-empty result as "this model needs `chunk_oversized_embeddings`".
-    """
-    return [
-        (name, param.numel())
-        for name, param in model.named_parameters()
-        if param.numel() > MPS_MAX_TENSOR_ELEMENTS
-    ]
-
-
-class ChunkedEmbedding(nn.Module):
-    """Drop-in `nn.Embedding` that splits its table to stay within 32-bit indexing.
-
-    The overflow lives in the kernel's arithmetic over the weight's *storage*,
-    so slicing is not enough -- views share that storage and stay broken. Each
-    chunk therefore has to be its own contiguous allocation. Gathering per
-    chunk and concatenating on the feature axis reproduces the full gather
-    exactly: verified bit-identical to CPU on MPS, where the single-tensor
-    lookup was off by up to 3.16.
-
-    Wraps rather than replaces the semantics of `Gemma4TextScaledWordEmbedding`,
-    whose forward multiplies the lookup by `embed_scale` -- dropping that would
-    shrink every per-layer embedding by sqrt(hidden_size_per_layer_input).
-    """
-
-    def __init__(self, embedding: nn.Embedding) -> None:
-        """Split `embedding`'s table into MPS-safe column chunks, keeping its metadata."""
-        super().__init__()
-        self.num_embeddings = embedding.num_embeddings
-        self.embedding_dim = embedding.embedding_dim
-        self.padding_idx = embedding.padding_idx
-        # Gemma 4's scaled embedding carries this; a plain nn.Embedding does not.
-        self.embed_scale = getattr(embedding, "embed_scale", None)
-
-        num_chunks = math.ceil(embedding.weight.numel() / MPS_MAX_TENSOR_ELEMENTS)
-        device = embedding.weight.device
-        # Chunk from a CPU copy: carving contiguous pieces out of an
-        # already-oversized MPS tensor would run the same overflowing kernel
-        # this class exists to avoid.
-        source = embedding.weight.detach().cpu()
-        self.chunks = nn.ParameterList(
-            [
-                nn.Parameter(
-                    chunk.contiguous().to(device), requires_grad=embedding.weight.requires_grad
-                )
-                for chunk in torch.chunk(source, num_chunks, dim=1)
-            ]
-        )
-        del source
-
-    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        """Look up `input_ids` in every chunk and concatenate along the feature dim."""
-        out = torch.cat([functional.embedding(input_ids, chunk) for chunk in self.chunks], dim=-1)
-        if self.embed_scale is not None:
-            out = out * self.embed_scale.to(out.dtype)
-        return out
-
-    @property
-    def weight(self) -> torch.Tensor:
-        """Reassembled table, for callers that introspect `.weight`.
-
-        Materializes the full tensor, so this is a debugging/serialization
-        convenience rather than something to call on a hot path.
-        """
-        return torch.cat(list(self.chunks), dim=1)
-
-    def extra_repr(self) -> str:
-        """Describe the table like `nn.Embedding` does, plus the chunk count."""
-        return (
-            f"{self.num_embeddings}, {self.embedding_dim}, "
-            f"chunks={len(self.chunks)} (MPS int32-indexing workaround)"
-        )
+from torch import nn
 
 
 def find_encoder_layer_stack(encoder: nn.Module) -> tuple[str, nn.ModuleList] | None:
@@ -182,19 +84,3 @@ def unfreeze_encoder_top_layers(
             param.requires_grad_(True)
             unfrozen.append(f"{name}.{pname}")
     return sorted(unfrozen)
-
-
-def chunk_oversized_embeddings(root: nn.Module) -> list[str]:
-    """Replace every embedding MPS cannot index with a `ChunkedEmbedding`.
-
-    Returns the qualified names that were replaced, empty when nothing needed
-    it. Idempotent: `ChunkedEmbedding` is not an `nn.Embedding`, so a second
-    pass finds nothing.
-    """
-    replaced: list[str] = []
-    for module_name, module in list(root.named_modules()):
-        for child_name, child in list(module.named_children()):
-            if isinstance(child, nn.Embedding) and child.weight.numel() > MPS_MAX_TENSOR_ELEMENTS:
-                setattr(module, child_name, ChunkedEmbedding(child))
-                replaced.append(f"{module_name}.{child_name}" if module_name else child_name)
-    return replaced

asr_modeling.py

@@ -1,5 +1,4 @@
 import contextlib
-import functools
 import importlib.util
 import inspect
 import json
@@ -13,9 +12,9 @@ from threading import Thread
 from typing import TYPE_CHECKING, Any, Self, cast, overload
 
 import torch
-import torch.nn as nn
-from peft import LoraConfig, PeftModel, get_peft_model
+from peft import LoraConfig, PeftModel, get_peft_model, set_peft_model_state_dict
 from safetensors.torch import load_file
+from torch import nn
 from torch.nn import functional
 from transformers import (
     AutoFeatureExtractor,
@@ -23,6 +22,10 @@ from transformers import (
     AutoModelForCausalLM,
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
+    GenerationConfig,
+    LogitsProcessor,
+    LogitsProcessorList,
+    PretrainedConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
     SequenceFeatureExtractor,
@@ -38,15 +41,20 @@ from transformers.utils.hub import (
 
 if TYPE_CHECKING:
     from transformers import GraniteSpeech5Encoder as _GraniteSpeech5Encoder
-    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
-    from transformers.models.gemma4.modeling_gemma4 import (
-        Gemma4ForCausalLM as _Gemma4ForCausalLM,
-    )
-    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
 
-    from .asr_attention import resolve_attn_implementation, resolve_decoder_attn_implementation
-    from .asr_config import ASRConfig, compute_encoder_output_length
-    from .asr_layers import chunk_oversized_embeddings, unfreeze_encoder_top_layers
+    from .asr_attention import (
+        active_attn_implementation,
+        resolve_attn_implementation,
+        resolve_decoder_attn_implementation,
+    )
+    from .asr_config import (
+        GENERATION_DEFAULTS,
+        ASRConfig,
+        compute_encoder_output_length,
+        is_qwen3_asr_encoder,
+    )
+    from .asr_layers import unfreeze_encoder_top_layers
+    from .asr_packing import pack_rows
     from .asr_processing import (
         DEFAULT_TRANSCRIBE_PROMPT,
         ASRProcessor,
@@ -62,16 +70,26 @@ if TYPE_CHECKING:
         OutputLengthProjector,
         StateDictT,
     )
-    from .projectors import PROJECTOR_CLASSES, MLPAudioProjector
+    from .projectors import MLPAudioProjector
+    from .qwen3_asr_encoder import Qwen3ASRAudioEncoder
 
     # None at runtime on a transformers without the architecture; see below.
     GraniteSpeech5Encoder: type[_GraniteSpeech5Encoder] | None
-    Gemma4ForCausalLM: type[_Gemma4ForCausalLM] | None
 else:
     try:
-        from .asr_attention import resolve_attn_implementation, resolve_decoder_attn_implementation
-        from .asr_config import ASRConfig, compute_encoder_output_length
-        from .asr_layers import chunk_oversized_embeddings, unfreeze_encoder_top_layers
+        from .asr_attention import (
+            active_attn_implementation,
+            resolve_attn_implementation,
+            resolve_decoder_attn_implementation,
+        )
+        from .asr_config import (
+            GENERATION_DEFAULTS,
+            ASRConfig,
+            compute_encoder_output_length,
+            is_qwen3_asr_encoder,
+        )
+        from .asr_layers import unfreeze_encoder_top_layers
+        from .asr_packing import pack_rows
         from .asr_processing import (
             DEFAULT_TRANSCRIBE_PROMPT,
             ASRProcessor,
@@ -79,11 +97,22 @@ else:
             render_audio_prompt,
         )
         from .asr_types import LoadStateDictResult, OutputLengthProjector, StateDictT
-        from .projectors import PROJECTOR_CLASSES
+        from .projectors import MLPAudioProjector
+        from .qwen3_asr_encoder import Qwen3ASRAudioEncoder
     except ImportError:  # flat layout on the Hub: sibling modules, no package
-        from asr_attention import resolve_attn_implementation, resolve_decoder_attn_implementation
-        from asr_config import ASRConfig, compute_encoder_output_length
-        from asr_layers import chunk_oversized_embeddings, unfreeze_encoder_top_layers
+        from asr_attention import (
+            active_attn_implementation,
+            resolve_attn_implementation,
+            resolve_decoder_attn_implementation,
+        )
+        from asr_config import (
+            GENERATION_DEFAULTS,
+            ASRConfig,
+            compute_encoder_output_length,
+            is_qwen3_asr_encoder,
+        )
+        from asr_layers import unfreeze_encoder_top_layers
+        from asr_packing import pack_rows
         from asr_processing import (
             DEFAULT_TRANSCRIBE_PROMPT,
             ASRProcessor,
@@ -91,22 +120,19 @@ else:
             render_audio_prompt,
         )
         from asr_types import LoadStateDictResult, OutputLengthProjector, StateDictT
-        from projectors import PROJECTOR_CLASSES
+        from projectors import MLPAudioProjector
+        from qwen3_asr_encoder import Qwen3ASRAudioEncoder
 
-    # Architectures newer than the transformers floor in pyproject.toml. Kept
-    # optional so a checkpoint that needs neither still loads on an older
-    # install; the branch that needs one raises at its use site instead.
+    # Newer than the transformers floor in pyproject.toml. Kept optional so a
+    # checkpoint that does not need it still loads on an older install; the
+    # branch that needs it raises at its use site instead.
     try:
         from transformers import GraniteSpeech5Encoder
     except ImportError:  # transformers < 5.16: no native granite_speech5
         GraniteSpeech5Encoder = None
-    try:
-        from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM
-    except ImportError:  # transformers without the gemma4 architecture
-        Gemma4ForCausalLM = None
 
 
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
 
 # Vocab dimension the embedding table is padded to after adding the audio
 # token. 128 is what transformers' own `pad_to_multiple_of` docs recommend for
@@ -220,44 +246,34 @@ def _splice_audio(
     )
 
 
-def _patch_gemma_decode_loop(model: "GenerativeDecoder") -> None:
-    """Make a Gemma4ForCausalLM drop `per_layer_inputs` after the first step.
-
-    `Gemma4ForConditionalGeneration` does this itself; the causal LM does not.
-    Without it, step 2 of generation hands PLE to forward alongside
-    `input_ids` and raises "You cannot specify per_layer_inputs if input_ids is
-    provided". Dropping is semantically correct rather than a workaround: PLE
-    describes the prompt, and from step 2 on generation feeds a single freshly
-    sampled token whose PLE Gemma derives from that token's own id.
+class TensorNoRepeatNGram(LogitsProcessor):
+    """`NoRepeatNGramLogitsProcessor`'s bans, computed as tensor ops on the ids' device.
 
-    Patched on the instance because subclassing Gemma4ForCausalLM trips
-    transformers' experts-implementation check during __init__
-    ("does not support setting experts implementation").
+    The stock processor rebuilds a Python dict of every n-gram in every row
+    each step, with a `.tolist()` (a GPU sync) per row: 0.3-0.7 ms of host work
+    a step at batch 8-32 against a 6.8 ms CUDA-graphed decode step. Same
+    semantics: a token is banned when the last n-1 ids plus it would repeat an
+    n-gram already in the row, prompt included.
     """
-    inner_prepare = model.prepare_inputs_for_generation
-
-    # functools.wraps is load-bearing, not cosmetic. Generation's
-    # _prepare_model_inputs decides whether a model can accept `inputs_embeds`
-    # by testing for it in
-    # inspect.signature(prepare_inputs_for_generation).parameters. A bare
-    # *args/**kwargs wrapper hides the parameter, and generate() then refuses
-    # inputs_embeds with "doesn't have its forwarding implemented" -- fatal
-    # here, since audio only ever reaches the decoder as inputs_embeds.
-    # `wraps` sets __wrapped__, which inspect.signature follows.
-    @functools.wraps(inner_prepare)
-    def prepare_inputs_for_generation(
-        *args: Any, is_first_iteration: bool = False, **kwargs: Any
-    ) -> dict[str, Any]:
-        """Drop stale `per_layer_inputs` on every decode step after the first."""
-        model_inputs = inner_prepare(*args, is_first_iteration=is_first_iteration, **kwargs)
-        if not is_first_iteration:
-            model_inputs.pop("per_layer_inputs", None)
-        return model_inputs
 
-    model.prepare_inputs_for_generation = prepare_inputs_for_generation
+    def __init__(self, ngram_size: int) -> None:
+        self.ngram_size = ngram_size
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        n = self.ngram_size
+        if input_ids.shape[1] < n:
+            return scores
+        windows = input_ids.unfold(1, n, 1)  # (batch, positions, n)
+        tail = input_ids[:, None, input_ids.shape[1] - n + 1 :]
+        repeats = (windows[..., :-1] == tail).all(dim=-1)
+        # Non-repeating windows point at a spare column, so every write is True.
+        vocab = scores.shape[-1]
+        banned = torch.zeros(scores.shape[0], vocab + 1, dtype=torch.bool, device=scores.device)
+        banned.scatter_(1, torch.where(repeats, windows[..., -1], vocab), True)
+        return cast("torch.FloatTensor", scores.masked_fill(banned[:, :vocab], -float("inf")))
 
 
-def _assert_projector_loaded(incompatible_keys: LoadStateDictResult, projector_type: str) -> None:
+def _assert_projector_loaded(incompatible_keys: LoadStateDictResult) -> None:
     """Fail loudly when a checkpoint's projector doesn't match the built one.
 
     `save_pretrained` serializes `projector.*` (plus the language model when
@@ -304,8 +320,7 @@ def _assert_projector_loaded(incompatible_keys: LoadStateDictResult, projector_t
         )
 
     raise RuntimeError(
-        f"Projector weights do not match the configured projector "
-        f"(projector_type={projector_type!r}): " + "; ".join(detail) + hint
+        "Projector weights do not match the configured projector: " + "; ".join(detail) + hint
     )
 
 
@@ -337,7 +352,7 @@ def _log_linear_attention_backends(level: int = logging.WARNING) -> None:
 class ASRModel(PreTrainedModel, GenerationMixin):
     """Audio-to-text model combining an audio encoder, projector, and language model."""
 
-    config_class = ASRConfig
+    config_class: type[PretrainedConfig] | None = ASRConfig
     base_model_prefix = "model"
     main_input_name = "input_features"
     _supports_flash_attn_2 = True
@@ -387,7 +402,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             if model_file is not None:
                 state_dict = load_file(model_file)
                 incompatible = model.load_state_dict(state_dict, strict=False)
-                _assert_projector_loaded(incompatible, getattr(config, "projector_type", "?"))
+                _assert_projector_loaded(incompatible)
 
             # Load LoRA adapters if use_lora is enabled
             if getattr(config, "use_lora", False):
@@ -442,7 +457,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         target_dtype = getattr(torch, config.model_dtype)
 
         # Audio encoder (frozen)
-        self.audio_tower = self._load_audio_encoder(config, target_dtype)
+        self.audio_tower: nn.Module = self._load_audio_encoder(config, target_dtype)
 
         # Language model (frozen). Typed as the decoder even with LoRA attached:
         # a PeftModel forwards every attribute it does not define to the
@@ -472,7 +487,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # text_model_id, liger missing (it is linux-only), or a patcher
             # that patched a different class than the one that got loaded.
             logger.log(
-                logging.INFO if type(self)._is_loading_from_pretrained else logging.WARNING,
+                logging.INFO if self._is_loading_from_pretrained else logging.WARNING,
                 "%s's forward does not accept `skip_logits` — liger's fused "
                 "linear cross-entropy is NOT active, so every training step "
                 "materializes a (batch, seq, %s) logits tensor. Check for an "
@@ -485,26 +500,25 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self._init_tokenizer(config)
 
         # Set up generation config with greedy decoding defaults
-        self.generation_config = self.language_model.generation_config
-        self.generation_config.max_new_tokens = config.max_new_tokens
-        self.generation_config.use_cache = config.use_cache
+        self.generation_config: GenerationConfig = self.language_model.generation_config
         # `generate()` reads the GenerationConfig and nothing else, so every
-        # decoding knob ASRConfig owns has to be copied across here. Omitting
-        # this line left ASRConfig's 12-gram loop guard inert: the value was
+        # decoding knob ASRConfig owns has to be copied across here. A missing
+        # copy line once left ASRConfig's 12-gram loop guard inert: the value was
         # set on the config, serialized into config.json, and never consulted.
-        self.generation_config.no_repeat_ngram_size = config.no_repeat_ngram_size
+        for key in GENERATION_DEFAULTS:
+            setattr(self.generation_config, key, getattr(config, key))
         self.generation_config.eos_token_id = self._resolve_eos_token_ids()
         self.generation_config.pad_token_id = self.tokenizer.pad_token_id
 
         # Feature extractor for audio preprocessing
-        self.feature_extractor = self._create_feature_extractor(config)
+        self.feature_extractor: SequenceFeatureExtractor = self._create_feature_extractor(config)
 
         # Audio projector (trainable unless freeze_projector is set)
-        self.projector = self._create_projector(config, target_dtype)
+        self.projector: MLPAudioProjector = self._create_projector(config, target_dtype)
 
         # Setup LoRA if enabled (Stage 2 fine-tuning)
         # Skip if loading from pretrained - from_pretrained will handle adapter loading
-        if getattr(config, "use_lora", False) and not type(self)._is_loading_from_pretrained:
+        if getattr(config, "use_lora", False) and not self._is_loading_from_pretrained:
             self._setup_lora(config)
 
         # Freeze projector if specified (for Stage 2 LoRA-only training)
@@ -530,12 +544,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         feature_extractor: SequenceFeatureExtractor = AutoFeatureExtractor.from_pretrained(
             config.audio_model_id
         )
-        # Whisper's encoder requires a fixed 3000 mel frames (30s) and the
-        # feature extractor pads to that by default — leave it alone. Other
-        # encoders (e.g. GLM-ASR) accept variable-length input, so we disable
-        # padding to avoid wasting compute on silent frames.
-        if "whisper" not in config.audio_model_id.lower():
-            feature_extractor.padding = False  # pyright: ignore[reportAttributeAccessIssue]
         return feature_extractor
 
     @classmethod
@@ -557,7 +565,12 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         }
 
         encoder: nn.Module
-        if "whisper" in config.audio_model_id.lower():
+        if is_qwen3_asr_encoder(config.audio_model_id):
+            # Only the audio tower of the ASR checkpoint; see qwen3_asr_encoder.py.
+            encoder = Qwen3ASRAudioEncoder.from_pretrained(
+                config.audio_model_id, dtype, encoder_kwargs["attn_implementation"]
+            )
+        elif "whisper" in config.audio_model_id.lower():
             full_model = WhisperModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
             encoder = full_model.encoder
             del full_model
@@ -633,8 +646,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # trainable blocks; it must be fp32 for Adam's step to survive
         # rounding. See ASRConfig.encoder_dtype for the arithmetic.
         encoder_dtype = _resolve_dtype(getattr(config, "encoder_dtype", None), dtype)
-        # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a method.
-        encoder = encoder.to(dtype=encoder_dtype)  # pyright: ignore[reportCallIssue]
+        # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a
+        # method; it casts in place and returns the module itself.
+        cast(nn.Module, encoder).to(dtype=encoder_dtype)
         if getattr(config, "freeze_audio_encoder", True):
             encoder.requires_grad_(False)
             encoder.train(False)  # equivalent to .eval(); avoids a security hook false-positive
@@ -714,95 +728,18 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             AutoModelForCausalLM.from_pretrained(config.text_model_id, **decoder_kwargs),
         )
 
-        # Gemma 4 checkpoints are natively multimodal: AutoModelForCausalLM maps
-        # `gemma4` to Gemma4ForConditionalGeneration, which is both wasteful
-        # (a vision tower + audio tower this model never calls) and unusable
-        # for inference. Convert it to the text-only causal LM.
-        #
-        # Gated on the model id, not on the `.model.language_model` layout:
-        # that layout is the generic transformers composite-model shape, shared
-        # by GLM-ASR and Llava/Qwen-Omni style checkpoints. Probing for it
-        # structurally would funnel those into _gemma_text_only, which imports
-        # transformers.models.gemma4 and transplants weights into a
-        # Gemma4ForCausalLM shell built from a foreign text_config. The
-        # hasattr check is kept as a guard so an already-text-only gemma-4
-        # checkpoint passes through untouched.
-        gemma_inner = getattr(decoder, "model", None)
-        is_gemma4 = "gemma-4" in (config.text_model_id or "").lower()
-        if is_gemma4 and gemma_inner is not None and hasattr(gemma_inner, "language_model"):
-            decoder = cls._gemma_text_only(decoder)
         # See _load_audio_encoder note: idempotent post-load cast to dodge the
         # FA2 "current dype is fp32" warning when from_pretrained's dtype kwarg
         # isn't fully propagated to every submodule.
         # `to` returns the module itself, so the cast happens in place.
         # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a method.
-        decoder.to(dtype=dtype)  # pyright: ignore[reportCallIssue]
+        cast(nn.Module, decoder).to(dtype=dtype)
         decoder.config.use_cache = getattr(config, "use_cache", True)
         if getattr(config, "freeze_language_model", True):
             decoder.requires_grad_(False)
             decoder.train(False)
         return decoder
 
-    @staticmethod
-    def _gemma_text_only(loaded: "GenerativeDecoder") -> "GenerativeDecoder":
-        """Turn a loaded Gemma4ForConditionalGeneration into a text-only CausalLM.
-
-        Two independent problems force this transplant rather than just loading
-        the right class directly:
-
-        1. Only `Gemma4ForCausalLM` supports `inputs_embeds` in `.generate()`.
-           Its `prepare_inputs_for_generation` accepts the argument; the
-           multimodal wrapper's does not, so generation dies with "You passed
-           `inputs_embeds` to `.generate()`, but the model class
-           Gemma4ForConditionalGeneration doesn't have its forwarding
-           implemented." Audio reaches this decoder exclusively as
-           inputs_embeds, so the wrapper cannot do inference at all.
-
-        2. `Gemma4ForCausalLM.from_pretrained` cannot read the published
-           checkpoint. Its keys are `model.language_model.*` (the multimodal
-           layout) while the causal LM expects `model.*`, and gemma4 ships no
-           checkpoint conversion mapping. Critically it fails *silently*:
-           every tensor is reported MISSING and randomly initialized
-           (embedding std lands at initializer_range), producing a model that
-           trains without error and learns nothing.
-
-        So the wrapper is loaded for its correct key layout, then its text
-        model and lm_head are moved into a shell built on the meta device --
-        which allocates nothing, because every real parameter is transplanted.
-        Dropping the wrapper also frees the vision and audio towers: the
-        returned text-only model holds 4,628,569,344 parameters against the
-        checkpoint's 5,123,178,979, so ~494M are released.
-        """
-        if Gemma4ForCausalLM is None:
-            msg = f"Gemma 4 decoders need the gemma4 architecture; found {transformers_version}."
-            raise ImportError(msg)
-
-        # Only reached for gemma-4 checkpoints, whose text config is Gemma4TextConfig.
-        text_config = cast("Gemma4TextConfig", loaded.config.get_text_config())
-        with torch.device("meta"):
-            shell = Gemma4ForCausalLM(text_config)
-
-        shell.model = cast("Gemma4TextModel", loaded.get_submodule("model.language_model"))
-        shell.lm_head = cast(nn.Linear, loaded.get_submodule("lm_head"))
-        shell.generation_config = loaded.generation_config
-        shell.tie_weights()
-
-        # See _patch_gemma_decode_loop for why PLE is dropped after step 1.
-        # Gemma4ForCausalLM is a PreTrainedModel with GenerationMixin, which is
-        # all `GenerativeDecoder` describes.
-        decoder = cast("GenerativeDecoder", shell)
-        _patch_gemma_decode_loop(decoder)
-
-        stranded = [n for n, t in shell.named_parameters() if t.device.type == "meta"]
-        stranded += [n for n, t in shell.named_buffers() if t.device.type == "meta"]
-        if stranded:
-            msg = (
-                "Gemma text-only conversion left parameters on the meta device "
-                f"(would silently produce garbage): {stranded[:5]}"
-            )
-            raise RuntimeError(msg)
-        return decoder
-
     def _create_projector(self, config: ASRConfig, dtype: torch.dtype) -> "MLPAudioProjector":
         """Create the trainable audio projector."""
         # Auto-detect dimensions if not specified
@@ -817,8 +754,8 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         if config.llm_dim is None:
             dec_cfg = self.language_model.config
-            # Composite configs (Gemma 4) keep hidden_size on `text_config`;
-            # get_text_config() returns self for plain decoders like Qwen3.
+            # Composite configs keep hidden_size on `text_config`;
+            # get_text_config() returns self for plain decoders.
             if hasattr(dec_cfg, "get_text_config"):
                 dec_cfg = dec_cfg.get_text_config()
             config.llm_dim = getattr(dec_cfg, "hidden_size", None) or getattr(
@@ -828,16 +765,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 msg = "Could not auto-detect llm_dim. Please specify in config."
                 raise ValueError(msg)
 
-        # Select projector type based on config
-        projector_type = getattr(config, "projector_type", "mlp")
-        projector_class = PROJECTOR_CLASSES.get(projector_type)
-        if projector_class is None:
-            msg = (
-                f"Unknown projector_type: {projector_type}. "
-                f"Valid options: {list(PROJECTOR_CLASSES.keys())}"
-            )
-            raise ValueError(msg)
-        projector = projector_class(config)
+        projector = MLPAudioProjector(config)
 
         # Move projector to same device as language model (important when using quantization)
         device = next(self.language_model.parameters()).device
@@ -855,13 +783,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             lora_alpha=config.lora_alpha,
             target_modules=config.lora_target_modules,
             lora_dropout=config.lora_dropout,
-            # Per-module overrides for targets whose shape makes the global
-            # rank wrong -- see ASRConfig.lora_rank_pattern. PEFT resolves the
-            # two patterns independently, so a recipe that sets one and not the
-            # other silently changes that module's alpha/r scale; the config
-            # field's docstring spells out the arithmetic.
-            rank_pattern=dict(getattr(config, "lora_rank_pattern", None) or {}),
-            alpha_pattern=dict(getattr(config, "lora_alpha_pattern", None) or {}),
             bias="none",
             task_type="CAUSAL_LM",
         )
@@ -891,7 +812,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 self.tokenizer.pad_token = self.tokenizer.eos_token
 
         # Audio placeholder token. Prefer a token the tokenizer already knows
-        # (Gemma 4 ships "<|audio|>") so no embedding resize is needed at all.
+        # so no embedding resize is needed at all.
         # `unk_token_id` is the sentinel convert_tokens_to_ids returns for an
         # unknown token, so it means "not in vocab" rather than a usable id.
         self.audio_token = getattr(config, "audio_token", None) or "<audio>"
@@ -1033,7 +954,8 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 tokenize=False,
                 add_generation_prompt=False,
             )
-        except Exception:
+        except Exception:  # any template that cannot render this probe has no usable tail
+            logger.debug("Chat template probe failed", exc_info=True)
             return None
         if not isinstance(rendered, str) or sentinel not in rendered:
             return None
@@ -1048,33 +970,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         first = tail_ids[0]
         return first if first in set(self.tokenizer.all_special_ids) else None
 
-    def _apply(self, *args: Any, **kwargs: Any) -> Self:
-        """Repair MPS-unsafe embeddings whenever the model lands on MPS.
-
-        Device placement happens after construction -- `pipeline` builds the
-        model, then moves it -- so the check belongs on the move, not in
-        __init__. `_apply` is the single funnel every relocation goes through
-        (`.to()`, `.cuda()`, `.mps()`, dtype casts), which keeps this off the
-        many individual entry points.
-
-        Without it, Gemma 4 E2B on MPS reads the wrong rows out of its 2.35B-
-        element per-layer embedding table and transcribes confident nonsense at
-        full speed. Chunking keeps the model on the GPU -- falling back to CPU
-        would be correct but costs roughly an order of magnitude in latency.
-        """
-        module: Self = super()._apply(*args, **kwargs)
-        try:
-            on_mps = any(p.device.type == "mps" for p in module.parameters())
-        except StopIteration:  # pragma: no cover - parameterless model
-            return module
-        if on_mps and (replaced := chunk_oversized_embeddings(module)):
-            logger.warning(
-                "MPS cannot index %s with 32-bit offsets; split into chunks to "
-                "keep the lookup correct on GPU.",
-                ", ".join(replaced),
-            )
-        return module
-
     def train(self, mode: bool = True) -> Self:
         """Set train/eval mode, but keep frozen submodules out of train mode.
 
@@ -1097,7 +992,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # statistics anyway. Granite Speech 5.0 carries 16 BatchNorm1d and
             # zero Dropout, and for non-Whisper extractors the collator pads to
             # the batch longest (`DataCollator._audio_padding`) with
-            # group_by_length off, so clips spanning 0.8-19.0s leave a large
+            # group_by_length off, so clips spanning 0.1-30s leave a large
             # pad fraction. Granite's conv module zeroes pad positions and THEN
             # runs BatchNorm1d over the full padded (B, C, T); BN reduces over
             # B*T, so those zeros are counted as data -- deflating the mean and
@@ -1223,8 +1118,30 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             projector=self.projector,
             encoder_conv_layers=self.config.encoder_conv_layers,
             audio_token=self.audio_token,
+            encoder_chunk_frames=self.config.encoder_chunk_frames,
         )
 
+    def load_lora_checkpoint(self, checkpoint_dir: str | PathLike[str]) -> bool:
+        """Load the LoRA weights a checkpoint saved as adapter_model.safetensors.
+
+        HF Trainer resumes by loading model.safetensors (the projector, see
+        `state_dict`) and, when that file exists, never reaches its PEFT
+        branch, so a resumed run kept the trained projector but restarted LoRA
+        from its initialization under the trained optimizer state
+        (granite-qwen-verbatim-aug-2 resumed at step 52,000 on 2026-10-19:
+        loss 0.28 -> 1.7). Returns whether adapter weights were loaded; raises
+        if a saved tensor matches no LoRA parameter.
+        """
+        path = Path(checkpoint_dir) / "adapter_model.safetensors"
+        if not path.is_file() or not isinstance(self.language_model, PeftModel):
+            return False
+        result = set_peft_model_state_dict(self.language_model, load_file(path))
+        if result.unexpected_keys:
+            msg = f"adapter weights match no LoRA parameter: {result.unexpected_keys[:3]}"
+            raise ValueError(msg)
+        logger.info("Loaded LoRA adapter weights from %s", path)
+        return True
+
     @overload
     def state_dict(
         self, *, destination: StateDictT, prefix: str = ..., keep_vars: bool = ...
@@ -1301,50 +1218,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         destination.update({f"{prefix}{k}": v for k, v in sd.items()})
         return destination
 
-    def _ple_text_model(self) -> "Gemma4TextModel | None":
-        """Locate the submodule owning `get_per_layer_inputs`, or None.
-
-        Only Gemma 4 style decoders have one. The attribute sits on
-        Gemma4TextModel, which is `language_model.model.language_model` for the
-        multimodal checkpoint and `language_model.model` for a text-only one.
-        """
-        candidates = (
-            getattr(getattr(self.language_model, "model", None), "language_model", None),
-            getattr(self.language_model, "model", None),
-            self.language_model,
-        )
-        for candidate in candidates:
-            if candidate is not None and hasattr(candidate, "get_per_layer_inputs"):
-                # Checked just above; only Gemma 4 text models define it.
-                return cast("Gemma4TextModel", candidate)
-        return None
-
-    def _per_layer_kwargs(self, input_ids: torch.Tensor | None) -> dict[str, torch.Tensor]:
-        """Precompute Gemma 4 per-layer embeddings (PLE) from clean `input_ids`.
-
-        Gemma 4 builds a token-identity PLE component by looking `input_ids` up
-        in `embed_tokens_per_layer`. It must be computed *before* masked_scatter
-        overwrites the audio positions, because Gemma's `inputs_embeds`-only
-        path tries to recover ids by reversing the main embedding -- an exact
-        row match against `embed_tokens.weight * sqrt(hidden_size)`. Projector
-        output matches no row, so that path raises RuntimeError (and would be
-        an O(B*T*vocab*hidden) comparison even if it matched).
-
-        Computing it here keeps every text token's PLE row intact and leaves
-        audio positions on the audio placeholder's own row, which is exactly
-        what Gemma does natively for image/audio placeholder tokens.
-
-        Returns {} for decoders without PLE, so callers can splat it blindly.
-        """
-        if input_ids is None:
-            return {}
-        text_model = self._ple_text_model()
-        if text_model is None:
-            return {}
-        if not getattr(text_model.config, "hidden_size_per_layer_input", None):
-            return {}
-        return {"per_layer_inputs": text_model.get_per_layer_inputs(input_ids, None)}
-
     def _embed_tokens(self, input_ids: torch.Tensor | None) -> torch.Tensor:
         """Look `input_ids` up in the decoder's token embedding table."""
         embeds: torch.Tensor = self.language_model.get_input_embeddings()(input_ids)
@@ -1358,6 +1231,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         return compute_encoder_output_length(
             audio_attention_mask.sum(dim=-1),
             self.config.encoder_conv_layers,
+            self.config.encoder_chunk_frames,
         )
 
     def _encode_audio(
@@ -1377,17 +1251,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         Returns:
             Flattened audio embeddings of shape (sum(expected_token_counts), hidden_dim).
         """
-        # SpecAugment is applied on the mel input, training-only. Most useful
-        # when the encoder is trainable; on the frozen-encoder path it still
-        # perturbs the projector's input slightly but with no gradient flowing
-        # back to the encoder to leverage the diversity.
-        if (
-            self.training
-            and getattr(self.config, "apply_spec_augment", False)
-            and audio_features.numel() > 0
-        ):
-            audio_features = self._mask_input_features(audio_features, audio_attention_mask)
-
         # When the encoder is frozen, skip gradient tracking through it — cuts
         # activation memory and matches the prior published recipe's behavior.
         # When trainable, we MUST allow gradients to flow back to encoder
@@ -1434,161 +1297,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
         return _gather_audio_embeds(audio_embeds, token_counts, max_tokens)
 
-    def _mask_input_features(
-        self,
-        input_features: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """SpecAugment on mel input (pure-torch, vectorized, compile-ready).
-
-        Follows the same semantics as
-        `transformers.models.whisper.modeling_whisper.WhisperModel._mask_input_features`
-        (wav2vec2-style mask sampling: sample N start positions per sample,
-        mask `mask_length` frames forward from each), but reimplemented in
-        pure torch so it stays inside the autograd graph without crossing
-        the numpy boundary. This avoids inductor codegen failures
-        (e.g. the `'zuf0' was not declared` error from the prior numpy ->
-        torch.tensor round-trip) AND avoids the per-forward host-to-GPU
-        sync that the numpy path required.
-
-        One minor semantic divergence vs the upstream helper: this version
-        allows mask spans to overlap, while upstream rejects overlapping
-        samples. For ASR purposes this is irrelevant — occasional region
-        double-coverage has no measurable effect on the regularization
-        signal.
-
-        Reads ASRConfig fields by Whisper naming convention: mask_time_prob,
-        mask_time_length, mask_time_min_masks.
-
-        Args:
-            input_features: (batch, n_mels, mel_len) log-mel features.
-            attention_mask: (batch, mel_len) mel-frame padding mask. Like
-                upstream, it confines time-axis span count and start
-                positions to each sample's unpadded length. Without it,
-                spans are drawn uniformly over the padded axis, so under
-                `padding="longest"` a 1s clip batched with a 19s clip gets
-                roughly 1/19 of its spans in real audio — short utterances
-                go effectively unaugmented while long ones take the full
-                dose. Whisper is unaffected either way (fixed 3000 frames);
-                the variable-length encoders are the ones that need it.
-                Pass None to treat the whole axis as valid.
-
-        Returns:
-            Same-shape tensor with time-axis masks zeroed.
-        """
-        input_features = input_features.clone()
-        config = self.config
-        # Whisper/GLM-ASR hand us (batch, n_mels, mel_len); Conformer encoders
-        # (Granite, Parakeet, Nemotron) hand us (batch, time, feature_dim).
-        # Masking the wrong axis is silent -- it still runs and still zeroes
-        # cells -- so the axis order is read from config rather than guessed.
-        time_major = getattr(config, "audio_features_time_major", False)
-        if time_major:
-            batch_size, sequence_length, _ = input_features.size()
-            feature_dim = 2
-        else:
-            batch_size, _, sequence_length = input_features.size()
-            feature_dim = 1
-        device = input_features.device
-
-        valid_lengths = None
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != sequence_length:
-                msg = (
-                    "SpecAugment attention_mask time axis "
-                    f"({attention_mask.shape[-1]}) does not match input_features "
-                    f"({sequence_length}). Masking against the wrong axis is silent, "
-                    "so this is raised rather than ignored."
-                )
-                raise ValueError(msg)
-            valid_lengths = attention_mask.to(device).sum(dim=-1)
-
-        if getattr(config, "mask_time_prob", 0.0) > 0:
-            mask_time = self._sample_mask_indices(
-                batch_size,
-                sequence_length,
-                mask_prob=config.mask_time_prob,
-                mask_length=config.mask_time_length,
-                min_masks=config.mask_time_min_masks,
-                device=device,
-                valid_lengths=valid_lengths,
-            )
-            # Broadcast (B, T) over the feature axis to mask all bins at
-            # masked times: (B, 1, T) feature-major, (B, T, 1) time-major.
-            input_features.masked_fill_(mask_time.unsqueeze(feature_dim), 0)
-
-        return input_features
-
-    @staticmethod
-    def _sample_mask_indices(
-        batch_size: int,
-        axis_length: int,
-        mask_prob: float,
-        mask_length: int,
-        min_masks: int,
-        device: torch.device,
-        valid_lengths: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Vectorized SpecAugment mask sampler — torch.compile-friendly.
-
-        Returns a (batch_size, axis_length) bool tensor where True marks
-        a position covered by at least one mask span. Spans may overlap
-        (see _mask_input_features docstring on the semantic difference vs
-        the upstream Whisper helper).
-
-        `valid_lengths` is a (batch_size,) tensor of unpadded lengths along
-        this axis; both the span count and the start positions are derived
-        per sample from it, so every span lands inside real content. Pass
-        None for an axis with no padding.
-        """
-        # Worst-case span count, from the full axis. Computed as a Python int
-        # from axis_length (an upper bound on any sample's valid length)
-        # rather than from valid_lengths.max(), which would need a .item()
-        # call and reintroduce the per-forward GPU->host sync this function
-        # was rewritten to avoid. Surplus spans are masked off per sample
-        # below, so over-allocating here is free apart from memory.
-        #
-        # Matches the upstream formula (ignoring the epsilon noise term, which
-        # only shifts the count by ±1 stochastically — negligible at the
-        # default mask_time_prob=0.05 / mask_length=10 setting which gives
-        # ~5 spans for a typical 1500-frame mel input).
-        max_spans = max(int(mask_prob * axis_length / mask_length + 0.5), min_masks)
-        if max_spans == 0:
-            return torch.zeros(batch_size, axis_length, device=device, dtype=torch.bool)
-
-        if valid_lengths is None:
-            lengths = torch.full((batch_size,), axis_length, device=device, dtype=torch.long)
-        else:
-            lengths = valid_lengths.to(device=device, dtype=torch.long).clamp(
-                min=1, max=axis_length
-            )
-
-        # Per-sample span count, so a short clip is not given a long clip's
-        # dose of masking. Bounded above by max_spans since lengths <= axis_length.
-        num_spans = torch.clamp(
-            (mask_prob * lengths.float() / mask_length + 0.5).long(), min=min_masks
-        )  # (B,)
-
-        # Sample start positions independently per sample x span, inside that
-        # sample's own valid range. Clamp so a span of length mask_length never
-        # starts past the end of short content. torch.rand is [0, 1), so the
-        # floor below never reaches max_start.
-        max_start = torch.clamp(lengths - mask_length + 1, min=1)  # (B,)
-        starts = (
-            torch.rand(batch_size, max_spans, device=device) * max_start.unsqueeze(1)
-        ).long()  # (B, N)
-
-        # Drop the spans a sample didn't earn (its num_spans < max_spans).
-        span_active = torch.arange(max_spans, device=device).unsqueeze(0) < num_spans.unsqueeze(1)
-
-        # For each (sample, span, position), True iff position ∈ [start, start+mask_length).
-        positions = torch.arange(axis_length, device=device).view(1, 1, -1)  # (1, 1, T)
-        starts_b = starts.unsqueeze(-1)  # (B, N, 1)
-        span_mask = (positions >= starts_b) & (positions < starts_b + mask_length)
-        span_mask &= span_active.unsqueeze(-1)
-        # Reduce over the span dim: True if ANY span covers this position.
-        return span_mask.any(dim=1)
-
     def forward(
         self,
         input_ids: torch.Tensor | None = None,
@@ -1608,10 +1316,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         if inputs_embeds is None:
             inputs_embeds = self._embed_tokens(input_ids)
 
-        # Gemma 4 PLE: derived from input_ids, so it must be captured before the
-        # masked_scatter below replaces the audio rows. No-op ({}) elsewhere.
-        ple_kwargs = self._per_layer_kwargs(input_ids)
-
         if input_features is not None and input_ids is not None:
             is_audio_token = input_ids == self.audio_token_id
             if audio_token_counts is None:
@@ -1627,14 +1331,22 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
             inputs_embeds = _splice_audio(inputs_embeds, is_audio_token, audio_embeds)
 
-        # Forward label_smoothing to the LM's loss_function via **kwargs.
-        # transformers.loss.loss_utils.ForCausalLMLoss → fixed_cross_entropy
-        # forwards extra kwargs to F.cross_entropy, which accepts label_smoothing.
-        # When apply_liger_kernel_to_qwen3() has patched the LM, the smoothing
-        # is consumed by liger's fused linear CE (no (B,T,V) materialization).
-        # Zeroed on eval so eval/loss is raw CE and comparable to LS=0 runs.
-        if labels is not None and self.training and self.config.label_smoothing > 0:
-            kwargs.setdefault("label_smoothing", self.config.label_smoothing)
+        # Padding-free packing (asr_packing.py): only the real tokens reach the
+        # decoder, so no attention or MLP compute is spent on padding.
+        if (
+            labels is not None
+            and attention_mask is not None
+            and self._can_pack(position_ids, past_key_values)
+        ):
+            packed = pack_rows(inputs_embeds, attention_mask, labels)
+            inputs_embeds, labels = packed.inputs_embeds, packed.labels
+            position_ids, attention_mask = packed.position_ids, None
+            kwargs.update(
+                cu_seq_lens_q=packed.cu_seq_lens,
+                cu_seq_lens_k=packed.cu_seq_lens,
+                max_length_q=packed.max_length,
+                max_length_k=packed.max_length,
+            )
 
         # Ask for the dataclass output explicitly. Every consumer below and in
         # generate() reads `outputs.loss` / `outputs.logits` by attribute, so
@@ -1653,11 +1365,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # the (B, T, V) tensor is pure cost on both train and eval.
         #
         # Requesting it EXPLICITLY on both, rather than leaning on liger's
-        # default, because that default is wrong here in both directions:
-        #
-        #   skip_logits = self.training and labels is not None
-        #
-        # where `self` is the LANGUAGE MODEL. Two ways that misfires:
+        # default, because that default is wrong here in both directions. Liger
+        # sets `skip_logits = self.training and labels is not None`, where
+        # `self` is the LANGUAGE MODEL. Two ways that misfires:
         #
         #   eval  -- runs under model.eval(), so the default is False and the
         #            full tensor gets built. granite_qwen.yaml's claim that
@@ -1681,11 +1391,25 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Done here rather than via `use_liger_kernel: true` because that flag
         # makes transformers call apply_liger_kernel(), which *raises* when
         # liger-kernel is unavailable -- and liger is a linux-only dependency
-        # (pyproject.toml), so it would break every mps_smoke run on a mac.
+        # (pyproject.toml), so it would break every local run on a mac.
         # Gating on the patched signature is equivalent where it matters and
         # inert everywhere else.
         if labels is not None and self._lm_accepts_skip_logits:
             kwargs.setdefault("skip_logits", True)
+        elif labels is not None:
+            # No liger (a mac, or liger missing): the stock forward projects
+            # every position through lm_head, audio and prompt included, though
+            # only the transcript carries labels. Keep just the positions that
+            # predict a label and hand the loss its already-shifted targets,
+            # which ForCausalLMLoss takes as-is. Same loss, same
+            # num_items_in_batch normalization, a (B, K, V) logits tensor with
+            # K ~ the longest transcript instead of T ~ the whole prompt. On
+            # Qwen3.5's 248K vocab that tensor and its fp32 upcast and gradient
+            # are what pushed a batch-4 mac run past 40 GB.
+            shift_labels = nn.functional.pad(labels, (0, 1), value=-100)[:, 1:]
+            keep = (shift_labels != -100).any(dim=0).nonzero().squeeze(-1)
+            kwargs.setdefault("logits_to_keep", keep)
+            kwargs.setdefault("shift_labels", shift_labels[:, keep])
 
         outputs: CausalLMOutputWithPast = self.language_model(
             attention_mask=attention_mask,
@@ -1695,11 +1419,23 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             labels=labels,
             use_cache=use_cache,
             cache_position=cache_position,
-            **ple_kwargs,
             **kwargs,
         )
         return outputs
 
+    def _can_pack(self, position_ids: torch.Tensor | None, past_key_values: object) -> bool:
+        """Whether a labelled, padded batch can run padding-free: uncached and on FA2.
+
+        Only flash_attention_2 keeps packed rows apart (by their cumulative
+        lengths); SDPA or eager would let them attend to each other.
+        """
+        return (
+            bool(getattr(self.config, "pack_sequences", False))
+            and position_ids is None
+            and past_key_values is None
+            and active_attn_implementation(self.language_model.config) == "flash_attention_2"
+        )
+
     def prepare_inputs_for_generation(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
         """Prepare inputs for generation, handling audio features for cached decoding."""
         input_features = kwargs.pop("input_features", None)
@@ -1732,7 +1468,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         input_ids, attention_mask = left_pad_prompt_rows(rows, self.tokenizer)
         return input_ids.to(device), attention_mask.to(device)
 
-    def _prepare_audio_inputs(
+    def prepare_audio_inputs(
         self,
         input_features: torch.Tensor,
         audio_attention_mask: torch.Tensor,
@@ -1793,10 +1529,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         """
         if attention_mask is None or attention_mask.device.type != "mps":
             return
-        if (
-            self.language_model.config._attn_implementation  # pyright: ignore[reportPrivateUsage]  # transformers has no public accessor
-            != "sdpa"
-        ):
+        if active_attn_implementation(self.language_model.config) != "sdpa":
             return
         if not bool((attention_mask == 0).any()):
             return
@@ -1811,7 +1544,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
         raise ValueError(msg)
 
-    @torch.no_grad()  # pyright: ignore[reportUntypedFunctionDecorator]
     def generate(
         self,
         *args: Any,
@@ -1831,79 +1563,83 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         signature stays compatible with `GenerationMixin.generate`, whose
         positional parameters mean something else entirely.
         """
-        if args:
-            msg = (
-                "ASRModel.generate takes input_ids, input_features, "
-                "audio_attention_mask and attention_mask as keyword arguments"
-            )
-            raise TypeError(msg)
-        if input_features is None:
-            msg = "input_features required for generation"
-            raise ValueError(msg)
-        if audio_attention_mask is None:
-            msg = "audio_attention_mask required for generation"
-            raise ValueError(msg)
+        # A context manager, not the @torch.no_grad() decorator, which pyright
+        # cannot type and so would erase this signature.
+        with torch.no_grad():
+            if args:
+                msg = (
+                    "ASRModel.generate takes input_ids, input_features, "
+                    "audio_attention_mask and attention_mask as keyword arguments"
+                )
+                raise TypeError(msg)
+            if input_features is None:
+                msg = "input_features required for generation"
+                raise ValueError(msg)
+            if audio_attention_mask is None:
+                msg = "audio_attention_mask required for generation"
+                raise ValueError(msg)
 
-        input_ids, attention_mask, inputs_embeds = self._prepare_audio_inputs(
-            input_features,
-            audio_attention_mask,
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-        )
-        self._assert_sdpa_safe_on_mps(attention_mask)
-
-        # Fold generation flags (max_new_tokens, num_beams, output_scores, ...)
-        # into a private copy of generation_config instead of passing them
-        # alongside it: transformers v5 still applies such kwargs but warns that
-        # combining them with `generation_config` is deprecated. update() returns
-        # the kwargs that aren't config fields (streamer, logits_processor, ...),
-        # which still go to generate() directly.
-        gen_cfg = deepcopy(self.generation_config)
-        generate_kwargs = gen_cfg.update(**generate_kwargs)
-        # output_scores requires return_dict_in_generate for HF generate to
-        # actually populate .scores on the output object.
-        if gen_cfg.output_scores:
-            gen_cfg.return_dict_in_generate = True
-
-        lm_inputs = self._lm_generate_inputs(input_ids, inputs_embeds)
-        output = self.language_model.generate(
-            **lm_inputs,
-            attention_mask=attention_mask,
-            generation_config=gen_cfg,
-            **generate_kwargs,
-        )
+            input_ids, attention_mask, inputs_embeds = self.prepare_audio_inputs(
+                input_features,
+                audio_attention_mask,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            self._assert_sdpa_safe_on_mps(attention_mask)
+
+            # Fold generation flags (max_new_tokens, num_beams, output_scores, ...)
+            # into a private copy of generation_config instead of passing them
+            # alongside it: transformers v5 still applies such kwargs but warns that
+            # combining them with `generation_config` is deprecated. update() returns
+            # the kwargs that aren't config fields (streamer, logits_processor, ...),
+            # which still go to generate() directly.
+            gen_cfg = deepcopy(self.generation_config)
+            generate_kwargs = gen_cfg.update(**generate_kwargs)
+            # output_scores requires return_dict_in_generate for HF generate to
+            # actually populate .scores on the output object.
+            if gen_cfg.output_scores:
+                gen_cfg.return_dict_in_generate = True
+            if gen_cfg.no_repeat_ngram_size:
+                processors = LogitsProcessorList(
+                    generate_kwargs.pop("logits_processor", None) or []
+                )
+                processors.append(TensorNoRepeatNGram(gen_cfg.no_repeat_ngram_size))
+                generate_kwargs["logits_processor"] = processors
+                # 0, not None: generate refills None from the decoder's own config
+                # and would build the stock processor alongside this one.
+                gen_cfg.no_repeat_ngram_size = 0
+
+            lm_inputs = self._lm_generate_inputs(input_ids, inputs_embeds)
+            output = self.language_model.generate(
+                **lm_inputs,
+                attention_mask=attention_mask,
+                generation_config=gen_cfg,
+                **generate_kwargs,
+            )
 
-        # With input_ids AND inputs_embeds, generate returns the full sequence
-        # (prompt + generated), so the prompt must be stripped. With
-        # inputs_embeds alone it already returns only the new tokens, so
-        # stripping would eat real output -- hence prompt_len = 0 there.
-        # When scores were requested, preserve the GenerateOutput so callers
-        # can read .scores; otherwise return the bare tensor for backward
-        # compatibility with existing callers.
-        prompt_len = input_ids.shape[1] if "input_ids" in lm_inputs else 0
-        if isinstance(output, torch.Tensor):
-            # Slicing a token-id tensor keeps its int64 dtype.
-            return cast("torch.LongTensor", output[:, prompt_len:])
-        output.sequences = cast("torch.LongTensor", output.sequences[:, prompt_len:])
-        return output
+            # With input_ids AND inputs_embeds, generate returns the full sequence
+            # (prompt + generated), so the prompt must be stripped. With
+            # inputs_embeds alone it already returns only the new tokens, so
+            # stripping would eat real output -- hence prompt_len = 0 there.
+            # When scores were requested, preserve the GenerateOutput so callers
+            # can read .scores; otherwise return the bare tensor for backward
+            # compatibility with existing callers.
+            prompt_len = input_ids.shape[1] if "input_ids" in lm_inputs else 0
+            if isinstance(output, torch.Tensor):
+                # Slicing a token-id tensor keeps its int64 dtype.
+                return cast("torch.LongTensor", output[:, prompt_len:])
+            output.sequences = cast("torch.LongTensor", output.sequences[:, prompt_len:])
+            return output
 
     def _lm_generate_inputs(
         self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor
     ) -> dict[str, torch.Tensor]:
         """The prompt arguments for `language_model.generate`.
 
-        Default path passes both input_ids and inputs_embeds so
-        repetition_penalty and no_repeat_ngram_size see the prompt (they need
-        input_ids to track which tokens have been used). Gemma 4 forbids that
-        combination outright -- "You must specify exactly one of input_ids or
-        inputs_embeds" -- so there we pass embeds plus the PLE precomputed from
-        the clean prompt ids. Gemma's prepare_inputs_for_generation drops
-        per_layer_inputs after the first step, so it is correct to hand it to
-        generate().
+        Passes both input_ids and inputs_embeds so repetition_penalty and
+        no_repeat_ngram_size see the prompt (they need input_ids to track which
+        tokens have been used).
         """
-        ple_kwargs = self._per_layer_kwargs(input_ids)
-        if ple_kwargs:
-            return {"inputs_embeds": inputs_embeds, **ple_kwargs}
         return {"input_ids": input_ids, "inputs_embeds": inputs_embeds}
 
     def generate_streaming(
@@ -1925,7 +1661,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         Yields:
             Partial transcript text as each token is generated
         """
-        input_ids, attention_mask, inputs_embeds = self._prepare_audio_inputs(
+        input_ids, attention_mask, inputs_embeds = self.prepare_audio_inputs(
             input_features, audio_attention_mask
         )
 
@@ -1997,11 +1733,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         save_dir = Path(save_directory)
         save_dir.mkdir(parents=True, exist_ok=True)
 
-        # Update config with actual vocab size. Composite decoder configs
-        # (Gemma 4 is multimodal, so Gemma4Config holds text/vision/audio
-        # sub-configs) have no top-level `vocab_size` -- reading it raises
-        # AttributeError, which used to surface only at the first checkpoint
-        # save, i.e. after the run had already spent real training time.
+        # Update config with actual vocab size. Composite decoder configs hold
+        # text/vision/audio sub-configs and have no top-level `vocab_size` --
+        # reading it raises AttributeError, which used to surface only at the
+        # first checkpoint save, after the run had already spent real time.
         lm_config = self.language_model.config
         if hasattr(lm_config, "get_text_config"):
             lm_config = lm_config.get_text_config()
@@ -2085,7 +1820,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         # Copy source files for auto-loading
         src_dir = Path(__file__).parent
-        siblings = ("projectors.py", "alignment.py", "diarization.py")
+        siblings = ("projectors.py", "alignment.py", "diarization.py", "qwen3_asr_encoder.py")
         for src in [*src_dir.glob("asr_*.py"), *(src_dir / name for name in siblings)]:
             shutil.copy(src, save_dir / src.name)
 
@@ -2135,7 +1870,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         """
         # Store repo_id in config so save_pretrained can access it
         self.config.pretrained_model_path = repo_id
-        # Call parent's push_to_hub
+        # `tags` only when given: transformers reads it as
+        # `kwargs.get("tags", [])` and iterates the result, so an explicit
+        # None is a TypeError after the whole run has trained.
+        extra: dict[str, Any] = {"tags": tags} if tags is not None else {}
         return PreTrainedModel.push_to_hub(
             self,
             repo_id,
@@ -2146,7 +1884,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             revision=revision,
             create_pr=create_pr,
             max_shard_size=max_shard_size,
-            tags=tags,
+            **extra,
         )
 
 

asr_packing.py

@@ -0,0 +1,58 @@
+"""Padding-free packing of a labelled decoder batch for FlashAttention's varlen kernels.
+
+Training batches are padded to their longest row, and the clips run 0.1-30 s,
+so a large share of every step's decoder compute (attention AND MLP) went to
+padding. Packing keeps only the real tokens: the batch's rows are concatenated
+into one `(1, total_tokens)` sequence, and FlashAttention's variable-length
+kernels keep each row's attention inside its own segment via the cumulative
+sequence lengths passed alongside.
+
+The batch's rows, their order and the per-token loss are unchanged:
+
+- each row's positions restart at 0. Training batches are left-padded, so
+  padded rows used to start at an offset; RoPE scores depend only on relative
+  position, and 0-based positions are what `generate()` uses at inference;
+- the first token of every row gets label -100, so no row is trained to
+  predict its first token from the previous row's last one. Every row starts
+  with a prompt token, which is already unlabelled, so this drops nothing; in
+  the padded layout such a target would have been predicted from a fully
+  masked pad position anyway.
+
+Only valid under flash_attention_2: SDPA and eager would see one unmasked
+sequence and let rows attend to each other. ASRModel checks before packing.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class PackedBatch:
+    """One flattened sequence and the varlen metadata FlashAttention needs."""
+
+    inputs_embeds: torch.Tensor  # (1, total_tokens, hidden)
+    labels: torch.Tensor  # (1, total_tokens)
+    position_ids: torch.Tensor  # (1, total_tokens), restarting at 0 per row
+    cu_seq_lens: torch.Tensor  # (rows + 1,) int32 row boundaries
+    max_length: int
+
+
+def pack_rows(
+    inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor
+) -> PackedBatch:
+    """Concatenate the real tokens of a padded `(rows, length)` batch into one sequence."""
+    valid = attention_mask.bool()
+    lengths = valid.sum(dim=-1)
+    position_ids = (valid.long().cumsum(dim=-1) - 1)[valid].unsqueeze(0)
+    packed_labels = labels[valid].clone()
+    starts = torch.cat([lengths.new_zeros(1), lengths.cumsum(dim=0)[:-1]])
+    packed_labels[starts[lengths > 0]] = -100
+    cu_seq_lens = torch.nn.functional.pad(lengths.cumsum(dim=0), (1, 0)).to(torch.int32)
+    return PackedBatch(
+        inputs_embeds=inputs_embeds[valid].unsqueeze(0),
+        labels=packed_labels.unsqueeze(0),
+        position_ids=position_ids,
+        cu_seq_lens=cu_seq_lens,
+        max_length=int(lengths.max()),
+    )

asr_pipeline.py

@@ -1,15 +1,14 @@
 """ASR pipeline for audio-to-text transcription with optional timestamps and diarization."""
 
+import logging
 import re
 from collections.abc import Callable, Iterator, MutableMapping, Sequence
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, TypedDict, cast
 
 import numpy as np
 import numpy.typing as npt
 import torch
 import transformers
-from transformers.pipelines.audio_utils import ffmpeg_read
 
 if TYPE_CHECKING:
     from transformers import PreTrainedTokenizerBase, SequenceFeatureExtractor
@@ -17,12 +16,15 @@ if TYPE_CHECKING:
     from .alignment import QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_processing import (
+        CHUNK_MAX_S,
         audible_chunks,
         chunk_bounds,
         collate_chunks,
+        decode_audio_input,
         is_silent,
         prepend_lead_in,
     )
+    from .asr_text import truncate_repetitions
     from .asr_types import PreparedChunk
     from .diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
 else:
@@ -30,24 +32,30 @@ else:
         from .alignment import QwenForcedAligner
         from .asr_modeling import ASRModel
         from .asr_processing import (
+            CHUNK_MAX_S,
             audible_chunks,
             chunk_bounds,
             collate_chunks,
+            decode_audio_input,
             is_silent,
             prepend_lead_in,
         )
+        from .asr_text import truncate_repetitions
         from .asr_types import PreparedChunk
         from .diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
     except ImportError:  # flat layout on the Hub: sibling modules, no package
         from alignment import QwenForcedAligner
         from asr_modeling import ASRModel
         from asr_processing import (
+            CHUNK_MAX_S,
             audible_chunks,
             chunk_bounds,
             collate_chunks,
+            decode_audio_input,
             is_silent,
             prepend_lead_in,
         )
+        from asr_text import truncate_repetitions
         from asr_types import PreparedChunk
         from diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
 
@@ -61,12 +69,9 @@ __all__ = [
 ]
 
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", flags=re.DOTALL)
-_MIN_REPEATS = 3
-_LOOSE_MIN_REPEATS = 6
-_TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{_MIN_REPEATS - 1},}}$")
-_TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{_MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
-_NON_WORD_RE = re.compile(r"\W+")
 
 
 class _RawAudio(TypedDict):
@@ -127,6 +132,30 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             **kwargs,
         )
 
+    @property
+    def chunk_max_seconds(self) -> float:
+        """Longest chunk to cut long audio into: the model's trained-for length.
+
+        `ASRConfig.inference_chunk_seconds`; checkpoints saved before it existed
+        fall back to the 18 s they were trained for.
+        """
+        value = getattr(self.model.config, "inference_chunk_seconds", None)
+        return float(value) if value else CHUNK_MAX_S
+
+    def _sampling_rate(self) -> int:
+        """The feature extractor's sample rate (`__init__` always sets the extractor)."""
+        if self.feature_extractor is None:
+            msg = "ASRPipeline has no feature extractor"
+            raise RuntimeError(msg)
+        return self.feature_extractor.sampling_rate
+
+    def _tokenizer(self) -> "PreTrainedTokenizerBase":
+        """The tokenizer (`__init__` always sets it; the base class types it Optional)."""
+        if self.tokenizer is None:
+            msg = "ASRPipeline has no tokenizer"
+            raise RuntimeError(msg)
+        return self.tokenizer
+
     def _sanitize_parameters(
         self,
         chunk_length_s: float | None = None,
@@ -246,7 +275,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             return super().__call__(inputs, **kwargs)
         array, sr = audio["array"], audio["sampling_rate"]
 
-        bounds = chunk_bounds(array, sr)
+        bounds = chunk_bounds(array, sr, self.chunk_max_seconds)
         if not kwargs:
             texts = self._transcribe_chunks(audible_chunks(array, bounds, sr), sr)
             return {"text": " ".join(t for t in texts if t)}
@@ -318,7 +347,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             audio_attention_mask=batch["attention_mask"].to(self.model.device),
             **generate_kwargs,
         )
-        assert torch.is_tensor(tokens)  # no output_scores: plain token ids
+        if not isinstance(tokens, torch.Tensor):  # no output_scores: plain token ids
+            msg = f"generate returned {type(tokens).__name__}, expected token ids"
+            raise TypeError(msg)
         return [self.postprocess({"tokens": row})["text"] for row in tokens.cpu()]
 
     def _transcribe_chunk(
@@ -362,7 +393,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         self, array: npt.NDArray[np.float32], sr: int, **kwargs: Any
     ) -> dict[str, Any]:
         """`{"text", "words"}` for the whole recording: chunk, transcribe, align, offset."""
-        bounds = chunk_bounds(array, sr)
+        bounds = chunk_bounds(array, sr, self.chunk_max_seconds)
         texts = self._transcribe_chunks(audible_chunks(array, bounds, sr), sr, **kwargs)
         result: dict[str, Any] = {"text": " ".join(t for t in texts if t)}
 
@@ -376,7 +407,10 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 for (s, _), chunk_words in zip(bounds, aligned, strict=True)
                 for w in chunk_words
             ]
-        except Exception as e:
+        except Exception as e:  # timestamps are optional: return the text without them
+            logger.warning(
+                "Word alignment failed; returning text without timestamps", exc_info=True
+            )
             result["words"] = []
             result["timestamp_error"] = str(e)
         return result
@@ -401,7 +435,8 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         try:
             activity = NemotronDiarizer.activity(array, sample_rate=sr)
             keep = NemotronDiarizer.top_speakers(activity, **diarization_params)
-        except Exception as e:
+        except Exception as e:  # speakers are optional: fall back to the unmasked transcript
+            logger.warning("Diarization failed; transcribing without speakers", exc_info=True)
             result = self._transcribe_aligned(array, sr, **kwargs)
             result["speaker_segments"] = []
             result["diarization_error"] = str(e)
@@ -448,7 +483,10 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 [(ch["audio"], text) for ch, text in zip(chunks, texts, strict=True)],
                 sample_rate=sr,
             )
-        except Exception as e:
+        except Exception as e:  # timestamps are optional: order chunks by start time instead
+            logger.warning(
+                "Word alignment failed; returning text without timestamps", exc_info=True
+            )
             order = sorted(range(len(chunks)), key=lambda k: chunks[k]["start"])
             text = " ".join(texts[k] for k in order if texts[k])
             return {"text": text, "words": [], "timestamp_error": str(e)}
@@ -463,9 +501,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 return None
             array, sr = inputs[key], inputs.get("sampling_rate", 16000)
         elif isinstance(inputs, (str, bytes)):
-            # File path or encoded bytes - decode with ffmpeg (same as HF pipeline)
-            data = Path(inputs).read_bytes() if isinstance(inputs, str) else inputs
-            array, sr = ffmpeg_read(data, sampling_rate=16000), 16000
+            array, sr = decode_audio_input(inputs), 16000
         elif isinstance(inputs, np.ndarray):
             array, sr = inputs, 16000
         else:
@@ -486,11 +522,10 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         inputs = _single_input(args, preprocess_params, "inputs")
         # Handle dict with "array" key (from datasets)
         if isinstance(inputs, dict) and "array" in inputs:
-            assert self.feature_extractor is not None  # always set by __init__
             sample = inputs
             inputs = {
                 "raw": sample["array"],
-                "sampling_rate": sample.get("sampling_rate", self.feature_extractor.sampling_rate),
+                "sampling_rate": sample.get("sampling_rate", self._sampling_rate()),
             }
 
         # Inference-only lead-in silence. See prepend_lead_in for the
@@ -508,11 +543,10 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             else 0.0
         )
         if lead_in > 0 and isinstance(inputs, dict) and "raw" in inputs:
-            assert self.feature_extractor is not None  # always set by __init__
             padded: dict[str, Any] = dict(inputs)
             padded["raw"] = prepend_lead_in(
                 padded["raw"],
-                padded.get("sampling_rate", self.feature_extractor.sampling_rate),
+                padded.get("sampling_rate", self._sampling_rate()),
                 lead_in,
             )
             inputs = padded
@@ -584,21 +618,15 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 "Call it with one clip per invocation."
             )
             raise ValueError(msg)
-        top1_logprobs: list[float] = []
-        top2_logprobs: list[float] = []
+        steps: list[list[float]] = []
         if scores:
             # One batched reduction over every step, then one device transfer.
             logits = torch.stack([step_logits[0] for step_logits in scores]).float()
-            steps: list[list[float]] = (
-                logits.log_softmax(dim=-1).topk(k=2, dim=-1).values.cpu().tolist()
-            )
-            for top1, top2 in steps:
-                top1_logprobs.append(top1)
-                top2_logprobs.append(top2)
+            steps = logits.log_softmax(dim=-1).topk(k=2, dim=-1).values.cpu().tolist()
         return {
             "tokens": sequences,
-            "top1_logprob": top1_logprobs,
-            "top2_logprob": top2_logprobs,
+            "top1_logprob": [step[0] for step in steps],
+            "top2_logprob": [step[1] for step in steps],
             "is_last": is_last,
         }
 
@@ -648,22 +676,22 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             token_ids: list[int] = np.asarray(tokens).tolist()
             tokens = [t for t in token_ids if t not in eos_set]
 
-        assert self.tokenizer is not None  # always set by __init__
-        decoded = self.tokenizer.decode(tokens, skip_special_tokens=True)
-        assert isinstance(decoded, str)  # one sequence in, one string out
+        decoded = self._tokenizer().decode(tokens, skip_special_tokens=True)
+        if not isinstance(decoded, str):  # one sequence in, one string out
+            msg = f"tokenizer.decode returned {type(decoded).__name__}, expected str"
+            raise TypeError(msg)
         text = decoded.strip()
         # Strip <think>...</think> tags (Qwen3 doesn't respect /no_think prompt)
         if "<think>" in text:
             text = _THINK_TAG_RE.sub("", text).strip()
-        text = _truncate_repetitions(text)
+        text = truncate_repetitions(text)
         out: dict[str, Any] = {"text": text}
         # Pass through per-step logprobs when _forward captured them (i.e. caller
         # passed output_scores=True). Lets eval harnesses compute confidence
         # stats without re-running the model.
-        if "top1_logprob" in model_outputs:
-            out["top1_logprob"] = model_outputs["top1_logprob"]
-        if "top2_logprob" in model_outputs:
-            out["top2_logprob"] = model_outputs["top2_logprob"]
+        for key in ("top1_logprob", "top2_logprob"):
+            if key in model_outputs:
+                out[key] = model_outputs[key]
         return out
 
 
@@ -681,71 +709,3 @@ def _single_input(args: tuple[Any, ...], kwargs: dict[str, Any], name: str) -> A
         return kwargs.pop(name)
     msg = f"expected exactly one {name!r} argument, got {len(args)} positional"
     raise TypeError(msg)
-
-
-def _truncate_repetitions(text: str) -> str:
-    """Truncate repeated words/phrases/characters at end of text.
-
-    Detects patterns like:
-    - Repeated words: "the the the the" -> "the"
-    - Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
-    - Repeated characters: "444444" -> "4"
-
-    Words match ignoring case. A run of `_LOOSE_MIN_REPEATS` also matches
-    ignoring punctuation, since the model punctuates inside a loop ("Yeah Yeah
-    Yeah. Yeah. Yeah,"); shorter punctuated runs are real speech ("No, no,
-    no,"), and collapsing those cost 40 errors over 47 rows of the served eval.
-
-    Args:
-        text: Input text to process
-
-    Returns:
-        Text with trailing repetitions removed
-    """
-    if not text:
-        return text
-
-    text = _TRAILING_CHAR_RE.sub(r"\1", text)
-    while _TRAILING_WORD_RE.search(text):
-        text = _TRAILING_WORD_RE.sub(r"\1", text)
-
-    # 3. Truncate repeated phrases (2-20 words) at end
-    # e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
-    words = text.split()
-    if len(words) < _MIN_REPEATS * 2:
-        return text
-
-    # Cheap pre-check: trailing window must contain duplicates for any phrase repeat
-    # to be possible. set(window) == window means all unique → no repetition.
-    window = words[-_MIN_REPEATS * 2 :]
-    if len(set(window)) == len(window):
-        return _truncate_punctuated_loop(text)
-
-    for phrase_len in range(2, min(21, len(words) // _MIN_REPEATS + 1)):
-        phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
-        phrase_pattern = re.compile(
-            rf"(^|.*?\s)({phrase_escaped})(?:\s+{phrase_escaped}){{{_MIN_REPEATS - 1},}}\s*$",
-            re.IGNORECASE,
-        )
-        match = phrase_pattern.match(text)
-        if match:
-            text = (match.group(1) + match.group(2)).strip()
-            break
-
-    return _truncate_punctuated_loop(text)
-
-
-def _truncate_punctuated_loop(text: str) -> str:
-    """Collapse a trailing 1-20 word phrase run `_LOOSE_MIN_REPEATS`+ times, punctuation ignored."""
-    words = text.split()
-    keys = [_NON_WORD_RE.sub("", w).lower() for w in words]
-    end = len(words)
-    # Shortest phrase first, so "Yeah Yeah. Yeah," stops at "Yeah".
-    for n in range(1, min(21, end // _LOOSE_MIN_REPEATS + 1)):
-        phrase = keys[end - n : end]
-        reps = 1
-        while end - (reps + 1) * n >= 0 and keys[end - (reps + 1) * n : end - reps * n] == phrase:
-            reps += 1
-        if reps >= _LOOSE_MIN_REPEATS:
-            return " ".join(words[: end - (reps - 1) * n])
-    return text

asr_processing.py

@@ -1,12 +1,18 @@
 """Processor that turns raw audio (and optional text) into model inputs."""
 
+import importlib
+import importlib.util
+import io
 from collections.abc import Mapping, Sequence
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, cast, overload
 
 import numpy as np
 import numpy.typing as npt
 import torch
+import torchaudio
 import transformers
+from huggingface_hub.utils import httpx
 from torch.nn.utils.rnn import pad_sequence
 from transformers import (
     BatchFeature,
@@ -14,6 +20,13 @@ from transformers import (
     ProcessorMixin,
     SequenceFeatureExtractor,
 )
+from transformers.pipelines.audio_utils import ffmpeg_read
+
+# soundfile is optional: it arrives with librosa, which the README installs, but
+# it is not in the model's requirements.txt, and transformers refuses to load a
+# remote-code file with a plain `import` of a missing package. Without it every
+# input goes to ffmpeg.
+_soundfile = importlib.import_module("soundfile") if importlib.util.find_spec("soundfile") else None
 
 if TYPE_CHECKING:
     from .asr_config import (
@@ -250,6 +263,35 @@ def is_silent(audio: npt.NDArray[np.float32]) -> bool:
     return audio.size == 0 or float(np.sqrt(np.mean(np.square(audio)))) < SILENCE_RMS
 
 
+def decode_audio_input(inputs: str | bytes) -> npt.NDArray[np.float32]:
+    """A URL, file path or encoded bytes to mono float32 at 16 kHz.
+
+    soundfile first: its bundled libsndfile reads WAV, FLAC, OGG and MP3 with
+    nothing else installed, so a fresh `pip install` can transcribe a file.
+    ffmpeg (a system binary the user may not have) only for what it can't
+    read, such as m4a or a video's audio track, or when soundfile is missing.
+    URLs are fetched the way transformers' own ASR pipeline fetches them.
+    """
+    if isinstance(inputs, bytes):
+        data = inputs
+    elif inputs.startswith(("http://", "https://")):
+        response = httpx.get(inputs, follow_redirects=True)
+        response.raise_for_status()
+        data = response.content
+    else:
+        data = Path(inputs).read_bytes()
+    try:
+        if _soundfile is None:
+            return ffmpeg_read(data, sampling_rate=16000)
+        array, sr = _soundfile.read(io.BytesIO(data), dtype="float32", always_2d=True)
+    except RuntimeError:  # LibsndfileError: a format libsndfile can't read
+        return ffmpeg_read(data, sampling_rate=16000)
+    mono = torch.from_numpy(array.mean(axis=1))
+    if sr != 16000:
+        mono = torchaudio.functional.resample(mono, orig_freq=sr, new_freq=16000)
+    return mono.numpy()
+
+
 class ASRProcessor(ProcessorMixin):
     """Processor for Whisper-based ASR models."""
 
@@ -258,12 +300,10 @@ class ASRProcessor(ProcessorMixin):
     tokenizer: PreTrainedTokenizerBase
     feature_extractor_class = "AutoFeatureExtractor"
     tokenizer_class = "AutoTokenizer"
-    # Fallback only. The real value comes from `ASRConfig.audio_token`, which
-    # resolves to the decoder's native placeholder where it has one (Gemma 4's
-    # pretrained "<|audio|>") and to "<audio>" otherwise. Hardcoding the
-    # fallback here fails silently on a native-token decoder: "<audio>" was
-    # never added to that vocab, so it tokenizes into ordinary subwords and
-    # the prompt ends up with zero scatter positions for N audio embeddings.
+    # Fallback only. The real value comes from `ASRConfig.audio_token`; a
+    # config naming a different token must win, or "<audio>" -- never added to
+    # that vocab -- tokenizes into ordinary subwords and the prompt ends up
+    # with zero scatter positions for N audio embeddings.
     AUDIO_TOKEN = "<audio>"
     TRANSCRIBE_PROMPT = DEFAULT_TRANSCRIBE_PROMPT
 
@@ -275,6 +315,7 @@ class ASRProcessor(ProcessorMixin):
         encoder_conv_layers: list[ConvLayerSpec] | None = None,
         audio_token: str | None = None,
         lead_in_seconds: float = 0.0,
+        encoder_chunk_frames: int = 0,
     ):
         """Initialize the ASR processor.
 
@@ -286,14 +327,23 @@ class ASRProcessor(ProcessorMixin):
             audio_token: Placeholder token scattered with audio embeddings.
                 Must match `ASRConfig.audio_token` / `ASRModel.audio_token`;
                 defaults to AUDIO_TOKEN.
+            lead_in_seconds: Silence prepended to every clip.
+            encoder_chunk_frames: `ASRConfig.encoder_chunk_frames` (0: no chunking).
         """
         self.feature_extractor = feature_extractor
         self.tokenizer = tokenizer
-        self.audio_token = audio_token or self.AUDIO_TOKEN
-        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
-        self.projector = projector
-        self.encoder_conv_layers = encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
-        self.lead_in_seconds = float(lead_in_seconds)
+        self.audio_token: str = audio_token or self.AUDIO_TOKEN
+        token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
+        if not isinstance(token_id, int):  # a str in always maps to one id
+            msg = f"audio token {self.audio_token!r} mapped to {token_id!r}, not one id"
+            raise TypeError(msg)
+        self.audio_token_id: int = token_id
+        self.projector: MLPAudioProjector | None = projector
+        self.encoder_conv_layers: list[ConvLayerSpec] = (
+            encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
+        )
+        self.lead_in_seconds: float = float(lead_in_seconds)
+        self.encoder_chunk_frames: int = encoder_chunk_frames
 
     def _render_prompt(self, num_audio_tokens: int, text: str | None) -> torch.Tensor:
         """Tokenize one chat prompt carrying exactly `num_audio_tokens` placeholders."""
@@ -301,10 +351,6 @@ class ASRProcessor(ProcessorMixin):
             self.tokenizer, self.audio_token, num_audio_tokens, self.TRANSCRIBE_PROMPT, text
         )
 
-    def _stack_prompt_rows(self, rows: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
-        """Stack per-sample prompt rows into a batch (see `left_pad_prompt_rows`)."""
-        return left_pad_prompt_rows(rows, self.tokenizer)
-
     def __call__(self, *args: Any, **kwargs: Any) -> BatchFeature:
         """Process audio and text inputs for inference; see `_process` for the arguments.
 
@@ -361,15 +407,17 @@ class ASRProcessor(ProcessorMixin):
             # returns batch-1 `input_ids` against batch-B `input_features`, and
             # gives every shorter row more `<audio>` placeholders than the
             # projector produced for it. `masked_scatter` then mis-scatters
-            # silently. This is the same failure `_prepare_audio_inputs`
+            # silently. This is the same failure `prepare_audio_inputs`
             # documents as fixed on the model side, and it only shows up on a
             # ragged batch, so batch-1 eval never sees it.
             mel_lengths = audio_inputs["attention_mask"].sum(dim=-1).reshape(-1).long()
-            encoder_lengths = compute_encoder_output_length(mel_lengths, self.encoder_conv_layers)
+            encoder_lengths = compute_encoder_output_length(
+                mel_lengths, self.encoder_conv_layers, self.encoder_chunk_frames
+            )
             token_counts = self.projector.get_output_length(encoder_lengths).tolist()
 
         rows = [self._render_prompt(n, text) for n in token_counts]
-        input_ids, attention_mask = self._stack_prompt_rows(rows)
+        input_ids, attention_mask = left_pad_prompt_rows(rows, self.tokenizer)
         result["input_ids"] = input_ids
         result["attention_mask"] = attention_mask
 

asr_text.py

@@ -0,0 +1,123 @@
+"""Transcript text cleanup: collapse the repetition loops a decoder can fall into.
+
+A language-model decoder that loses the audio sometimes repeats its last
+character, word or phrase until it hits the token limit. These functions trim
+such a trailing run back to one copy; ASRPipeline applies them to every output.
+`truncate_runaway_loop` also catches a loop that starts mid-transcript and
+drifts ("It's It's It it's It's"), so it never forms an exact trailing run.
+"""
+
+import re
+
+MIN_REPEATS = 3
+LOOSE_MIN_REPEATS = 6
+TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{MIN_REPEATS - 1},}}$")
+TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
+NON_WORD_RE = re.compile(r"\W+")
+# A 1-4 word unit repeated this many times in a row, after which at most this
+# share of the words are distinct: the decoder lost the audio and never came
+# back. Measured on the Open ASR leaderboard runs of 2026-10-16/17: URGENT
+# -0.97 WER (13.12 -> 12.11 on an independent rerun), and no set worse in
+# either run. Cutting at every 6x run regardless of what follows cost up to
+# +0.05 (Monsoon, GigaSpeech): "No, no, no, no, no, no, no, sir, which..." is
+# real speech that recovers.
+RUNAWAY_MIN_REPEATS = 6
+RUNAWAY_MAX_UNIT = 4
+RUNAWAY_MAX_DISTINCT = 0.5
+
+
+def truncate_repetitions(text: str) -> str:
+    """Truncate repeated words/phrases/characters at end of text.
+
+    Detects patterns like:
+    - Repeated words: "the the the the" -> "the"
+    - Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
+    - Repeated characters: "444444" -> "4"
+
+    Words match ignoring case. A run of `LOOSE_MIN_REPEATS` also matches
+    ignoring punctuation, since the model punctuates inside a loop ("Yeah Yeah
+    Yeah. Yeah. Yeah,"); shorter punctuated runs are real speech ("No, no,
+    no,"), and collapsing those cost 40 errors over 47 rows of the served eval.
+
+    Args:
+        text: Input text to process
+
+    Returns:
+        Text with trailing repetitions removed
+    """
+    if not text:
+        return text
+
+    text = truncate_runaway_loop(text)
+    text = TRAILING_CHAR_RE.sub(r"\1", text)
+    while TRAILING_WORD_RE.search(text):
+        text = TRAILING_WORD_RE.sub(r"\1", text)
+
+    # 3. Truncate repeated phrases (2-20 words) at end
+    # e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
+    words = text.split()
+    if len(words) < MIN_REPEATS * 2:
+        return text
+
+    # Cheap pre-check: trailing window must contain duplicates for any phrase repeat
+    # to be possible. set(window) == window means all unique → no repetition.
+    window = words[-MIN_REPEATS * 2 :]
+    if len(set(window)) == len(window):
+        return truncate_punctuated_loop(text)
+
+    for phrase_len in range(2, min(21, len(words) // MIN_REPEATS + 1)):
+        phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
+        phrase_pattern = re.compile(
+            rf"(^|.*?\s)({phrase_escaped})(?:\s+{phrase_escaped}){{{MIN_REPEATS - 1},}}\s*$",
+            re.IGNORECASE,
+        )
+        match = phrase_pattern.match(text)
+        if match:
+            text = (match.group(1) + match.group(2)).strip()
+            break
+
+    return truncate_punctuated_loop(text)
+
+
+def truncate_punctuated_loop(text: str) -> str:
+    """Collapse a trailing 1-20 word phrase run `LOOSE_MIN_REPEATS`+ times, punctuation ignored."""
+    words = text.split()
+    keys = [NON_WORD_RE.sub("", w).lower() for w in words]
+    end = len(words)
+    # Shortest phrase first, so "Yeah Yeah. Yeah," stops at "Yeah".
+    for n in range(1, min(21, end // LOOSE_MIN_REPEATS + 1)):
+        phrase = keys[end - n : end]
+        reps = 1
+        while end - (reps + 1) * n >= 0 and keys[end - (reps + 1) * n : end - reps * n] == phrase:
+            reps += 1
+        if reps >= LOOSE_MIN_REPEATS:
+            return " ".join(words[: end - (reps - 1) * n])
+    return text
+
+
+def truncate_runaway_loop(text: str) -> str:
+    """Cut at a loop the decoder never leaves, keeping the first copy of its unit.
+
+    Finds the first 1-`RUNAWAY_MAX_UNIT` word unit repeated
+    `RUNAWAY_MIN_REPEATS`+ times in a row (case and punctuation ignored). It
+    cuts there only when the rest of the transcript, from that run on, has at
+    most `RUNAWAY_MAX_DISTINCT` distinct words per word: a loop that drifts
+    ("It's It's It it's It's ...") rather than speech that moves on.
+    """
+    words = text.split()
+    keys = [NON_WORD_RE.sub("", w).lower() for w in words]
+    for start in range(len(keys)):
+        for size in range(1, RUNAWAY_MAX_UNIT + 1):
+            unit = keys[start : start + size]
+            if len(unit) < size or not all(unit):
+                break
+            reps = 1
+            while keys[start + reps * size : start + (reps + 1) * size] == unit:
+                reps += 1
+            if reps < RUNAWAY_MIN_REPEATS:
+                continue
+            rest = keys[start:]
+            if len(set(rest)) / len(rest) > RUNAWAY_MAX_DISTINCT:
+                return text  # the transcript recovers after the first run
+            return " ".join(words[: start + size])
+    return text

chat_template.jinja

@@ -1,90 +1,34 @@
-{%- set image_count = namespace(value=0) %}
-{%- set video_count = namespace(value=0) %}
-{%- macro render_content(content, do_vision_count, is_system_content=false) %}
-    {%- if content is string %}
-        {{- content }}
-    {%- elif content is iterable and content is not mapping %}
-        {%- for item in content %}
-            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
-                {%- if is_system_content %}
-                    {{- raise_exception('System message cannot contain images.') }}
-                {%- endif %}
-                {%- if do_vision_count %}
-                    {%- set image_count.value = image_count.value + 1 %}
-                {%- endif %}
-                {%- if add_vision_id %}
-                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
-                {%- endif %}
-                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
-            {%- elif 'video' in item or item.type == 'video' %}
-                {%- if is_system_content %}
-                    {{- raise_exception('System message cannot contain videos.') }}
-                {%- endif %}
-                {%- if do_vision_count %}
-                    {%- set video_count.value = video_count.value + 1 %}
-                {%- endif %}
-                {%- if add_vision_id %}
-                    {{- 'Video ' ~ video_count.value ~ ': ' }}
-                {%- endif %}
-                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
-            {%- elif 'text' in item %}
-                {{- item.text }}
-            {%- else %}
-                {{- raise_exception('Unexpected item type in content.') }}
-            {%- endif %}
-        {%- endfor %}
-    {%- elif content is none or content is undefined %}
-        {{- '' }}
-    {%- else %}
-        {{- raise_exception('Unexpected content type.') }}
-    {%- endif %}
-{%- endmacro %}
-{%- if not messages %}
-    {{- raise_exception('No messages provided.') }}
-{%- endif %}
-{%- if tools and tools is iterable and tools is not mapping %}
+{%- if tools %}
     {{- '<|im_start|>system\n' }}
-    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
     {%- for tool in tools %}
         {{- "\n" }}
         {{- tool | tojson }}
     {%- endfor %}
-    {{- "\n</tools>" }}
-    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
-    {%- if messages[0].role == 'system' %}
-        {%- set content = render_content(messages[0].content, false, true)|trim %}
-        {%- if content %}
-            {{- '\n\n' + content }}
-        {%- endif %}
-    {%- endif %}
-    {{- '<|im_end|>\n' }}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
 {%- else %}
     {%- if messages[0].role == 'system' %}
-        {%- set content = render_content(messages[0].content, false, true)|trim %}
-        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
     {%- endif %}
 {%- endif %}
 {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
 {%- for message in messages[::-1] %}
     {%- set index = (messages|length - 1) - loop.index0 %}
-    {%- if ns.multi_step_tool and message.role == "user" %}
-        {%- set content = render_content(message.content, false)|trim %}
-        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
-            {%- set ns.multi_step_tool = false %}
-            {%- set ns.last_query_index = index %}
-        {%- endif %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
     {%- endif %}
 {%- endfor %}
-{%- if ns.multi_step_tool %}
-    {{- raise_exception('No user query found in messages.') }}
-{%- endif %}
 {%- for message in messages %}
-    {%- set content = render_content(message.content, true)|trim %}
-    {%- if message.role == "system" %}
-        {%- if not loop.first %}
-            {{- raise_exception('System message must be at the beginning.') }}
-        {%- endif %}
-    {%- elif message.role == "user" %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
         {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
     {%- elif message.role == "assistant" %}
         {%- set reasoning_content = '' %}
@@ -96,59 +40,50 @@
                 {%- set content = content.split('</think>')[-1].lstrip('\n') %}
             {%- endif %}
         {%- endif %}
-        {%- set reasoning_content = reasoning_content|trim %}
         {%- if loop.index0 > ns.last_query_index %}
-            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
         {%- else %}
             {{- '<|im_start|>' + message.role + '\n' + content }}
         {%- endif %}
-        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+        {%- if message.tool_calls %}
             {%- for tool_call in message.tool_calls %}
-                {%- if tool_call.function is defined %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
                     {%- set tool_call = tool_call.function %}
                 {%- endif %}
-                {%- if loop.first %}
-                    {%- if content|trim %}
-                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
-                    {%- else %}
-                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
-                    {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
                 {%- else %}
-                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
-                {%- endif %}
-                {%- if tool_call.arguments is defined %}
-                    {%- for args_name, args_value in tool_call.arguments|items %}
-                        {{- '<parameter=' + args_name + '>\n' }}
-                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
-                        {{- args_value }}
-                        {{- '\n</parameter>\n' }}
-                    {%- endfor %}
+                    {{- tool_call.arguments | tojson }}
                 {%- endif %}
-                {{- '</function>\n</tool_call>' }}
+                {{- '}\n</tool_call>' }}
             {%- endfor %}
         {%- endif %}
         {{- '<|im_end|>\n' }}
     {%- elif message.role == "tool" %}
-        {%- if loop.previtem and loop.previtem.role != "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
             {{- '<|im_start|>user' }}
         {%- endif %}
         {{- '\n<tool_response>\n' }}
         {{- content }}
         {{- '\n</tool_response>' }}
-        {%- if not loop.last and loop.nextitem.role != "tool" %}
-            {{- '<|im_end|>\n' }}
-        {%- elif loop.last %}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
             {{- '<|im_end|>\n' }}
         {%- endif %}
-    {%- else %}
-        {{- raise_exception('Unexpected message role.') }}
     {%- endif %}
 {%- endfor %}
 {%- if add_generation_prompt %}
     {{- '<|im_start|>assistant\n' }}
-    {%- if enable_thinking is defined and enable_thinking is true %}
-        {{- '<think>\n' }}
-    {%- else %}
+    {%- if true %}
         {{- '<think>\n\n</think>\n\n' }}
     {%- endif %}
 {%- endif %}

config.json

@@ -1,17 +1,16 @@
 {
-  "apply_spec_augment": false,
   "architectures": [
     "ASRModel"
   ],
-  "attn_implementation": "sdpa",
+  "attn_implementation": "flash_attention_2",
   "audio_config": {
-    "_name_or_path": "ibm-granite/granite-speech-5.0-470m-turboctc",
+    "_name_or_path": "ibm-granite/granite-speech-5.0-470m-turboctc-nc",
     "architectures": [
       "GraniteSpeech5ForCTC"
     ],
     "ctc_loss_reduction": "mean",
     "ctc_zero_infinity": true,
-    "dtype": "bfloat16",
+    "dtype": "float32",
     "encoder_config": {
       "_name_or_path": "",
       "activation_dropout": 0.0,
@@ -60,14 +59,13 @@
     "tie_word_embeddings": true,
     "vocab_size": 16384
   },
-  "audio_features_time_major": true,
-  "audio_model_id": "ibm-granite/granite-speech-5.0-470m-turboctc",
+  "audio_model_id": "ibm-granite/granite-speech-5.0-470m-turboctc-nc",
   "audio_sample_rate": 16000,
   "audio_token": "<audio>",
   "auto_map": {
     "AutoConfig": "asr_config.ASRConfig",
-    "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
     "AutoModel": "asr_modeling.ASRModel",
+    "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
     "AutoProcessor": "asr_processing.ASRProcessor"
   },
   "custom_pipelines": {
@@ -82,13 +80,13 @@
   },
   "dtype": "bfloat16",
   "encoder": {
-    "_name_or_path": "ibm-granite/granite-speech-5.0-470m-turboctc",
+    "_name_or_path": "ibm-granite/granite-speech-5.0-470m-turboctc-nc",
     "architectures": [
       "GraniteSpeech5ForCTC"
     ],
     "ctc_loss_reduction": "mean",
     "ctc_zero_infinity": true,
-    "dtype": "bfloat16",
+    "dtype": "float32",
     "encoder_config": {
       "_name_or_path": "",
       "activation_dropout": 0.0,
@@ -138,6 +136,7 @@
     "vocab_size": 16384
   },
   "encoder_attention_mask": true,
+  "encoder_chunk_frames": 0,
   "encoder_conv_layers": [
     [
       0,
@@ -154,179 +153,94 @@
   "encoder_dtype": "bfloat16",
   "encoder_trainable_post_projections": false,
   "encoder_trainable_top_layers": 0,
-  "eos_token_id": 248046,
   "freeze_audio_encoder": true,
-  "freeze_language_model": true,
+  "freeze_language_model": false,
   "freeze_projector": false,
-  "freeze_text_embed_tokens": true,
+  "freeze_text_embed_tokens": false,
+  "inference_chunk_seconds": 28.0,
   "inference_lead_in_seconds": 0.25,
-  "label_smoothing": 0.0,
-  "llm_dim": 2048,
+  "llm_dim": 1024,
   "lora_alpha": 64,
-  "lora_alpha_pattern": {},
   "lora_dropout": 0.0,
   "lora_rank": 64,
-  "lora_rank_pattern": {},
   "lora_target_modules": "all-linear",
-  "mask_time_length": 10,
-  "mask_time_min_masks": 2,
-  "mask_time_prob": 0.05,
   "max_new_tokens": 256,
   "model_dtype": "bfloat16",
   "model_type": "asr_model",
   "no_repeat_ngram_size": 12,
-  "pad_token_id": 248044,
+  "pack_sequences": true,
   "pipeline_tag": "automatic-speech-recognition",
   "pretrained_model_path": "mazesmazes/tiny-audio",
   "projector_dtype": "float32",
   "projector_hidden_dim": 4096,
   "projector_pool_stride": 1,
-  "projector_type": "mlp",
   "text_config": {
-    "_name_or_path": "Qwen/Qwen3.5-2B",
+    "_name_or_path": "Qwen/Qwen3-0.6B",
     "architectures": [
-      "Qwen3_5ForConditionalGeneration"
+      "Qwen3ForCausalLM"
     ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
     "bos_token_id": null,
-    "dtype": "bfloat16",
-    "eos_token_id": 248046,
-    "image_token_id": 248056,
-    "model_type": "qwen3_5",
-    "pad_token_id": 248044,
-    "text_config": {
-      "_name_or_path": "",
-      "architectures": null,
-      "attention_bias": false,
-      "attention_dropout": 0.0,
-      "attn_output_gate": true,
-      "bos_token_id": null,
-      "chunk_size_feed_forward": 0,
-      "dtype": "bfloat16",
-      "eos_token_id": 248044,
-      "full_attention_interval": 4,
-      "head_dim": 256,
-      "hidden_act": "silu",
-      "hidden_size": 2048,
-      "id2label": {
-        "0": "LABEL_0",
-        "1": "LABEL_1"
-      },
-      "initializer_range": 0.02,
-      "intermediate_size": 6144,
-      "is_encoder_decoder": false,
-      "label2id": {
-        "LABEL_0": 0,
-        "LABEL_1": 1
-      },
-      "layer_types": [
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention",
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention",
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention",
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention",
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention",
-        "linear_attention",
-        "linear_attention",
-        "linear_attention",
-        "full_attention"
-      ],
-      "linear_conv_kernel_dim": 4,
-      "linear_key_head_dim": 128,
-      "linear_num_key_heads": 16,
-      "linear_num_value_heads": 16,
-      "linear_value_head_dim": 128,
-      "mamba_ssm_dtype": "float32",
-      "max_position_embeddings": 262144,
-      "mlp_only_layers": [],
-      "model_type": "qwen3_5_text",
-      "mtp_num_hidden_layers": 1,
-      "mtp_use_dedicated_embeddings": false,
-      "num_attention_heads": 8,
-      "num_hidden_layers": 24,
-      "num_key_value_heads": 2,
-      "output_attentions": false,
-      "output_hidden_states": false,
-      "pad_token_id": null,
-      "partial_rotary_factor": 0.25,
-      "problem_type": null,
-      "return_dict": true,
-      "rms_norm_eps": 1e-06,
-      "rope_parameters": {
-        "mrope_interleaved": true,
-        "mrope_section": [
-          11,
-          11,
-          10
-        ],
-        "partial_rotary_factor": 0.25,
-        "rope_theta": 10000000,
-        "rope_type": "default"
-      },
-      "tie_word_embeddings": true,
-      "use_cache": true,
-      "vocab_size": 248320
+    "dtype": "float32",
+    "eos_token_id": 151645,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 40960,
+    "max_window_layers": 28,
+    "model_type": "qwen3",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 8,
+    "pad_token_id": 151643,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 1000000,
+      "rope_type": "default"
     },
+    "sliding_window": null,
     "tie_word_embeddings": true,
-    "video_token_id": 248057,
-    "vision_config": {
-      "_name_or_path": "",
-      "architectures": null,
-      "chunk_size_feed_forward": 0,
-      "deepstack_visual_indexes": [],
-      "depth": 24,
-      "dtype": null,
-      "hidden_act": "gelu_pytorch_tanh",
-      "hidden_size": 1024,
-      "id2label": {
-        "0": "LABEL_0",
-        "1": "LABEL_1"
-      },
-      "in_channels": 3,
-      "initializer_range": 0.02,
-      "intermediate_size": 4096,
-      "is_encoder_decoder": false,
-      "label2id": {
-        "LABEL_0": 0,
-        "LABEL_1": 1
-      },
-      "model_type": "qwen3_5_vision",
-      "num_heads": 16,
-      "num_position_embeddings": 2304,
-      "out_hidden_size": 2048,
-      "output_attentions": false,
-      "output_hidden_states": false,
-      "patch_size": 16,
-      "problem_type": null,
-      "return_dict": true,
-      "rope_parameters": {
-        "rope_theta": 10000.0,
-        "rope_type": "axial"
-      },
-      "spatial_merge_size": 2,
-      "temporal_patch_size": 2
-    },
-    "vision_end_token_id": 248054,
-    "vision_start_token_id": 248053,
-    "vocab_size": 248192
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151680
   },
-  "text_model_id": "Qwen/Qwen3.5-2B",
+  "text_model_id": "Qwen/Qwen3-0.6B",
   "transcribe_prompt": "Transcribe the speech with proper punctuation and capitalization",
-  "transformers_version": "5.17.0",
+  "transformers_version": "5.19.0.dev0",
   "use_cache": true,
-  "use_lora": true,
-  "vocab_size": 248192
+  "use_lora": false,
+  "vocab_size": 151680
 }

diarization.py

@@ -2,7 +2,7 @@
 
 import importlib.util
 import re
-from typing import TYPE_CHECKING, Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -161,7 +161,7 @@ class NemotronDiarizer:
                 AutoModelForAudioFrameClassification.from_pretrained(cls.MODEL_ID)
             )
             # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a method.
-            model.to(get_device())  # pyright: ignore[reportArgumentType]
+            cast(torch.nn.Module, model).to(get_device())
             model.eval()
             cls._model = model
             processor: Nemotron3DiarizationProcessor = AutoProcessor.from_pretrained(cls.MODEL_ID)

generation_config.json

@@ -1,13 +1,16 @@
 {
-  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "do_sample": true,
   "eos_token_id": [
-    248046,
-    248046,
-    248044
+    151645,
+    151643
   ],
   "max_new_tokens": 256,
   "no_repeat_ngram_size": 12,
-  "pad_token_id": 248044,
-  "transformers_version": "5.17.0",
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "5.19.0.dev0",
   "use_cache": true
 }

handler.py

@@ -1,112 +0,0 @@
-"""Custom inference handler for HuggingFace Inference Endpoints."""
-
-import base64
-import binascii
-import os
-from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from .alignment import QwenForcedAligner
-    from .asr_modeling import ASRModel
-    from .asr_pipeline import ASRPipeline
-    from .diarization import NemotronDiarizer
-    from .diarization import get_device as _best_device
-else:
-    try:
-        # For remote execution, imports are relative
-        from .alignment import QwenForcedAligner
-        from .asr_modeling import ASRModel
-        from .asr_pipeline import ASRPipeline
-        from .diarization import NemotronDiarizer
-        from .diarization import get_device as _best_device
-    except ImportError:
-        # For local execution, imports are not relative
-        from alignment import QwenForcedAligner
-        from asr_modeling import ASRModel
-        from asr_pipeline import ASRPipeline
-        from diarization import NemotronDiarizer
-        from diarization import get_device as _best_device
-
-
-def decode_inputs(inputs: Any) -> Any:
-    """A JSON request's base64 audio as bytes; paths, raw bytes and arrays unchanged.
-
-    The toolkit hands `inputs` over as raw bytes for an `audio/*` body, but a
-    JSON body can only carry audio as a base64 string, which the pipeline
-    would otherwise read as a file path.
-    """
-    if not isinstance(inputs, str) or Path(inputs).exists():
-        return inputs
-    try:
-        return base64.b64decode(inputs, validate=True)
-    except (binascii.Error, ValueError):
-        return inputs
-
-
-class EndpointHandler:
-    """HuggingFace Inference Endpoints handler for ASR model.
-
-    Handles model loading, warmup, and inference requests for deployment
-    on HuggingFace Inference Endpoints or similar services.
-    """
-
-    def __init__(self, path: str = ""):
-        """Initialize the endpoint handler.
-
-        Args:
-            path: Path to model directory or HuggingFace model ID
-        """
-        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-
-        # `ASRModel.from_pretrained` constructs its own submodules and forwards
-        # **kwargs to `ASRModel.__init__`, which discards them -- no loader
-        # reads `device_map`, `torch_dtype` or `low_cpu_mem_usage`, and it sets
-        # `_is_loading_from_pretrained` precisely to keep `device_map="auto"`
-        # out of the sub-model loaders. Passing them here did nothing, so the
-        # model stayed on CPU and a GPU Inference Endpoint silently decoded on
-        # CPU. Place it explicitly instead. dtype comes from
-        # `config.model_dtype` and the attention backend from
-        # `config.attn_implementation`, which already downgrades FA2 when
-        # flash_attn is missing.
-        self.model = ASRModel.from_pretrained(path)
-        self.device = _best_device()
-        # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a method.
-        self.model.to(self.device)  # pyright: ignore[reportArgumentType]
-        self.model.eval()
-
-        self.pipe = ASRPipeline(
-            model=self.model,
-            feature_extractor=self.model.feature_extractor,
-            tokenizer=self.model.tokenizer,
-            device=self.device,
-        )
-        # Load the forced aligner (return_timestamps) and Nemotron diarizer
-        # (return_speakers) at boot, as the Spaces demo does, so the first
-        # request that asks for them doesn't pay the download and load.
-        QwenForcedAligner.get_instance()
-        NemotronDiarizer.get_instance()
-
-    def __call__(self, data: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
-        """Process an inference request.
-
-        Args:
-            data: Request data containing 'inputs' (audio path, bytes or base64
-                string) and optional 'parameters' for the pipeline, e.g.
-                `return_timestamps`, `return_speakers`, `num_speakers`,
-                `max_speakers` -- the same options as the Spaces demo
-
-        Returns:
-            Transcription result with 'text' key; 'words' with timestamps, plus
-            speaker labels and 'speaker_segments' when diarizing
-        """
-        inputs = data.get("inputs")
-        if inputs is None:
-            msg = "Missing 'inputs' in request data"
-            raise ValueError(msg)
-
-        # Pass through any parameters from request, let model config provide defaults
-        params = data.get("parameters", {})
-
-        result: dict[str, Any] | list[dict[str, Any]] = self.pipe(decode_inputs(inputs), **params)
-        return result

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2aae926199ca8125d15aa878e4c51cc69086ce150e9bb1c5dcb648f90be97cee
-size 50360824
+oid sha256:4b17a2bb72d661cd79a5d7cb2772b1d817a667816f81c933dc73c6c126915559
+size 2416770552

preprocessor_config.json

@@ -6,7 +6,6 @@
   "logmel_floor_db": 8.0,
   "n_fft": 512,
   "num_mel_bins": 80,
-  "padding": false,
   "padding_side": "right",
   "padding_value": 0.0,
   "return_attention_mask": true,

projectors.py

@@ -7,7 +7,7 @@ This module contains all projector architectures:
 from typing import overload
 
 import torch
-import torch.nn as nn
+from torch import nn
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
 # =============================================================================
@@ -108,10 +108,10 @@ class MLPAudioProjector(nn.Module):
         in_dim = encoder_dim * self.k
         # Hidden dim defaults to llm_dim, can be overridden via config
         hidden_dim: int = getattr(config, "projector_hidden_dim", None) or llm_dim
-        self.input_norm = LlamaRMSNorm(in_dim, eps=1e-6)
-        self.linear_1 = nn.Linear(in_dim, hidden_dim)
-        self.act = nn.GELU()
-        self.linear_2 = nn.Linear(hidden_dim, llm_dim)
+        self.input_norm: LlamaRMSNorm = LlamaRMSNorm(in_dim, eps=1e-6)
+        self.linear_1: nn.Linear = nn.Linear(in_dim, hidden_dim)
+        self.act: nn.GELU = nn.GELU()
+        self.linear_2: nn.Linear = nn.Linear(hidden_dim, llm_dim)
 
     @overload
     def get_output_length(self, input_length: int) -> int: ...
@@ -160,12 +160,3 @@ def _frame_stack(x: torch.Tensor, k: int) -> torch.Tensor:
     batch, seq, dim = x.shape
     out_len = _frame_stack_length(seq, k)
     return x[:, : out_len * k, :].reshape(batch, out_len, dim * k)
-
-
-# =============================================================================
-# Projector Registry
-# =============================================================================
-
-PROJECTOR_CLASSES = {
-    "mlp": MLPAudioProjector,
-}

qwen3_asr_encoder.py

@@ -0,0 +1,152 @@
+"""Qwen3-ASR's audio encoder behind ASRModel's encoder interface.
+
+`Qwen/Qwen3-ASR-1.7B-hf` (Apache-2.0) is #9 on the Open ASR leaderboard (mean
+of 9 4.936, 2026-10-20), the best system whose encoder loads natively in our
+transformers, and its encoder was trained to feed a Qwen3 decoder.
+
+transformers' `Qwen3ASREncoder` owns its preprocessing contract, which is why
+it is used as-is rather than reimplemented:
+
+- its convolutions run on independent 1 s chunks (100 mel frames, 13 tokens
+  each) and attention is limited to 8 s windows, both inside `forward`;
+- it takes the feature extractor's mel mask and drops padding itself.
+
+Two things differ from what ASRModel expects, and this module bridges them:
+
+- the encoder returns the valid frames of the whole batch packed into one
+  `(sum(tokens), d_model)` tensor, which is split back into a padded
+  `(batch, tokens, d_model)` tensor here;
+- it lives inside the full ASR checkpoint, so only its `model.audio_tower.*`
+  tensors are read.
+"""
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import cast
+
+import torch
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import EntryNotFoundError
+from safetensors import safe_open
+from torch import nn
+from transformers.models.qwen3_asr import Qwen3ASRConfig, Qwen3ASREncoder, Qwen3ASREncoderConfig
+from transformers.utils.generic import ModelOutput
+
+try:
+    from .asr_config import (
+        QWEN3_ASR_CHUNK_FRAMES,
+        QWEN3_ASR_CONV_LAYERS,
+        compute_encoder_output_length,
+    )
+except ImportError:  # flat layout on the Hub: sibling modules, no package
+    from asr_config import (
+        QWEN3_ASR_CHUNK_FRAMES,
+        QWEN3_ASR_CONV_LAYERS,
+        compute_encoder_output_length,
+    )
+
+# The encoder's tensors inside the ASR checkpoint.
+_WEIGHT_PREFIX = "model.audio_tower."
+
+
+@dataclass
+class Qwen3ASRAudioOutput(ModelOutput):
+    """Padded encoder states and the mask of their valid frames."""
+
+    last_hidden_state: torch.Tensor | None = None
+    # (batch, tokens), True on valid frames. ASRModel checks the prompt's audio
+    # token counts against it, as it does for Granite's own returned mask.
+    attention_mask: torch.Tensor | None = None
+
+
+class Qwen3ASRAudioEncoder(nn.Module):
+    """`Qwen3ASREncoder` with its packed output split back into a padded batch."""
+
+    def __init__(self, encoder: Qwen3ASREncoder) -> None:
+        """Wrap a loaded encoder; `config` is exposed for encoder_dim detection."""
+        super().__init__()
+        self.encoder: Qwen3ASREncoder = encoder
+        self.config: Qwen3ASREncoderConfig = encoder.config
+
+    def forward(
+        self, input_features: torch.Tensor, attention_mask: torch.Tensor | None = None
+    ) -> Qwen3ASRAudioOutput:
+        """Encode `(batch, n_mels, frames)` mel features to `(batch, tokens, d_model)`.
+
+        `frames` must be a multiple of 100, which Qwen3-ASR's feature extractor
+        guarantees. Without `attention_mask` every frame counts as audio.
+        """
+        batch = input_features.shape[0]
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                batch, input_features.shape[-1], dtype=torch.long, device=input_features.device
+            )
+        mask = attention_mask.to(device=input_features.device, dtype=torch.long)
+        packed = self.encoder(input_features=input_features, input_features_mask=mask)
+        flat = packed.last_hidden_state
+
+        lengths = compute_encoder_output_length(
+            mask.sum(dim=-1), QWEN3_ASR_CONV_LAYERS, QWEN3_ASR_CHUNK_FRAMES
+        )
+        if int(lengths.sum()) != flat.shape[0]:
+            msg = (
+                f"Qwen3-ASR encoder packed {flat.shape[0]} frames but the chunked length "
+                f"formula expects {int(lengths.sum())} ({lengths.tolist()} per clip)."
+            )
+            raise ValueError(msg)
+        rows = list(flat.split(lengths.tolist()))
+        hidden = nn.utils.rnn.pad_sequence(rows, batch_first=True)
+        valid = torch.arange(hidden.shape[1], device=hidden.device)[None, :] < lengths[:, None]
+        return Qwen3ASRAudioOutput(last_hidden_state=hidden, attention_mask=valid)
+
+    @classmethod
+    def from_pretrained(
+        cls, repo_id: str, dtype: torch.dtype, attn_implementation: str | None
+    ) -> "Qwen3ASRAudioEncoder":
+        """Build the encoder from `repo_id`'s audio config and load only its tensors.
+
+        Built on CPU, not the meta device: the sinusoidal position table is a
+        non-persistent buffer that is absent from the checkpoint, so a meta
+        build would leave it without values.
+        """
+        audio_config = Qwen3ASRConfig.from_pretrained(repo_id).audio_config
+        if not isinstance(audio_config, Qwen3ASREncoderConfig):
+            msg = f"{repo_id} has no Qwen3-ASR audio_config (got {type(audio_config).__name__})"
+            raise TypeError(msg)
+        encoder = Qwen3ASREncoder(audio_config)
+
+        state: dict[str, torch.Tensor] = {}
+        for path, names in _encoder_shards(repo_id).items():
+            with safe_open(path, framework="pt") as f:
+                for name in names:
+                    state[name.removeprefix(_WEIGHT_PREFIX)] = f.get_tensor(name)
+        # strict: a renamed or missing tensor must fail here, not train on random weights.
+        encoder.load_state_dict(state, strict=True)
+        # PreTrainedModel.to is functools.wraps'd, which pyright cannot bind as a
+        # method; it casts in place.
+        cast(nn.Module, encoder).to(dtype=dtype)
+        # After the cast: transformers warns when FA2 is set on an fp32 model.
+        # Its other warning, "not initialized on GPU", is expected here; the
+        # Trainer moves the whole ASRModel to the GPU before the first step.
+        if attn_implementation:
+            encoder.set_attn_implementation(attn_implementation)
+        return cls(encoder)
+
+
+def _encoder_shards(repo_id: str) -> dict[str, list[str]]:
+    """Local path of each checkpoint file holding encoder tensors -> those tensor names."""
+    try:
+        index_path = hf_hub_download(repo_id, "model.safetensors.index.json")
+    except EntryNotFoundError:  # a single-file checkpoint (the 1.7B and 0.6B are) has no index
+        path = hf_hub_download(repo_id, "model.safetensors")
+        with safe_open(path, framework="pt") as f:
+            keys = cast("list[str]", f.keys())
+            names = [k for k in keys if k.startswith(_WEIGHT_PREFIX)]
+        return {path: names}
+    weight_map: dict[str, str] = json.loads(Path(index_path).read_text())["weight_map"]
+    shards: dict[str, list[str]] = {}
+    for name, shard in weight_map.items():
+        if name.startswith(_WEIGHT_PREFIX):
+            shards.setdefault(shard, []).append(name)
+    return {hf_hub_download(repo_id, shard): names for shard, names in shards.items()}

requirements.txt

@@ -8,3 +8,5 @@ transformers @ git+https://github.com/huggingface/transformers@91ea1192a00e0ebb6
 torch
 torchaudio
 peft
+# Nemotron's streaming feature extractor (return_speakers=True) imports it.
+librosa

tokenizer.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a454aca253aa94d056abd50384d2b15bc468030ad401f9d644b20416b1d46055
-size 19989509
+oid sha256:33b674fb8444e2553eae8f1b261093371920a28ef75b5c18f4deb3f9217ed0ba
+size 11422834

tokenizer_config.json

@@ -1,8 +1,5 @@
 {
   "add_prefix_space": false,
-  "audio_bos_token": "<|audio_start|>",
-  "audio_eos_token": "<|audio_end|>",
-  "audio_token": "<|audio_pad|>",
   "backend": "tokenizers",
   "bos_token": null,
   "clean_up_tokenization_spaces": false,
@@ -11,25 +8,11 @@
   "extra_special_tokens": [
     "<audio>"
   ],
-  "image_token": "<|image_pad|>",
   "is_local": false,
   "local_files_only": false,
-  "model_max_length": 262144,
-  "model_specific_special_tokens": {
-    "audio_bos_token": "<|audio_start|>",
-    "audio_eos_token": "<|audio_end|>",
-    "audio_token": "<|audio_pad|>",
-    "image_token": "<|image_pad|>",
-    "video_token": "<|video_pad|>",
-    "vision_bos_token": "<|vision_start|>",
-    "vision_eos_token": "<|vision_end|>"
-  },
+  "model_max_length": 131072,
   "pad_token": "<|endoftext|>",
-  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",
-  "unk_token": null,
-  "video_token": "<|video_pad|>",
-  "vision_bos_token": "<|vision_start|>",
-  "vision_eos_token": "<|vision_end|>"
+  "unk_token": null
 }

training_args.bin

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