Update custom model files, README, and requirements

Browse files

Files changed (12) hide show

README.md +128 -87
alignment.py +76 -261
asr_attention.py +171 -0
asr_config.py +51 -25
asr_layers.py +201 -0
asr_modeling.py +433 -569
asr_pipeline.py +188 -68
asr_processing.py +108 -47
asr_types.py +108 -0
diarization.py +62 -712
handler.py +28 -25
projectors.py +17 -7

README.md CHANGED Viewed

@@ -1,42 +1,46 @@
 ---
 license: mit
 language:
-- en
 datasets:
-- mythicinfinity/libriheavy
-- MLCommons/peoples_speech
-- fixie-ai/common_voice_17_0
-- speechcolab/gigaspeech
-- kensho/spgispeech
-- facebook/voxpopuli
-- edinburghcstr/ami
-- sanchit-gandhi/tedlium-data
 base_model:
-- ibm-granite/granite-speech-5.0-470m-turboctc
-- Qwen/Qwen3.5-2B
 pipeline_tag: automatic-speech-recognition
 tags:
-- asr
-- speech-recognition
-- audio
-- qwen
-- granite-speech
-- lora
 library_name: transformers
 ---
 # Tiny Audio
-An English speech recognition model that outputs punctuated, capitalized, formatted text. Built with [Tiny Audio](https://github.com/alexkroman/tiny-audio)—a minimal, hackable ASR framework.
-A frozen Granite Speech encoder is connected to a frozen Qwen3.5-2B decoder through a trained MLP projector, with LoRA adapters on the decoder. Only ~80M parameters are trained.
 ## Quick Start
 ```python
 from transformers import pipeline
-pipe = pipeline("automatic-speech-recognition", model="mazesmazes/tiny-audio", trust_remote_code=True)
 result = pipe("audio.wav")
 print(result["text"])
 # The quarterly revenue grew by 12% according to Dr. Smith.
@@ -49,7 +53,9 @@ print(result["text"])
 ```python
 from transformers import pipeline
-pipe = pipeline("automatic-speech-recognition", model="mazesmazes/tiny-audio", trust_remote_code=True)
 # From file
 result = pipe("audio.wav")
@@ -60,6 +66,7 @@ result = pipe("https://example.com/audio.mp3")
 # From numpy array (must be 16kHz)
 import numpy as np
 audio = np.random.randn(16000).astype(np.float32)  # 1 second
 result = pipe(audio)
 ```
@@ -75,7 +82,11 @@ for r in results:
 ### Word-Level Timestamps
-`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.
 ```python
 result = pipe("audio.wav", return_timestamps=True)
@@ -87,7 +98,11 @@ for word in result["words"]:
 ### Speaker Diarization
-`return_speakers=True` labels every word with a speaker (it implies `return_timestamps=True`). Speakers come from [Nemotron-3-Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization), run once over the whole recording, so labels stay consistent across hour-long meetings; it handles overlapping speech and up to 8 speakers. Each word goes to the speaker most active while it was spoken.
 Nemotron-3-Diarization is in transformers `main`, not yet a release:
@@ -113,13 +128,17 @@ for w in result["words"]:
 #   0.22-  0.69  SPEAKER_0  Daniel.
 ```
-Speakers are numbered by when they first speak. The number of speakers is detected automatically. If you know it, pass `num_speakers` (exact) or `max_speakers` (an upper bound); any extra voices' words go to the closest kept speaker:
 ```python
 result = pipe("call.wav", return_speakers=True, num_speakers=2)
 ```
-If alignment or diarization fails (for example, transformers without Nemotron), the transcript is still returned and the error is reported under `result["timestamp_error"]` or `result["diarization_error"]`.
 ### GPU Inference
@@ -137,78 +156,90 @@ pipe = pipeline(
 ## Benchmarks
-Word error rate (%, lower is better) on 1,000 samples per dataset, scored 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.80 |
-| SPGISpeech | 2.24 |
-| LibriSpeech test-other | 3.09 |
-| TED-LIUM | 3.74 |
-| LoquaciousSet † | 6.10 |
-| Common Voice | 6.62 |
-| VoxPopuli | 6.96 |
-| AMI (IHM) | 8.88 |
-| GigaSpeech | 9.07 |
-| Earnings22 † | 10.54 |
-| People's Speech | 17.69 |
-| AMI (SDM) | 23.59 |
-| **Mean (12 sets)** | **8.36** |
 † Held out: no data from this source was used in training.
 ## Architecture
-```
-Audio (16kHz) → Granite Speech encoder (frozen) → MLP projector (trained) → Qwen3.5-2B + LoRA (trained adapters) → Text
 ```
-| Component | Model | Parameters | Status |
-|-----------|-------|------------|--------|
-| Audio Encoder | [granite-speech-5.0-470m-turboctc](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc) | ~470M | Frozen |
-| Projector | 2-layer MLP (hidden 4096) | 12.6M | Trained |
-| Language Model | [Qwen3.5-2B](https://huggingface.co/Qwen/Qwen3.5-2B) | ~2B | Frozen |
-| LoRA adapters | r=64, alpha=64, all linear layers | 67.3M | Trained |
 ### How It Works
 1. **Audio encoder**: Granite Speech turns 16kHz audio into frame-level embeddings.
-2. **Projector**: A 2-layer MLP maps those embeddings into the decoder's embedding space. Each projected frame replaces an `<audio>` placeholder token in the prompt.
-3. **Language model**: Qwen3.5-2B, adapted with LoRA, generates the transcript conditioned on the projected audio and the prompt *"Transcribe the speech with proper punctuation and capitalization"*.
-At inference, 0.25s of silence is prepended to each clip (`inference_lead_in_seconds`). This keeps the first word from being dropped on clips that start mid-speech.
 ## Model Specifications
-| Specification | Value |
-|---------------|-------|
-| Input | Audio (16kHz mono) |
-| Output | Punctuated, capitalized text with formatted numbers |
-| Max Audio Length | ~30 seconds per call |
-| Vocabulary | Qwen3.5 tokenizer |
-| Languages | English only |
-| Generation | Greedy decoding (num_beams=1, do_sample=False), max 256 new tokens |
 ## 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) |
-The training recipe is [`configs/experiments/granite_qwen_frozen.yaml`](https://github.com/alexkroman/tiny-audio/blob/main/configs/experiments/granite_qwen_frozen.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.
-- **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 a release includes Nemotron-3-Diarization.
 - **Accuracy**: May degrade on:
   - Far-field and overlapping speech (see AMI SDM)
   - Noisy or low-quality audio
@@ -216,21 +247,27 @@ The training recipe is [`configs/experiments/granite_qwen_frozen.yaml`](https://
 ## Files
-| File | Description |
-|------|-------------|
-| `config.json` | Model configuration |
-| `model.safetensors` | Projector weights (~50MB) |
-| `adapter_config.json` / `adapter_model.safetensors` | LoRA adapters for the decoder (~270MB) |
-| `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 previous GLM-ASR + Qwen3-0.6B model is still available at revision `glm-asr-qwen3-0.6b`:
 ```python
-pipe = pipeline("automatic-speech-recognition", model="mazesmazes/tiny-audio", revision="glm-asr-qwen3-0.6b", trust_remote_code=True)
 ```
 ## Citation
@@ -250,16 +287,20 @@ If you use this model, please cite:
 ## Links
 - [GitHub Repository](https://github.com/alexkroman/tiny-audio) - Train your own model
-- [Free 3.5-hour Course](https://github.com/alexkroman/tiny-audio/blob/main/docs/course/0-course-overview.md) - Learn ASR from scratch
 - [Live Demo](https://huggingface.co/spaces/mazesmazes/tiny-audio) - Try it in your browser
 ## Acknowledgments
-- [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc) for the audio encoder
 - [Qwen3.5](https://huggingface.co/Qwen/Qwen3.5-2B) 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
 ## License

 ---
 license: mit
 language:
+  - en
 datasets:
+  - mythicinfinity/libriheavy
+  - MLCommons/peoples_speech
+  - fixie-ai/common_voice_17_0
+  - speechcolab/gigaspeech
+  - kensho/spgispeech
+  - facebook/voxpopuli
+  - edinburghcstr/ami
+  - sanchit-gandhi/tedlium-data
 base_model:
+  - ibm-granite/granite-speech-5.0-470m-turboctc
+  - Qwen/Qwen3.5-2B
 pipeline_tag: automatic-speech-recognition
 tags:
+  - asr
+  - speech-recognition
+  - audio
+  - qwen
+  - granite-speech
+  - lora
 library_name: transformers
 ---
 # Tiny Audio
+An English speech recognition model that outputs punctuated, capitalized, formatted text. Built with
+[Tiny Audio](https://github.com/alexkroman/tiny-audio)—a minimal, hackable ASR framework.
+A frozen Granite Speech encoder is connected to a frozen Qwen3.5-2B decoder through a trained MLP
+projector, with LoRA adapters on the decoder. Only ~80M parameters are trained.
 ## Quick Start
 ```python
 from transformers import pipeline
+pipe = pipeline(
+    "automatic-speech-recognition", model="mazesmazes/tiny-audio", trust_remote_code=True
+)
 result = pipe("audio.wav")
 print(result["text"])
 # The quarterly revenue grew by 12% according to Dr. Smith.
 ```python
 from transformers import pipeline
+pipe = pipeline(
+    "automatic-speech-recognition", model="mazesmazes/tiny-audio", trust_remote_code=True
+)
 # From file
 result = pipe("audio.wav")
 # From numpy array (must be 16kHz)
 import numpy as np
 audio = np.random.randn(16000).astype(np.float32)  # 1 second
 result = pipe(audio)
 ```
 ### Word-Level Timestamps
+`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.
 ```python
 result = pipe("audio.wav", return_timestamps=True)
 ### Speaker Diarization
+`return_speakers=True` labels every word with a speaker (it implies `return_timestamps=True`).
+Speakers come from [Nemotron-3-Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization),
+run once over the whole recording, so labels stay consistent across hour-long meetings; it handles
+overlapping speech and up to 8 speakers. Each word goes to the speaker most active while it was
+spoken.
 Nemotron-3-Diarization is in transformers `main`, not yet a release:
 #   0.22-  0.69  SPEAKER_0  Daniel.
 ```
+Speakers are numbered by when they first speak. The number of speakers is detected automatically. If
+you know it, pass `num_speakers` (exact) or `max_speakers` (an upper bound); any extra voices' words
+go to the closest kept speaker:
 ```python
 result = pipe("call.wav", return_speakers=True, num_speakers=2)
 ```
+If alignment or diarization fails (for example, transformers without Nemotron), the transcript is
+still returned and the error is reported under `result["timestamp_error"]` or
+`result["diarization_error"]`.
 ### GPU Inference
 ## Benchmarks
+Word error rate (%, lower is better) on 1,000 samples per dataset, scored 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.80 |
+| SPGISpeech             |     2.24 |
+| LibriSpeech test-other |     3.09 |
+| TED-LIUM               |     3.74 |
+| LoquaciousSet †        |     6.10 |
+| Common Voice           |     6.62 |
+| VoxPopuli              |     6.96 |
+| AMI (IHM)              |     8.88 |
+| GigaSpeech             |     9.07 |
+| Earnings22 †           |    10.54 |
+| People's Speech        |    17.69 |
+| AMI (SDM)              |    23.59 |
+| **Mean (12 sets)**     | **8.36** |
 † Held out: no data from this source was used in training.
 ## Architecture
+```text
+Audio (16kHz) → Granite Speech encoder (frozen) → MLP projector (trained)
+  → Qwen3.5-2B + LoRA (trained adapters) → Text
 ```
+| Component      | Model                                                                                                   | Parameters | Status  |
+| -------------- | ------------------------------------------------------------------------------------------------------- | ---------- | ------- |
+| Audio Encoder  | [granite-speech-5.0-470m-turboctc](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc) | ~470M      | Frozen  |
+| Projector      | 2-layer MLP (hidden 4096)                                                                               | 12.6M      | Trained |
+| Language Model | [Qwen3.5-2B](https://huggingface.co/Qwen/Qwen3.5-2B)                                                    | ~2B        | Frozen  |
+| LoRA adapters  | r=64, alpha=64, all linear layers                                                                       | 67.3M      | Trained |
 ### How It Works
 1. **Audio encoder**: Granite Speech turns 16kHz audio into frame-level embeddings.
+1. **Projector**: A 2-layer MLP maps those embeddings into the decoder's embedding space. Each
+   projected frame replaces an `<audio>` placeholder token in the prompt.
+1. **Language model**: Qwen3.5-2B, adapted with LoRA, generates the transcript conditioned on the
+   projected audio and the prompt *"Transcribe the speech with proper punctuation and
+   capitalization"*.
+At inference, 0.25s of silence is prepended to each clip (`inference_lead_in_seconds`). This keeps
+the first word from being dropped on clips that start mid-speech.
 ## Model Specifications
+| Specification    | Value                                                              |
+| ---------------- | ------------------------------------------------------------------ |
+| Input            | Audio (16kHz mono)                                                 |
+| Output           | Punctuated, capitalized text with formatted numbers                |
+| Max Audio Length | ~30 seconds per call                                               |
+| Vocabulary       | Qwen3.5 tokenizer                                                  |
+| Languages        | English only                                                       |
+| Generation       | Greedy decoding (num_beams=1, do_sample=False), max 256 new tokens |
 ## 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)                                                                                                                                   |
+The training recipe is
+[`configs/experiments/granite_qwen_frozen.yaml`](https://github.com/alexkroman/tiny-audio/blob/main/configs/experiments/granite_qwen_frozen.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.
+- **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 a release includes Nemotron-3-Diarization.
 - **Accuracy**: May degrade on:
   - Far-field and overlapping speech (see AMI SDM)
   - Noisy or low-quality audio
 ## Files
+| File                                                               | Description                                              |
+| ------------------------------------------------------------------ | -------------------------------------------------------- |
+| `config.json`                                                      | Model configuration                                      |
+| `model.safetensors`                                                | Projector weights (~50MB)                                |
+| `adapter_config.json` / `adapter_model.safetensors`                | LoRA adapters for the decoder (~270MB)                   |
+| `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 previous GLM-ASR + Qwen3-0.6B model is still available at revision `glm-asr-qwen3-0.6b`:
 ```python
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model="mazesmazes/tiny-audio",
+    revision="glm-asr-qwen3-0.6b",
+    trust_remote_code=True,
+)
 ```
 ## Citation
 ## Links
 - [GitHub Repository](https://github.com/alexkroman/tiny-audio) - Train your own model
+- [Free 3.5-hour Course](https://github.com/alexkroman/tiny-audio/blob/main/docs/course/0-course-overview.md)
+  \- Learn ASR from scratch
 - [Live Demo](https://huggingface.co/spaces/mazesmazes/tiny-audio) - Try it in your browser
 ## Acknowledgments
+- [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-5.0-470m-turboctc) for the
+  audio encoder
 - [Qwen3.5](https://huggingface.co/Qwen/Qwen3.5-2B) 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
 ## License

alignment.py CHANGED Viewed

@@ -1,11 +1,45 @@
-"""Forced alignment for word-level timestamps: Qwen3-ForcedAligner, wav2vec2 as fallback."""
 import numpy as np
 import torch
 def _get_device() -> str:
-    """Get best available device for non-transformers models."""
     if torch.cuda.is_available():
         return "cuda"
     if torch.backends.mps.is_available():
@@ -13,235 +47,6 @@ def _get_device() -> str:
     return "cpu"
-class ForcedAligner:
-    """Lazy-loaded forced aligner for word-level timestamps using torchaudio wav2vec2.
-    The CTC Viterbi search itself is `torchaudio.functional.forced_align`; this
-    class owns the model singleton, the word-level tokenization and the
-    frame-to-seconds conversion around it.
-    """
-    _bundle = None
-    _model = None
-    _labels = None
-    _dictionary = None
-    @classmethod
-    def get_instance(cls, device: str = "cuda"):
-        """Get or create the forced alignment model (singleton).
-        Args:
-            device: Device to run model on ("cuda" or "cpu")
-        Returns:
-            Tuple of (model, labels, dictionary)
-        """
-        if cls._model is None:
-            import torchaudio
-            cls._bundle = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H
-            cls._model = cls._bundle.get_model().to(device)
-            cls._model.eval()
-            cls._labels = cls._bundle.get_labels()
-            cls._dictionary = {c: i for i, c in enumerate(cls._labels)}
-        return cls._model, cls._labels, cls._dictionary
-    @staticmethod
-    def _align_tokens(
-        emission: torch.Tensor, tokens: list[int], blank_id: int = 0
-    ) -> list[tuple[int, float, float]]:
-        """Viterbi-align `tokens` to `emission` and return one span per token.
-        Delegates to `torchaudio.functional.forced_align` (a C++/CUDA CTC
-        Viterbi) and `merge_tokens`, which collapse the per-frame path into
-        `(token_id, start_frame, end_frame)` spans with `end_frame` exclusive.
-        Guarantees:
-        - All tokens are emitted exactly once, in order (strictly monotonic).
-        - When no valid path exists -- more tokens than frames, or an emission
-          that assigns every path zero probability -- the tokens are spread
-          uniformly over the frames instead of raising, so a timestamp request
-          degrades to coarse timings rather than failing the transcription.
-        Args:
-            emission: Log-softmax emission matrix of shape (num_frames, num_classes)
-            tokens: Target token indices
-            blank_id: Index of the CTC blank (default 0)
-        """
-        import torchaudio.functional as F  # noqa: N812
-        num_frames = emission.size(0)
-        num_tokens = len(tokens)
-        if num_tokens == 0:
-            return []
-        def _uniform() -> list[tuple[int, float, float]]:
-            frames_per_token = num_frames / num_tokens
-            return [
-                (tokens[i], i * frames_per_token, (i + 1) * frames_per_token)
-                for i in range(num_tokens)
-            ]
-        log_probs = emission.detach().float().unsqueeze(0)
-        targets = torch.tensor([tokens], dtype=torch.int32, device=log_probs.device)
-        try:
-            aligned, scores = F.forced_align(log_probs, targets, blank=blank_id)
-        except RuntimeError:
-            # "targets length is too long for CTC": no monotonic path fits.
-            return _uniform()
-        if not torch.isfinite(scores).all():
-            # Every path has -inf log-probability; the returned path is junk.
-            return _uniform()
-        spans = F.merge_tokens(aligned[0], scores[0].exp(), blank=blank_id)
-        if len(spans) != num_tokens:
-            return _uniform()
-        return [(int(span.token), float(span.start), float(span.end)) for span in spans]
-    # Offset compensation for Wav2Vec2-BASE systematic bias (in seconds)
-    # Calibrated on librispeech-alignments dataset
-    START_OFFSET = 0.06  # Subtract from start times (shift earlier)
-    END_OFFSET = -0.03  # Add to end times (shift later)
-    @staticmethod
-    def _tokenize_words(
-        text: str, dictionary: dict, blank_id: int = 0
-    ) -> tuple[list[str], list[int]]:
-        """Split `text` into alignable words and the CTC token sequence for them.
-        Returns `(words, tokens)` where `tokens` is each word's character ids
-        joined by the separator id, and `words` lists only the words that
-        contributed at least one token -- so the token groups the Viterbi path
-        produces line up 1:1 with `words`.
-        That pairing is the whole point. A word can produce no usable tokens in
-        two ways, and both used to leave an empty group behind while still
-        emitting the separators on either side of it:
-        - every character is outside the label set (digits, "...");
-        - every character maps to a control id. In wav2vec2's label set index 0
-          is spelled "-", so a hyphen is the *blank* symbol, and "--" became
-          two blank target tokens rather than a word.
-        Pairing groups against a plain `text.split()` then shifted every later
-        word onto the previous word's timing and dropped the last one:
-        "hello -- world" returned `hello, --` with `--` carrying WORLD's frames.
-        `blank_id` defaults to 0 to match the `blank_id=0` that `align` passes
-        to `_align_tokens`.
-        """
-        separator_id = dictionary.get("|", dictionary.get(" ", 0))
-        # Neither may appear as a target token: the separator would split the
-        # word in two, and the blank is what the trellis emits *between*
-        # tokens. Either one desyncs the group-to-word pairing.
-        control_ids = {separator_id, blank_id}
-        words: list[str] = []
-        tokens: list[int] = []
-        for word in text.split():
-            word_tokens = [
-                token_id
-                for c in word.upper()
-                if (token_id := dictionary.get(c)) is not None and token_id not in control_ids
-            ]
-            if not word_tokens:
-                continue
-            if tokens:
-                tokens.append(separator_id)
-            tokens.extend(word_tokens)
-            words.append(word)
-        return words, tokens
-    @classmethod
-    def align(
-        cls,
-        audio: np.ndarray | torch.Tensor,
-        text: str,
-        sample_rate: int = 16000,
-    ) -> list[dict]:
-        """Align transcript to audio and return word-level timestamps.
-        Args:
-            audio: Audio waveform as a numpy array or torch tensor
-            text: Transcript text to align
-            sample_rate: Audio sample rate (default 16000)
-        Returns:
-            List of dicts with 'word', 'start', 'end' keys
-        """
-        import torchaudio
-        device = _get_device()
-        model, _labels, dictionary = cls.get_instance(device)
-        assert cls._bundle is not None
-        assert dictionary is not None
-        # Convert audio to tensor (copy to ensure array is writable)
-        if isinstance(audio, np.ndarray):
-            waveform = torch.from_numpy(audio.copy()).float()
-        else:
-            waveform = audio.clone().float()
-        # Ensure 2D (channels, time)
-        if waveform.dim() == 1:
-            waveform = waveform.unsqueeze(0)
-        # Resample if needed (wav2vec2 expects 16kHz)
-        if sample_rate != cls._bundle.sample_rate:
-            waveform = torchaudio.functional.resample(
-                waveform, sample_rate, cls._bundle.sample_rate
-            )
-        waveform = waveform.to(device)
-        # Get emissions from model
-        with torch.inference_mode():
-            emissions, _ = model(waveform)
-            emissions = torch.log_softmax(emissions, dim=-1)
-        emission = emissions[0].cpu()
-        # Tokenize per word so each Viterbi token group maps back to the word
-        # it came from; words with no representable characters are dropped.
-        words, tokens = cls._tokenize_words(text, dictionary)
-        if not tokens:
-            return []
-        alignment_path = cls._align_tokens(emission, tokens, blank_id=0)
-        # Convert frame indices to time (model stride is 320 samples at 16kHz = 20ms)
-        frame_duration = 320 / cls._bundle.sample_rate
-        # Group token spans into words on the separator token. `_tokenize_words`
-        # guarantees one non-empty group per entry in `words`, so the groups
-        # zip 1:1 with the word list.
-        separator_id = dictionary.get("|", dictionary.get(" ", 0))
-        groups: list[list[tuple[float, float]]] = []
-        current: list[tuple[float, float]] = []
-        for token_id, start_frame, end_frame in alignment_path:
-            if token_id == separator_id:
-                if current:
-                    groups.append(current)
-                    current = []
-            else:
-                current.append((start_frame, end_frame))
-        if current:
-            groups.append(current)
-        # Apply separate offset compensation for start/end (Wav2Vec2 systematic bias)
-        word_timestamps = []
-        for word, frames in zip(words, groups):
-            start_frame, end_frame = frames[0][0], frames[-1][1]
-            word_timestamps.append(
-                {
-                    "word": word,
-                    "start": max(0.0, start_frame * frame_duration - cls.START_OFFSET),
-                    "end": max(0.0, end_frame * frame_duration - cls.END_OFFSET),
-                }
-            )
-        return word_timestamps
 class QwenForcedAligner:
     """Word timestamps from Qwen3-ForcedAligner-0.6B (transformers >= 5.17).
@@ -259,67 +64,71 @@ class QwenForcedAligner:
     MODEL_ID = "Qwen/Qwen3-ForcedAligner-0.6B-hf"
     MAX_SECONDS = 300.0
     BATCH_SIZE = 8
-    _model = None
-    _processor = None
     @classmethod
-    def get_instance(cls):
-        if cls._model is None:
-            from transformers import AutoProcessor, Qwen3ASRForTokenClassification
             device = _get_device()
             # Batches are padded, and Metal sdpa returns NaN for fully masked rows.
             attn = "eager" if device == "mps" else "sdpa"
             model = Qwen3ASRForTokenClassification.from_pretrained(
                 cls.MODEL_ID, dtype=torch.bfloat16, attn_implementation=attn
             )
-            cls._model = model.to(device).eval()
-            cls._processor = AutoProcessor.from_pretrained(cls.MODEL_ID)
         return cls._model, cls._processor
     @staticmethod
-    def _to_16k(audio, sample_rate: int) -> np.ndarray:
         if isinstance(audio, torch.Tensor):
-            audio = audio.cpu().numpy()
-        audio = np.asarray(audio, dtype=np.float32).reshape(-1)
         if sample_rate != 16000:
-            import torchaudio
-            audio = torchaudio.functional.resample(
-                torch.from_numpy(audio), sample_rate, 16000
-            ).numpy()
-        return audio
     @classmethod
     @torch.inference_mode()
     def align_chunks(
         cls,
-        chunks: list[tuple[np.ndarray, str]],
         sample_rate: int = 16000,
         language: str = "English",
-    ) -> list[list[dict]]:
         """Align each `(audio, text)` pair (each <= 5 min); one word list per pair."""
-        from transformers.models.qwen3_asr.processing_qwen3_asr import _clean_tokens
-        results: list[list[dict]] = [[] for _ in chunks]
-        todo = []
         for i, (raw, text) in enumerate(chunks):
             kept = [w for w in text.split() if _clean_tokens([w])]
             if not kept:
                 continue
             audio = cls._to_16k(raw, sample_rate)
             if len(audio) / 16000 > cls.MAX_SECONDS:
-                raise ValueError(
                     f"chunk {i} is {len(audio) / 16000:.0f} s; Qwen3-ForcedAligner takes at most "
                     f"{cls.MAX_SECONDS:.0f} s -- cut the audio first"
                 )
             todo.append((i, audio, kept))
         if not todo:  # nothing to time: don't load the model
             return results
         model, processor = cls.get_instance()
-        for b in range(0, len(todo), cls.BATCH_SIZE):
-            batch = todo[b : b + cls.BATCH_SIZE]
-            inputs, word_lists = processor.prepare_forced_aligner_inputs(
                 [audio for _, audio, _ in batch],
                 [" ".join(kept) for _, _, kept in batch],
                 language=language,
@@ -332,7 +141,7 @@ class QwenForcedAligner:
                     if v.is_floating_point()
                     else v.to(model.device)
                 )
-                for k, v in inputs.items()
             }
             logits = model(**inputs).logits
             timed = processor.decode_forced_alignment(
@@ -346,6 +155,12 @@ class QwenForcedAligner:
         return results
     @classmethod
-    def align(cls, audio, text: str, sample_rate: int = 16000, language: str = "English"):
-        """ForcedAligner.align's interface for one clip of at most 5 minutes."""
         return cls.align_chunks([(audio, text)], sample_rate, language)[0]

+"""Forced alignment for word-level timestamps with Qwen3-ForcedAligner."""
+from __future__ import annotations
+import itertools
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, TypedDict
 import numpy as np
+import numpy.typing as npt
 import torch
+import torchaudio
+if TYPE_CHECKING:
+    from transformers import Qwen3ASRForTokenClassification, Qwen3ASRProcessor
+    from transformers.models.qwen3_asr.processing_qwen3_asr import (
+        _clean_tokens,  # pyright: ignore[reportPrivateUsage]  # no public word filter
+    )
+    _QWEN3_ASR_IMPORT_ERROR: ImportError | None = None
+else:
+    # Qwen3-ForcedAligner needs transformers >= 5.17; keep this module importable
+    # on older versions and only fail when alignment is actually requested.
+    try:
+        from transformers import Qwen3ASRForTokenClassification, Qwen3ASRProcessor
+        from transformers.models.qwen3_asr.processing_qwen3_asr import _clean_tokens
+        _QWEN3_ASR_IMPORT_ERROR = None
+    except ImportError as e:
+        _QWEN3_ASR_IMPORT_ERROR = e
+class AlignedWord(TypedDict):
+    """One aligned word: the original word and its span in seconds."""
+    word: str
+    start: float
+    end: float
 def _get_device() -> str:
+    """Get best available device for inference."""
     if torch.cuda.is_available():
         return "cuda"
     if torch.backends.mps.is_available():
     return "cpu"
 class QwenForcedAligner:
     """Word timestamps from Qwen3-ForcedAligner-0.6B (transformers >= 5.17).
     MODEL_ID = "Qwen/Qwen3-ForcedAligner-0.6B-hf"
     MAX_SECONDS = 300.0
     BATCH_SIZE = 8
+    _model: Qwen3ASRForTokenClassification | None = None
+    _processor: Qwen3ASRProcessor | None = None
     @classmethod
+    def get_instance(cls) -> tuple[Qwen3ASRForTokenClassification, Qwen3ASRProcessor]:
+        """Load the aligner model and processor once, then return the cached pair."""
+        if cls._model is None or cls._processor is None:
+            if _QWEN3_ASR_IMPORT_ERROR is not None:
+                raise _QWEN3_ASR_IMPORT_ERROR
             device = _get_device()
             # Batches are padded, and Metal sdpa returns NaN for fully masked rows.
             attn = "eager" if device == "mps" else "sdpa"
             model = Qwen3ASRForTokenClassification.from_pretrained(
                 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]
+            model.eval()
+            cls._model = model
+            cls._processor = Qwen3ASRProcessor.from_pretrained(cls.MODEL_ID)
         return cls._model, cls._processor
     @staticmethod
+    def _to_16k(audio: npt.ArrayLike | torch.Tensor, sample_rate: int) -> npt.NDArray[np.float32]:
+        """Flatten `audio` to a float32 mono array resampled to 16 kHz."""
         if isinstance(audio, torch.Tensor):
+            audio = audio.cpu()
+        out: npt.NDArray[np.float32] = np.asarray(audio, dtype=np.float32).reshape(-1)
         if sample_rate != 16000:
+            out = np.asarray(
+                torchaudio.functional.resample(torch.as_tensor(out), sample_rate, 16000),
+                dtype=np.float32,
+            )
+        return out
     @classmethod
     @torch.inference_mode()
     def align_chunks(
         cls,
+        chunks: Sequence[tuple[npt.ArrayLike | torch.Tensor, str]],
         sample_rate: int = 16000,
         language: str = "English",
+    ) -> list[list[AlignedWord]]:
         """Align each `(audio, text)` pair (each <= 5 min); one word list per pair."""
+        if _QWEN3_ASR_IMPORT_ERROR is not None:
+            raise _QWEN3_ASR_IMPORT_ERROR
+        results: list[list[AlignedWord]] = [[] for _ in chunks]
+        todo: list[tuple[int, npt.NDArray[np.float32], list[str]]] = []
         for i, (raw, text) in enumerate(chunks):
             kept = [w for w in text.split() if _clean_tokens([w])]
             if not kept:
                 continue
             audio = cls._to_16k(raw, sample_rate)
             if len(audio) / 16000 > cls.MAX_SECONDS:
+                msg = (
                     f"chunk {i} is {len(audio) / 16000:.0f} s; Qwen3-ForcedAligner takes at most "
                     f"{cls.MAX_SECONDS:.0f} s -- cut the audio first"
                 )
+                raise ValueError(msg)
             todo.append((i, audio, kept))
         if not todo:  # nothing to time: don't load the model
             return results
         model, processor = cls.get_instance()
+        for batch in itertools.batched(todo, cls.BATCH_SIZE):
+            features, word_lists = processor.prepare_forced_aligner_inputs(
                 [audio for _, audio, _ in batch],
                 [" ".join(kept) for _, _, kept in batch],
                 language=language,
                     if v.is_floating_point()
                     else v.to(model.device)
                 )
+                for k, v in features.items()
             }
             logits = model(**inputs).logits
             timed = processor.decode_forced_alignment(
         return results
     @classmethod
+    def align(
+        cls,
+        audio: npt.ArrayLike | torch.Tensor,
+        text: str,
+        sample_rate: int = 16000,
+        language: str = "English",
+    ) -> list[AlignedWord]:
+        """Align one clip of at most 5 minutes; `[{"word", "start", "end"}]`."""
         return cls.align_chunks([(audio, text)], sample_rate, language)[0]

asr_attention.py ADDED Viewed

	@@ -0,0 +1,171 @@

+"""Attention-implementation selection for ASRModel's encoder and decoder loaders.
+Decides which `attn_implementation` each sub-model is loaded with: FA2 only
+where CUDA and a working flash-attn are present and every head fits its
+kernels, eager on MPS for models that build a sliding-window mask, and the
+requested implementation otherwise.
+"""
+import functools
+import importlib
+import logging
+import torch
+from transformers import AutoConfig
+logger = logging.getLogger(__name__)
+# FlashAttention's kernels are compiled for head dimensions up to 256.
+FLASH_ATTENTION_MAX_HEAD_DIM = 256
+def _max_attention_head_dim(text_config: object) -> int | None:
+    """Largest attention head dim across layers, or None if undeterminable.
+    Has to cope with heterogeneous configs: Gemma 4 varies head_dim per layer,
+    and reading `config.head_dim` on one of those raises
+    AmbiguousGlobalPerLayerAttributeError instead of returning a number, so the
+    per-layer list must be consulted explicitly.
+    """
+    per_layer = getattr(text_config, "per_layer_config", None)
+    dims: list[object] = []
+    if per_layer:
+        dims = [getattr(layer, "head_dim", None) for layer in per_layer]
+    else:
+        try:
+            dims = [getattr(text_config, "head_dim", None)]
+        except Exception:
+            # Raised by 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
+@functools.lru_cache(maxsize=16)
+def _has_sliding_window_attention(model_id: str) -> bool:
+    """Whether `model_id`'s text stack has any sliding-window attention layer.
+    Read from the config rather than the loaded module so the answer is
+    available before `from_pretrained` picks an attn implementation.
+    Two independent spellings, because transformers has both. Gemma 4 declares
+    `sliding_window: 512` AND `layer_types: [full_attention,
+    sliding_attention]` (verified against google/gemma-4-E2B-it); Qwen2-style
+    configs carry a non-null `sliding_window` that is inert unless
+    `use_sliding_window` is true, so that flag has to veto the window.
+    Returns True when the config can't be read, so an unknown architecture
+    keeps the conservative eager path on MPS.
+    """
+    try:
+        probe = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+        text_config = probe.get_text_config() if hasattr(probe, "get_text_config") else probe
+    except Exception:
+        logger.warning(
+            "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,
+        )
+        return True
+    layer_types: list[object] = getattr(text_config, "layer_types", None) or []
+    if any("sliding" in str(layer_type) for layer_type in layer_types):
+        return True
+    if getattr(text_config, "use_sliding_window", None) is False:
+        return False
+    return bool(getattr(text_config, "sliding_window", None))
+def resolve_attn_implementation(requested: str | None, model_id: str | None = None) -> str | None:
+    """Coerce flash_attention_2 to sdpa when CUDA isn't available, and avoid sdpa on MPS.
+    FA2 is CUDA-only. On MPS/CPU, requesting it either errors at load or
+    silently falls back to a slower path; either way the user pays the FA2
+    install + import cost for no win. Coerce here so a saved config that
+    pins flash_attention_2 still loads on Mac / CPU-only Linux boxes.
+    MPS needs eager for SOME models. PyTorch's Metal sdpa kernel returns wrong
+    results for cached single-token decode against a sliding-window mask, which
+    is exactly Gemma 4's layout (`sliding_window=512`, four sliding layers per
+    full-attention layer). A full-sequence forward is fine, so the damage shows
+    up only during generation: greedy decode with no cache transcribed
+    correctly while the identical cached decode produced "Mr. and a".
+    The mask is the trigger, so the coercion is scoped to models that actually
+    build one. `model_id` opts into that check; without it every MPS load falls
+    back to eager, which is what this function used to do unconditionally. That
+    blanket version taxed the models it was never protecting: Qwen3.5-2B
+    declares `sliding_window: None` and cannot hit the bug, and eager cost it
+    41% of decode throughput on an M4 Max (19.9 -> 28.1 tok/s, bf16, 64 tokens
+    from a 200-token prompt). Audio encoders are likewise unaffected -- they run
+    full-sequence forwards with no cache at all, and the Granite branch below
+    has been pinning sdpa on MPS for every granite_qwen run to date.
+    SECOND, SEPARATE MPS sdpa HAZARD, and the reason `generate` carries a guard:
+    Metal's sdpa returns NaN for a row containing any fully-masked position, so
+    a LEFT-PADDED batch comes back as NaN logits, argmax 0, and a transcript of
+    "!!!!!!". One pad token is enough. Measured on bare transformers with no
+    tiny_audio code involved (Qwen3.5-2B, bf16): the identical forward is exact
+    to 0.000 on CPU/fp32, and eager is correct on MPS at any padding. Batch 1
+    pads nothing, which is why the eval harness is safe under sdpa -- but this
+    is what the blanket eager was accidentally also protecting, so anything
+    that batches ragged audio on a Mac must ask for eager explicitly. See
+    `_assert_sdpa_safe_on_mps`.
+    """
+    if (
+        torch.backends.mps.is_available()
+        and requested in (None, "sdpa", "flash_attention_2")
+        and (model_id is None or _has_sliding_window_attention(model_id))
+    ):
+        return "eager"
+    # Otherwise fall through: on MPS, sdpa and None (let transformers pick,
+    # which is sdpa where available) are both safe for a model with no
+    # sliding-window layer, and a config pinning flash_attention_2 is still
+    # downgraded to sdpa below.
+    if requested != "flash_attention_2":
+        return requested
+    if not torch.cuda.is_available():
+        return "sdpa"
+    # CUDA alone isn't enough -- the flash_attn package must actually be
+    # importable, and new enough for transformers. It is a source build that
+    # routinely fails on RunPod images (its metadata hook imports torch, so any
+    # broken torch install takes it down with it). Without this check a CUDA
+    # box with no flash-attn raises at from_pretrained instead of quietly
+    # using sdpa, which is the same numerics at lower throughput.
+    #
+    # Looked up on `transformers.utils` at call time rather than bound at
+    # import, so patching it there (as the tests do) reaches this probe.
+    if not importlib.import_module("transformers.utils").is_flash_attn_2_available():
+        logger.warning(
+            "flash_attention_2 requested but flash_attn is not installed or too old; "
+            "falling back to sdpa."
+        )
+        return "sdpa"
+    return requested
+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.
+    FlashAttention only has kernels for head_dim <= 256, and it fails at
+    the first forward rather than at load. Gemma 4 E2B trips this: 7 of
+    its 35 layers use head_dim=512 (the other 28 use 256), so FA2 can
+    never run this architecture regardless of flash-attn version. Detect
+    it from the config and fall back instead of dying a step into
+    training.
+    """
+    attn_implementation = resolve_attn_implementation(requested, model_id)
+    if attn_implementation == "flash_attention_2":
+        probe = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+        text_probe = probe.get_text_config() if hasattr(probe, "get_text_config") else probe
+        head_dim = _max_attention_head_dim(text_probe)
+        if head_dim is not None and head_dim > FLASH_ATTENTION_MAX_HEAD_DIM:
+            logger.warning(
+                "%s has max head_dim=%d, above FlashAttention's limit of %d; "
+                "using sdpa for the decoder.",
+                model_id,
+                head_dim,
+                FLASH_ATTENTION_MAX_HEAD_DIM,
+            )
+            attn_implementation = "sdpa"
+    return attn_implementation

asr_config.py CHANGED Viewed

@@ -1,9 +1,17 @@
 """Configuration for the ASR model: encoder, decoder, projector and training options."""
 import transformers
 # Default conv layers for Whisper/GLM-ASR audio encoders: [(pad, kernel, stride), ...]
-DEFAULT_ENCODER_CONV_LAYERS = [(1, 3, 1), (1, 3, 2)]
 # Granite Speech 5.0 TurboCTC. Its feature extractor already stacks mel frame
 # pairs (100 Hz -> 50 Hz, 320-dim output), so the `mel_length` handed to
@@ -14,7 +22,7 @@ DEFAULT_ENCODER_CONV_LAYERS = [(1, 3, 1), (1, 3, 2)]
 # generic formula below. Verified against the real encoder at 1/2/5/10/20s
 # (T_fe 50/100/250/500/1000 -> 12/25/62/125/250).
 # Full path: 100 Hz mel -> 2x FE stacking -> 4x encoder -> 12.5 Hz.
-GRANITE_ENCODER_CONV_LAYERS = [(0, 2, 2), (0, 2, 2)]
 # Encoders whose `input_features` are (batch, time, feature_dim) rather than
 # Whisper/GLM-ASR's (batch, n_mels, mel_len). Conformer-family checkpoints
@@ -43,7 +51,17 @@ def is_time_major_encoder(audio_model_id: str | None) -> bool:
     return any(m in (audio_model_id or "").lower() for m in _TIME_MAJOR_ENCODER_MARKERS)
-def compute_encoder_output_length(mel_length, conv_layers=None):
     """Apply encoder conv layer formulas to compute output length.
     Works with both Python ints and torch tensors of mel lengths; the formula
@@ -56,6 +74,13 @@ def compute_encoder_output_length(mel_length, conv_layers=None):
     return length
 class ASRConfig(transformers.PretrainedConfig):
     """Configuration class for the ASR model.
@@ -87,7 +112,7 @@ class ASRConfig(transformers.PretrainedConfig):
         llm_dim: int | None = None,
         # 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 | None = None,
         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
@@ -151,7 +176,7 @@ class ASRConfig(transformers.PretrainedConfig):
         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 | 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.
         #
@@ -171,8 +196,8 @@ class ASRConfig(transformers.PretrainedConfig):
         # 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 | None = None,
-        lora_alpha_pattern: dict | 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
@@ -235,8 +260,8 @@ class ASRConfig(transformers.PretrainedConfig):
         max_new_tokens: int | None = None,
         use_cache: bool | None = None,
         no_repeat_ngram_size: int | None = None,
-        **kwargs,
-    ):
         """Initialize ASR model configuration.
         Args:
@@ -336,33 +361,34 @@ class ASRConfig(transformers.PretrainedConfig):
             value = explicit_generation_args[key]
             setattr(self, key, value if value is not None else default)
         if "audio_config" not in kwargs:
-            self.audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
             # Override dtype to match model_dtype
-            self.audio_config.dtype = model_dtype
         else:
-            self.audio_config = kwargs.pop("audio_config")
         if "text_config" not in kwargs:
-            self.text_config = transformers.AutoConfig.from_pretrained(
                 text_model_id, trust_remote_code=True
             )
             # Override dtype to match model_dtype
-            self.text_config.dtype = model_dtype
         else:
-            self.text_config = kwargs.pop("text_config")
-        if isinstance(self.text_config, dict):
             # Reconstruct config from dict using the model_type stored in the dict
-            model_type = self.text_config["model_type"]
-            config_class = transformers.AutoConfig.for_model(model_type).__class__
-            self.text_config = config_class(**self.text_config)
-        if isinstance(self.audio_config, dict):
-            model_type = self.audio_config.get("model_type")
-            if model_type:
-                config_class = transformers.AutoConfig.for_model(model_type).__class__
-                self.audio_config = config_class(**self.audio_config)
         super().__init__(**kwargs)

 """Configuration for the ASR model: encoder, decoder, projector and training options."""
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, overload
 import transformers
+if TYPE_CHECKING:
+    import torch
+ConvLayerSpec = tuple[int, int, int]
 # Default conv layers for Whisper/GLM-ASR audio encoders: [(pad, kernel, stride), ...]
+DEFAULT_ENCODER_CONV_LAYERS: list[ConvLayerSpec] = [(1, 3, 1), (1, 3, 2)]
 # Granite Speech 5.0 TurboCTC. Its feature extractor already stacks mel frame
 # pairs (100 Hz -> 50 Hz, 320-dim output), so the `mel_length` handed to
 # generic formula below. Verified against the real encoder at 1/2/5/10/20s
 # (T_fe 50/100/250/500/1000 -> 12/25/62/125/250).
 # 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)]
 # Encoders whose `input_features` are (batch, time, feature_dim) rather than
 # Whisper/GLM-ASR's (batch, n_mels, mel_len). Conformer-family checkpoints
     return any(m in (audio_model_id or "").lower() for m in _TIME_MAJOR_ENCODER_MARKERS)
+@overload
+def compute_encoder_output_length(
+    mel_length: int, conv_layers: Sequence[ConvLayerSpec] | None = None
+) -> int: ...
+@overload
+def compute_encoder_output_length(
+    mel_length: "torch.Tensor", conv_layers: Sequence[ConvLayerSpec] | None = None
+) -> "torch.Tensor": ...
+def compute_encoder_output_length(
+    mel_length: "int | torch.Tensor", conv_layers: Sequence[ConvLayerSpec] | 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
     return length
+def _config_from_dict(model_type: str, values: dict[str, Any]) -> transformers.PretrainedConfig:
+    """Rebuild a serialized sub-config as the config class registered for `model_type`."""
+    auto_config = transformers.AutoConfig
+    default = auto_config.for_model(model_type)
+    return default.__class__(**values)
 class ASRConfig(transformers.PretrainedConfig):
     """Configuration class for the ASR model.
         llm_dim: int | None = None,
         # 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,
         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
         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.
         #
         # 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
         max_new_tokens: int | None = None,
         use_cache: bool | None = None,
         no_repeat_ngram_size: int | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Initialize ASR model configuration.
         Args:
             value = explicit_generation_args[key]
             setattr(self, key, value if value is not None else default)
+        audio_config: transformers.PretrainedConfig | dict[str, Any]
         if "audio_config" not in kwargs:
+            audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
             # Override dtype to match model_dtype
+            audio_config.dtype = model_dtype
         else:
+            audio_config = kwargs.pop("audio_config")
+        self.audio_config: transformers.PretrainedConfig | dict[str, Any] = audio_config
+        text_config: transformers.PretrainedConfig | dict[str, Any]
         if "text_config" not in kwargs:
+            text_config = transformers.AutoConfig.from_pretrained(
                 text_model_id, trust_remote_code=True
             )
             # Override dtype to match model_dtype
+            text_config.dtype = model_dtype
         else:
+            text_config = kwargs.pop("text_config")
+        if isinstance(text_config, dict):
             # Reconstruct config from dict using the model_type stored in the dict
+            text_config = _config_from_dict(text_config["model_type"], text_config)
+        self.text_config: transformers.PretrainedConfig = text_config
+        if isinstance(audio_config, dict):
+            audio_model_type = audio_config.get("model_type")
+            if audio_model_type:
+                self.audio_config = _config_from_dict(audio_model_type, audio_config)
         super().__init__(**kwargs)

asr_layers.py ADDED Viewed

	@@ -0,0 +1,201 @@

+"""Module-surgery helpers for ASRModel: MPS-safe embeddings and 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.
+"""
+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)"
+        )
+def find_encoder_layer_stack(encoder: nn.Module) -> tuple[str, nn.ModuleList] | None:
+    """Locate the encoder's ordered transformer-block list.
+    Returns `(attribute_path, ModuleList)` or None. Encoders disagree on the
+    attribute name -- Granite Speech uses `layers`, Whisper `encoder.layers`,
+    others `blocks` or `encoder_layers` -- so probe the known names rather
+    than hardcoding one and silently unfreezing nothing on a new encoder.
+    """
+    candidates = ("layers", "blocks", "encoder_layers", "encoder.layers", "encoder.blocks")
+    for path in candidates:
+        obj: nn.Module | None = encoder
+        for part in path.split("."):
+            obj = getattr(obj, part, None)
+            if obj is None:
+                break
+        if isinstance(obj, nn.ModuleList) and len(obj) > 0:
+            return path, obj
+    return None
+def unfreeze_encoder_top_layers(
+    encoder: nn.Module, top_n: int, include_post_projections: bool = True
+) -> list[str]:
+    """Unfreeze the top `top_n` encoder blocks, optionally with the output projections.
+    Assumes the caller has already frozen the whole encoder. Returns the sorted
+    names of every parameter switched back on, so callers (and tests) can
+    assert on exactly what moved rather than trusting a count.
+    The post-stack projections (Granite's `out` / `out_mid`) are included by
+    default: they sit between the top block and the projector, so leaving them
+    frozen forces the newly-trainable blocks to adapt through a fixed output
+    map. That is a reasonable prior, but on this stack it is not what the
+    gradient says. Probing the step-2000 granite_qwen_lora checkpoint on real
+    audio, `out`/`out_mid` carry 33.57M params -- 23% of a top-4 budget -- for
+    an RMS-per-param gradient of 1.14e-05 against 1.4-1.7e-04 in the top
+    blocks, i.e. **0.27% of the selection's squared gradient**. Dropping them
+    moved the aggregate norm from 1.5840 to 1.5827. Pass False to reclaim the
+    parameters; the default stays True so existing recipes are unchanged.
+    The pre-stack `input_linear` stays frozen either way -- it is the feature
+    front-end, the part most expensive to damage and least specialised to the
+    encoder's own CTC head.
+    Raises if the layer stack cannot be found or `top_n` exceeds its depth:
+    silently unfreezing nothing would produce a run that looks like a
+    partial-unfreeze experiment and is actually the frozen baseline.
+    """
+    if top_n <= 0:
+        return []
+    found = find_encoder_layer_stack(encoder)
+    if found is None:
+        msg = (
+            f"Could not locate a transformer-block ModuleList on "
+            f"{type(encoder).__name__}; cannot unfreeze its top {top_n} layers. "
+            f"Children: {[n for n, _ in encoder.named_children()]}"
+        )
+        raise ValueError(msg)
+    path, stack = found
+    depth = len(stack)
+    if top_n > depth:
+        msg = f"encoder_trainable_top_layers={top_n} exceeds the encoder's {depth} blocks"
+        raise ValueError(msg)
+    unfrozen: list[str] = []
+    for idx in range(depth - top_n, depth):
+        for name, param in stack[idx].named_parameters():
+            param.requires_grad_(True)
+            unfrozen.append(f"{path}.{idx}.{name}")
+    # Post-stack projections, identified positionally: direct children that are
+    # not the stack itself and not the pre-stack input projection.
+    for name, child in encoder.named_children() if include_post_projections else ():
+        if name == path.split(".")[0] or name.startswith("input"):
+            continue
+        for pname, param in child.named_parameters():
+            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 CHANGED Viewed

@@ -1,156 +1,107 @@
 import functools
 import inspect
 import json
 import logging
-import math
-from collections.abc import Iterator
 from pathlib import Path
 from threading import Thread
 import torch
 import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-from torch.nn.utils.rnn import pad_sequence
 from transformers import (
-    AutoConfig,
     AutoModel,
     AutoModelForCausalLM,
     AutoTokenizer,
     PreTrainedModel,
     TextIteratorStreamer,
 )
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
-try:
     from .asr_config import ASRConfig, compute_encoder_output_length
-    from .projectors import PROJECTOR_CLASSES
-except ImportError:
-    from asr_config import ASRConfig, compute_encoder_output_length  # type: ignore[no-redef]
-    from projectors import PROJECTOR_CLASSES  # type: ignore[no-redef]
 logger = logging.getLogger(__name__)
-# FlashAttention's kernels are compiled for head dimensions up to 256.
-FLASH_ATTENTION_MAX_HEAD_DIM = 256
 # 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
 # tensor cores on sm_75+ and is what Qwen already pads its published tables to.
 VOCAB_PAD_MULTIPLE = 128
-# 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([F.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)"
-        )
-def find_encoder_layer_stack(encoder: nn.Module) -> tuple[str, nn.Module] | None:
-    """Locate the encoder's ordered transformer-block list.
-    Returns `(attribute_path, ModuleList)` or None. Encoders disagree on the
-    attribute name -- Granite Speech uses `layers`, Whisper `encoder.layers`,
-    others `blocks` or `encoder_layers` -- so probe the known names rather
-    than hardcoding one and silently unfreezing nothing on a new encoder.
-    """
-    candidates = ("layers", "blocks", "encoder_layers", "encoder.layers", "encoder.blocks")
-    for path in candidates:
-        obj: nn.Module | None = encoder
-        for part in path.split("."):
-            obj = getattr(obj, part, None)
-            if obj is None:
-                break
-        if isinstance(obj, nn.ModuleList) and len(obj) > 0:
-            return path, obj
-    return None
-def _resolve_dtype(name, fallback: torch.dtype) -> torch.dtype:
     """Resolve a config dtype name to a `torch.dtype`, falling back when unusable.
     Only a genuine `str` naming a real floating-point dtype is honoured. Stand-in
@@ -168,206 +119,6 @@ def _resolve_dtype(name, fallback: torch.dtype) -> torch.dtype:
     )
-def unfreeze_encoder_top_layers(
-    encoder: nn.Module, top_n: int, include_post_projections: bool = True
-) -> list[str]:
-    """Unfreeze the top `top_n` encoder blocks, optionally with the output projections.
-    Assumes the caller has already frozen the whole encoder. Returns the sorted
-    names of every parameter switched back on, so callers (and tests) can
-    assert on exactly what moved rather than trusting a count.
-    The post-stack projections (Granite's `out` / `out_mid`) are included by
-    default: they sit between the top block and the projector, so leaving them
-    frozen forces the newly-trainable blocks to adapt through a fixed output
-    map. That is a reasonable prior, but on this stack it is not what the
-    gradient says. Probing the step-2000 granite_qwen_lora checkpoint on real
-    audio, `out`/`out_mid` carry 33.57M params -- 23% of a top-4 budget -- for
-    an RMS-per-param gradient of 1.14e-05 against 1.4-1.7e-04 in the top
-    blocks, i.e. **0.27% of the selection's squared gradient**. Dropping them
-    moved the aggregate norm from 1.5840 to 1.5827. Pass False to reclaim the
-    parameters; the default stays True so existing recipes are unchanged.
-    The pre-stack `input_linear` stays frozen either way -- it is the feature
-    front-end, the part most expensive to damage and least specialised to the
-    encoder's own CTC head.
-    Raises if the layer stack cannot be found or `top_n` exceeds its depth:
-    silently unfreezing nothing would produce a run that looks like a
-    partial-unfreeze experiment and is actually the frozen baseline.
-    """
-    if top_n <= 0:
-        return []
-    found = find_encoder_layer_stack(encoder)
-    if found is None:
-        raise ValueError(
-            f"Could not locate a transformer-block ModuleList on "
-            f"{type(encoder).__name__}; cannot unfreeze its top {top_n} layers. "
-            f"Children: {[n for n, _ in encoder.named_children()]}"
-        )
-    path, stack = found
-    depth = len(stack)
-    if top_n > depth:
-        raise ValueError(
-            f"encoder_trainable_top_layers={top_n} exceeds the encoder's {depth} blocks"
-        )
-    unfrozen: list[str] = []
-    for idx in range(depth - top_n, depth):
-        for name, param in stack[idx].named_parameters():
-            param.requires_grad_(True)
-            unfrozen.append(f"{path}.{idx}.{name}")
-    # Post-stack projections, identified positionally: direct children that are
-    # not the stack itself and not the pre-stack input projection.
-    for name, child in encoder.named_children() if include_post_projections else ():
-        if name == path.split(".")[0] or name.startswith("input"):
-            continue
-        for pname, param in child.named_parameters():
-            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
-def _max_attention_head_dim(text_config) -> int | None:
-    """Largest attention head dim across layers, or None if undeterminable.
-    Has to cope with heterogeneous configs: Gemma 4 varies head_dim per layer,
-    and reading `config.head_dim` on one of those raises
-    AmbiguousGlobalPerLayerAttributeError instead of returning a number, so the
-    per-layer list must be consulted explicitly.
-    """
-    per_layer = getattr(text_config, "per_layer_config", None)
-    dims: list = []
-    if per_layer:
-        dims = [getattr(layer, "head_dim", None) for layer in per_layer]
-    else:
-        try:
-            dims = [getattr(text_config, "head_dim", None)]
-        except Exception:
-            # Raised by 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
-@functools.lru_cache(maxsize=16)
-def _has_sliding_window_attention(model_id: str) -> bool:
-    """Whether `model_id`'s text stack has any sliding-window attention layer.
-    Read from the config rather than the loaded module so the answer is
-    available before `from_pretrained` picks an attn implementation.
-    Two independent spellings, because transformers has both. Gemma 4 declares
-    `sliding_window: 512` AND `layer_types: [full_attention,
-    sliding_attention]` (verified against google/gemma-4-E2B-it); Qwen2-style
-    configs carry a non-null `sliding_window` that is inert unless
-    `use_sliding_window` is true, so that flag has to veto the window.
-    Returns True when the config can't be read, so an unknown architecture
-    keeps the conservative eager path on MPS.
-    """
-    try:
-        probe = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
-        text_config = probe.get_text_config() if hasattr(probe, "get_text_config") else probe
-    except Exception:
-        logger.warning(
-            "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,
-        )
-        return True
-    layer_types = getattr(text_config, "layer_types", None) or []
-    if any("sliding" in str(layer_type) for layer_type in layer_types):
-        return True
-    if getattr(text_config, "use_sliding_window", None) is False:
-        return False
-    return bool(getattr(text_config, "sliding_window", None))
-def _resolve_attn_implementation(requested: str | None, model_id: str | None = None) -> str | None:
-    """Coerce flash_attention_2 to sdpa when CUDA isn't available, and avoid sdpa on MPS.
-    FA2 is CUDA-only. On MPS/CPU, requesting it either errors at load or
-    silently falls back to a slower path; either way the user pays the FA2
-    install + import cost for no win. Coerce here so a saved config that
-    pins flash_attention_2 still loads on Mac / CPU-only Linux boxes.
-    MPS needs eager for SOME models. PyTorch's Metal sdpa kernel returns wrong
-    results for cached single-token decode against a sliding-window mask, which
-    is exactly Gemma 4's layout (`sliding_window=512`, four sliding layers per
-    full-attention layer). A full-sequence forward is fine, so the damage shows
-    up only during generation: greedy decode with no cache transcribed
-    correctly while the identical cached decode produced "Mr. and a".
-    The mask is the trigger, so the coercion is scoped to models that actually
-    build one. `model_id` opts into that check; without it every MPS load falls
-    back to eager, which is what this function used to do unconditionally. That
-    blanket version taxed the models it was never protecting: Qwen3.5-2B
-    declares `sliding_window: None` and cannot hit the bug, and eager cost it
-    41% of decode throughput on an M4 Max (19.9 -> 28.1 tok/s, bf16, 64 tokens
-    from a 200-token prompt). Audio encoders are likewise unaffected -- they run
-    full-sequence forwards with no cache at all, and the Granite branch below
-    has been pinning sdpa on MPS for every granite_qwen run to date.
-    SECOND, SEPARATE MPS sdpa HAZARD, and the reason `generate` carries a guard:
-    Metal's sdpa returns NaN for a row containing any fully-masked position, so
-    a LEFT-PADDED batch comes back as NaN logits, argmax 0, and a transcript of
-    "!!!!!!". One pad token is enough. Measured on bare transformers with no
-    tiny_audio code involved (Qwen3.5-2B, bf16): the identical forward is exact
-    to 0.000 on CPU/fp32, and eager is correct on MPS at any padding. Batch 1
-    pads nothing, which is why the eval harness is safe under sdpa -- but this
-    is what the blanket eager was accidentally also protecting, so anything
-    that batches ragged audio on a Mac must ask for eager explicitly. See
-    `_assert_sdpa_safe_on_mps`.
-    """
-    if (
-        torch.backends.mps.is_available()
-        and requested in (None, "sdpa", "flash_attention_2")
-        and (model_id is None or _has_sliding_window_attention(model_id))
-    ):
-        return "eager"
-    # Otherwise fall through: on MPS, sdpa and None (let transformers pick,
-    # which is sdpa where available) are both safe for a model with no
-    # sliding-window layer, and a config pinning flash_attention_2 is still
-    # downgraded to sdpa below.
-    if requested != "flash_attention_2":
-        return requested
-    if not torch.cuda.is_available():
-        return "sdpa"
-    # CUDA alone isn't enough -- the flash_attn package must actually be
-    # importable, and new enough for transformers. It is a source build that
-    # routinely fails on RunPod images (its metadata hook imports torch, so any
-    # broken torch install takes it down with it). Without this check a CUDA
-    # box with no flash-attn raises at from_pretrained instead of quietly
-    # using sdpa, which is the same numerics at lower throughput.
-    from transformers.utils import is_flash_attn_2_available
-    if not is_flash_attn_2_available():
-        logger.warning(
-            "flash_attention_2 requested but flash_attn is not installed or too old; "
-            "falling back to sdpa."
-        )
-        return "sdpa"
-    return requested
 def _gather_audio_embeds(
     audio_embeds: torch.Tensor,
     token_counts: torch.Tensor,
@@ -387,7 +138,7 @@ def _gather_audio_embeds(
     _, max_len, _ = audio_embeds.shape
     needed = int(token_counts.max().item()) if max_tokens is None else max_tokens
     if needed > max_len:
-        audio_embeds = F.pad(audio_embeds, (0, 0, 0, needed - max_len))
         max_len = needed
     indices = torch.arange(max_len, device=audio_embeds.device).unsqueeze(0)
     mask = indices < token_counts.unsqueeze(1)
@@ -397,7 +148,7 @@ def _gather_audio_embeds(
 def _assert_audio_token_counts(
     audio_embeds: torch.Tensor,
     token_counts: torch.Tensor,
-    projector,
     encoder_valid_lengths: torch.Tensor | None = None,
     max_tokens: int | None = None,
 ) -> None:
@@ -424,12 +175,13 @@ def _assert_audio_token_counts(
         actual = torch.as_tensor(actual, device=token_counts.device).to(torch.long).reshape(-1)
         if actual.shape == token_counts.shape and not torch.equal(actual, token_counts):
             rows = (actual != token_counts).nonzero().flatten()[:8].tolist()
-            raise ValueError(
                 "Audio token count mismatch between prompt and projector. Rows "
                 f"{rows}: prompt expects {token_counts[rows].tolist()}, encoder+projector "
                 f"produced {actual[rows].tolist()}. A wrong `encoder_conv_layers` for this "
                 "encoder is the usual cause."
             )
     available = audio_embeds.shape[1]
     if max_tokens is not None:
@@ -437,14 +189,15 @@ def _assert_audio_token_counts(
     else:
         needed = int(token_counts.max().item()) if token_counts.numel() else 0
     if needed > available:
-        raise ValueError(
             f"Projector produced {available} audio frames but the prompt expects up to "
             f"{needed}. Without this check `_gather_audio_embeds` would zero-pad the "
             "deficit, silently feeding the decoder zero vectors in place of audio."
         )
-def _patch_gemma_decode_loop(model) -> None:
     """Make a Gemma4ForCausalLM drop `per_layer_inputs` after the first step.
     `Gemma4ForConditionalGeneration` does this itself; the causal LM does not.
@@ -469,7 +222,9 @@ def _patch_gemma_decode_loop(model) -> None:
     # 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, is_first_iteration: bool = False, **kwargs):
         """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:
@@ -479,7 +234,7 @@ def _patch_gemma_decode_loop(model) -> None:
     model.prepare_inputs_for_generation = prepare_inputs_for_generation
-def _assert_projector_loaded(incompatible_keys, projector_type: str) -> None:
     """Fail loudly when a checkpoint's projector doesn't match the built one.
     `save_pretrained` serializes `projector.*` (plus the language model when
@@ -499,7 +254,7 @@ def _assert_projector_loaded(incompatible_keys, projector_type: str) -> None:
     if not missing and not unexpected:
         return
-    detail = []
     if missing:
         detail.append(f"missing from the checkpoint: {missing}")
     if unexpected:
@@ -539,8 +294,6 @@ def _log_linear_attention_backends(level: int = logging.WARNING) -> None:
     touch these functions at all. `level` is INFO for an inference load
     (`from_pretrained`), where the slower path is not something to act on.
     """
-    import importlib.util
     found = {
         pkg: importlib.util.find_spec(mod) is not None
         for pkg, mod in (("causal-conv1d", "causal_conv1d"), ("flash-linear-attention", "fla"))
@@ -571,10 +324,16 @@ class ASRModel(PreTrainedModel, GenerationMixin):
     TRANSCRIBE_PROMPT = "Transcribe the speech to text"
     @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> "ASRModel":
         """Load model from pretrained, handling device placement correctly."""
-        from safetensors.torch import load_file
-        from transformers.utils.hub import cached_file
         config = kwargs.pop("config", None)
         if config is None:
@@ -589,7 +348,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # Load projector weights from safetensors
             subfolder = kwargs.get("subfolder")
             revision = kwargs.get("revision")
-            cache_kwargs = {}
             if subfolder:
                 cache_kwargs["subfolder"] = subfolder
             if revision:
@@ -618,21 +377,23 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 )
                 if adapter_config_file is not None:
                     # Load saved adapter weights using the original repo_id/path
-                    # PEFT handles Hub downloads and caching internally
-                    from peft import PeftModel
                     # Adapter weights load onto the base model's device. Left to
                     # PEFT they go straight to "cuda" whenever CUDA reports
                     # available, which ZeroGPU Spaces do at startup without a
                     # GPU attached ("No CUDA GPUs are available").
                     base_device = next(model.language_model.parameters()).device
-                    model.language_model = PeftModel.from_pretrained(
                         model.language_model,
                         pretrained_model_name_or_path,
                         is_trainable=True,
                         torch_device=str(base_device),
                         **cache_kwargs,
                     )
                 else:
                     # No saved adapters - initialize fresh LLM LoRA for training.
                     # __init__ skips _setup_lora while loading, so call it here.
@@ -642,7 +403,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         finally:
             cls._is_loading_from_pretrained = False
-    def __init__(self, config: ASRConfig, **kwargs) -> None:
         """Build encoder, projector, decoder and tokenizer from `config`.
         `**kwargs` are the loader arguments `from_pretrained` forwards (e.g.
@@ -652,15 +413,18 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Shadows the class attribute when the config names one, so a run that
         # trained under a specific instruction decodes under the same one.
-        if getattr(config, "transcribe_prompt", None) is not None:
-            self.TRANSCRIBE_PROMPT = config.transcribe_prompt
         target_dtype = getattr(torch, config.model_dtype)
         # Audio encoder (frozen)
         self.audio_tower = self._load_audio_encoder(config, target_dtype)
-        # Language model (frozen)
-        self.language_model = self._load_language_model(config, target_dtype)
         # Does the decoder's forward take `skip_logits`? Only liger's patched
         # `lce_forward` declares it; the stock transformers forward does not.
@@ -706,51 +470,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # this line 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
-        # Set EOS tokens, filtering out any that don't exist in the tokenizer.
-        # `convert_tokens_to_ids` reports "not in vocab" inconsistently: Qwen-
-        # style tokenizers return None, while Gemma's returns unk_token_id.
-        # Filtering on None alone left Gemma with eos_token_id=[unk, unk], so
-        # generation never stopped and every sample ran to max_new_tokens.
-        #
-        # The membership test is a name round-trip, NOT `id == unk_token_id`:
-        # in GPT-2-lineage tokenizers (SmolLM2, Qwen) unk_token *is*
-        # "<|endoftext|>", so comparing ids would discard a legitimate stop
-        # token here while fixing Gemma. A token that fell back to unk comes
-        # back under unk's name; one that is really in the vocab comes back
-        # as itself.
-        #
-        # The name list alone is not enough, and silently produced garbage on
-        # Gemma 4: that template dropped Gemma 2/3's "<end_of_turn>" for a new
-        # "<|turn>role\n ... <turn|>\n" scheme, so all three probes missed and
-        # eos collapsed to the bare "<eos>" -- a token the template never
-        # emits. Training taught the model to close its turn with "<turn|>"
-        # (id 106) and generation then ignored it, running every sample to
-        # max_new_tokens and burying the transcript under 250 copies of
-        # "<turn|>" that `skip_special_tokens` silently swallowed.
-        #
-        # So derive the real terminator from the template instead of guessing
-        # its name: render an assistant turn holding a sentinel and take the
-        # first token that follows it. That token is the turn closer by
-        # construction, for any decoder, including ones released after this
-        # code was written. The name probes stay as a fallback for tokenizers
-        # whose template renders no trailing special token.
-        eos_ids: list[int] = []
-        derived = self._derive_turn_end_token_id()
-        if derived is not None:
-            eos_ids.append(derived)
-        for token in ("<|im_end|>", "<|endoftext|>", "<end_of_turn>"):
-            token_id = self.tokenizer.convert_tokens_to_ids(token)
-            if token_id is None or self.tokenizer.convert_ids_to_tokens(token_id) != token:
-                continue
-            if token_id not in eos_ids:
-                eos_ids.append(token_id)
-        # The chat-template stop tokens above are additions, not replacements:
-        # the tokenizer's own EOS is always a valid stop and is the only one
-        # left for a decoder using none of those three templates.
-        tokenizer_eos = self.tokenizer.eos_token_id
-        if tokenizer_eos is not None and tokenizer_eos not in eos_ids:
-            eos_ids.append(tokenizer_eos)
-        self.generation_config.eos_token_id = eos_ids
         self.generation_config.pad_token_id = self.tokenizer.pad_token_id
         # Feature extractor for audio preprocessing
@@ -784,17 +504,17 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # For model parallelism
         self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
-    def _create_feature_extractor(self, config: ASRConfig):
         """Create the appropriate feature extractor for the audio encoder."""
-        from transformers import AutoFeatureExtractor
-        feature_extractor = 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
         return feature_extractor
     @classmethod
@@ -807,17 +527,16 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         from its pretrained features). See `ASRTrainer.create_optimizer` for the
         LR routing.
         """
-        encoder_kwargs = {
-            "attn_implementation": _resolve_attn_implementation(
                 config.attn_implementation, config.audio_model_id
             ),
             "low_cpu_mem_usage": True,
             "dtype": dtype,
         }
         if "whisper" in config.audio_model_id.lower():
-            from transformers import WhisperModel
             full_model = WhisperModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
             encoder = full_model.encoder
             del full_model
@@ -832,15 +551,21 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # checkpoint ships modeling .py files but its config.json has no
             # `auto_map`, so Auto* classes can't resolve them on older
             # versions -- they fail with KeyError('granite_speech5_ctc').
-            from transformers import GraniteSpeech5Encoder
             # Granite's block-attention implementation has no FA2 kernel and
             # raises ValueError on attn_implementation="flash_attention_2".
-            # _resolve_attn_implementation only downgrades FA2 when CUDA is
             # absent, so on a CUDA box the inherited default would hard-fail
             # at load. Pin sdpa regardless of what the config asks for.
-            granite_kwargs = {**encoder_kwargs, "attn_implementation": "sdpa"}
-            encoder = GraniteSpeech5Encoder.from_pretrained(config.audio_model_id, **granite_kwargs)
         elif "glm" in config.audio_model_id.lower():
             # GLM-ASR stores its encoder at audio_tower (GlmAsrEncoder), but
             # which object owns that attribute depends on the transformers
@@ -851,21 +576,21 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # flat checkpoints hung them off the top-level model. Resolve the
             # owner instead of assuming, so neither layout AttributeErrors
             # at load.
-            from transformers import AutoModelForSeq2SeqLM
-            from transformers import __version__ as transformers_version
-            full_model = AutoModelForSeq2SeqLM.from_pretrained(
                 config.audio_model_id, trust_remote_code=True, **encoder_kwargs
             )
-            inner = getattr(full_model, "model", None)
-            holder = inner if hasattr(inner, "audio_tower") else full_model
             if not hasattr(holder, "audio_tower"):
-                raise AttributeError(
-                    f"{type(full_model).__name__} exposes no audio_tower at "
                     "`.audio_tower` or `.model.audio_tower`; GLM-ASR encoder "
                     f"extraction needs updating for transformers "
                     f"{transformers_version}."
                 )
             encoder = holder.audio_tower
             # Drop the LLM decoder and projector to free their VRAM. These must
             # go through the same owner as audio_tower: assigning None to a name
@@ -873,7 +598,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # a plain attribute and frees nothing.
             holder.language_model = None
             holder.multi_modal_projector = None
-            del full_model
         else:
             encoder = AutoModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
@@ -886,7 +611,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # `encoder_dtype` overrides the stack dtype when the encoder has
         # trainable blocks; it must be fp32 for Adam's step to survive
         # rounding. See ASRConfig.encoder_dtype for the arithmetic.
-        encoder = encoder.to(dtype=_resolve_dtype(getattr(config, "encoder_dtype", None), dtype))
         if getattr(config, "freeze_audio_encoder", True):
             encoder.requires_grad_(False)
             encoder.train(False)  # equivalent to .eval(); avoids a security hook false-positive
@@ -928,34 +655,12 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         return encoder
     @classmethod
-    def _load_language_model(cls, config: ASRConfig, dtype: torch.dtype) -> PreTrainedModel:
         """Load and freeze the language model."""
-        attn_implementation = _resolve_attn_implementation(
-            config.attn_implementation, config.text_model_id
-        )
-        # FlashAttention only has kernels for head_dim <= 256, and it fails at
-        # the first forward rather than at load. Gemma 4 E2B trips this: 7 of
-        # its 35 layers use head_dim=512 (the other 28 use 256), so FA2 can
-        # never run this architecture regardless of flash-attn version. Detect
-        # it from the config and fall back instead of dying a step into
-        # training.
-        if attn_implementation == "flash_attention_2":
-            probe = AutoConfig.from_pretrained(config.text_model_id, trust_remote_code=True)
-            text_probe = probe.get_text_config() if hasattr(probe, "get_text_config") else probe
-            head_dim = _max_attention_head_dim(text_probe)
-            if head_dim is not None and head_dim > FLASH_ATTENTION_MAX_HEAD_DIM:
-                logger.warning(
-                    "%s has max head_dim=%d, above FlashAttention's limit of %d; "
-                    "using sdpa for the decoder.",
-                    config.text_model_id,
-                    head_dim,
-                    FLASH_ATTENTION_MAX_HEAD_DIM,
-                )
-                attn_implementation = "sdpa"
-        decoder_kwargs = {
-            "attn_implementation": attn_implementation,
             "trust_remote_code": True,
             "low_cpu_mem_usage": True,
             "dtype": dtype,
@@ -983,7 +688,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 logging.INFO if cls._is_loading_from_pretrained else logging.WARNING
             )
-        decoder = AutoModelForCausalLM.from_pretrained(config.text_model_id, **decoder_kwargs)
         # Gemma 4 checkpoints are natively multimodal: AutoModelForCausalLM maps
         # `gemma4` to Gemma4ForConditionalGeneration, which is both wasteful
@@ -1005,7 +713,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # 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.
-        decoder = decoder.to(dtype=dtype)
         decoder.config.use_cache = getattr(config, "use_cache", True)
         if getattr(config, "freeze_language_model", True):
             decoder.requires_grad_(False)
@@ -1013,7 +723,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         return decoder
     @staticmethod
-    def _gemma_text_only(loaded: PreTrainedModel) -> PreTrainedModel:
         """Turn a loaded Gemma4ForConditionalGeneration into a text-only CausalLM.
         Two independent problems force this transplant rather than just loading
@@ -1042,14 +752,17 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         returned text-only model holds 4,628,569,344 parameters against the
         checkpoint's 5,123,178,979, so ~494M are released.
         """
-        from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM
-        text_config = loaded.config.get_text_config()
         with torch.device("meta"):
             shell = Gemma4ForCausalLM(text_config)
-        shell.model = loaded.model.language_model
-        shell.lm_head = loaded.lm_head
         shell.generation_config = loaded.generation_config
         shell.tie_weights()
@@ -1064,18 +777,22 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Patched on the instance rather than via a subclass: subclassing
         # Gemma4ForCausalLM trips transformers' experts-implementation check
         # ("does not support setting experts implementation") during __init__.
-        _patch_gemma_decode_loop(shell)
         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:
-            raise RuntimeError(
                 "Gemma text-only conversion left parameters on the meta device "
                 f"(would silently produce garbage): {stranded[:5]}"
             )
-        return shell
-    def _create_projector(self, config: ASRConfig, dtype: torch.dtype) -> nn.Module:
         """Create the trainable audio projector."""
         # Auto-detect dimensions if not specified
         if config.encoder_dim is None:
@@ -1084,7 +801,8 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 enc_cfg, "d_model", None
             )
             if config.encoder_dim is None:
-                raise ValueError("Could not auto-detect encoder_dim. Please specify in config.")
         if config.llm_dim is None:
             dec_cfg = self.language_model.config
@@ -1096,16 +814,18 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 dec_cfg, "d_model", None
             )
             if config.llm_dim is None:
-                raise ValueError("Could not auto-detect llm_dim. Please specify in config.")
         # 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:
-            raise ValueError(
                 f"Unknown projector_type: {projector_type}. "
                 f"Valid options: {list(PROJECTOR_CLASSES.keys())}"
             )
         projector = projector_class(config)
         # Move projector to same device as language model (important when using quantization)
@@ -1117,10 +837,8 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self._projector_dtype = proj_dtype
         return projector.to(device=device, dtype=proj_dtype)
-    def _setup_lora(self, config: ASRConfig):
         """Apply LoRA adapters to the language model for Stage 2 fine-tuning."""
-        from peft import LoraConfig, get_peft_model
         lora_config = LoraConfig(
             r=config.lora_rank,
             lora_alpha=config.lora_alpha,
@@ -1136,11 +854,17 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             bias="none",
             task_type="CAUSAL_LM",
         )
-        self.language_model = get_peft_model(self.language_model, lora_config)
-    def _init_tokenizer(self, config: ASRConfig):
         """Initialize tokenizer with audio token."""
-        self.tokenizer = AutoTokenizer.from_pretrained(config.text_model_id, trust_remote_code=True)
         # Set pad token. Prefer a dedicated pad token if the tokenizer has one
         # (e.g. Qwen's <|finetune_right_pad_id|>); otherwise fall back to
@@ -1164,7 +888,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         already_in_vocab = existing_id is not None and existing_id != self.tokenizer.unk_token_id
         if not already_in_vocab:
-            existing_special = getattr(self.tokenizer, "additional_special_tokens", None) or []
             self.tokenizer.add_special_tokens(
                 {"additional_special_tokens": [*existing_special, self.audio_token]}
             )
@@ -1203,7 +929,12 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 mean_resizing=True,
             )
-        self.audio_token_id = self.tokenizer.convert_tokens_to_ids(self.audio_token)
         self.tokenizer.padding_side = "right"
         # Sync token IDs to configs
@@ -1213,6 +944,58 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 cfg.eos_token_id = self.tokenizer.eos_token_id
                 cfg.bos_token_id = self.tokenizer.bos_token_id
     def _derive_turn_end_token_id(self) -> int | None:
         """Return the token id the chat template uses to close an assistant turn.
@@ -1243,9 +1026,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             return None
         if not isinstance(rendered, str) or sentinel not in rendered:
             return None
-        tail_ids = self.tokenizer(rendered.split(sentinel)[-1], add_special_tokens=False)[
-            "input_ids"
-        ]
         if not tail_ids:
             return None
         # Require a special token: a template that ends the turn with plain
@@ -1254,7 +1037,7 @@ 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, **kwargs):
         """Repair MPS-unsafe embeddings whenever the model lands on MPS.
         Device placement happens after construction -- `pipeline` builds the
@@ -1268,7 +1051,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         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 = super()._apply(*args, **kwargs)
         try:
             on_mps = any(p.device.type == "mps" for p in module.parameters())
         except StopIteration:  # pragma: no cover - parameterless model
@@ -1281,7 +1064,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             )
         return module
-    def train(self, mode: bool = True):
         """Set train/eval mode, but keep frozen submodules out of train mode.
         HF Trainer calls `model.train()` at the top of every training step, which
@@ -1325,7 +1108,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             # "Norm"). This mirrors what the partial-unfreeze path already gets
             # for free via the `freeze_audio_encoder: true` branch above.
             for module in self.audio_tower.modules():
-                if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
                     module.eval()
         if getattr(self.config, "freeze_language_model", True):
             self.language_model.train(False)
@@ -1333,10 +1119,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
     def _set_gradient_checkpointing(
         self,
         enable: bool = True,
-        gradient_checkpointing_func=None,
         every_n_layers: int = 1,
-        **kwargs,
-    ):
         """Enable/disable gradient checkpointing on the trainable submodules.
         Routes the request to whichever components are actually trainable in
@@ -1359,9 +1145,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         Upstream sniffs for it to detect the pre-4.35 checkpointing format and
         would silently take the legacy `self.apply(...)` path instead.
         """
-        import inspect
-        forwardable = dict(kwargs)
         forwardable["every_n_layers"] = every_n_layers
         if gradient_checkpointing_func is not None:
             forwardable["gradient_checkpointing_func"] = gradient_checkpointing_func
@@ -1370,16 +1154,16 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Gradient checkpointing trades compute for memory by recomputing activations.
         for submodule in self._gradient_checkpointing_targets():
             setter = getattr(submodule, "_set_gradient_checkpointing", None)
             if setter is not None:
                 accepted = inspect.signature(setter).parameters
                 passthrough = {k: v for k, v in forwardable.items() if k in accepted}
                 setter(enable=enable, **passthrough)
-            elif hasattr(submodule, "gradient_checkpointing_enable") and enable:
-                submodule.gradient_checkpointing_enable(
-                    gradient_checkpointing_kwargs={"use_reentrant": False}
-                )
-            elif hasattr(submodule, "gradient_checkpointing_disable") and not enable:
-                submodule.gradient_checkpointing_disable()
     def _gradient_checkpointing_targets(self) -> list[nn.Module]:
         """Return the submodules that should respond to gradient_checkpointing
@@ -1401,21 +1185,16 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         """Replace the decoder's token embedding module."""
         self.language_model.set_input_embeddings(value)
-    def get_output_embeddings(self) -> nn.Module:
-        """Return the decoder's LM head."""
         return self.language_model.get_output_embeddings()
-    def set_output_embeddings(self, value: nn.Module) -> None:
         """Replace the decoder's LM head."""
-        self.language_model.set_output_embeddings(value)
-    def get_processor(self):
         """Get the processor for this model."""
-        try:
-            from .asr_processing import ASRProcessor
-        except ImportError:
-            from asr_processing import ASRProcessor  # type: ignore[no-redef]
         return ASRProcessor(
             feature_extractor=self.feature_extractor,
             tokenizer=self.tokenizer,
@@ -1424,7 +1203,21 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             audio_token=self.audio_token,
         )
-    def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
         """Save trainable weights: projector, plus the encoder/LM when unfrozen.
         Every module this returns is gated on the same freeze flag that decides
@@ -1445,7 +1238,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         `freeze_text_embed_tokens` freezes an embedding that must still
         round-trip with the rest of the LM.
         """
-        sd = {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
         # Keyed on autograd state, not `freeze_audio_encoder`. Under the
         # documented partial-unfreeze recipe that flag stays True while the top
         # N blocks train, so gating on it meant a trained encoder was never
@@ -1458,14 +1253,14 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             sd.update(
                 {
                     f"audio_tower.{k}": v
-                    for k, v in self.audio_tower.state_dict().items()
                     if fully or k in enc_trainable
                 }
             )
         if not getattr(self.config, "freeze_language_model", True):
             lm = self.language_model
             if hasattr(lm, "peft_config"):
-                for name, v in lm.state_dict().items():
                     if "lora_" in name:
                         continue
                     k = name.removeprefix("base_model.model.")
@@ -1473,10 +1268,18 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                     k = k.replace(".base_layer.", ".")
                     sd[f"language_model.{k}"] = v
             else:
-                sd.update({f"language_model.{k}": v for k, v in lm.state_dict().items()})
-        return sd
-    def _ple_text_model(self):
         """Locate the submodule owning `get_per_layer_inputs`, or None.
         Only Gemma 4 style decoders have one. The attribute sits on
@@ -1490,10 +1293,11 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
         for candidate in candidates:
             if candidate is not None and hasattr(candidate, "get_per_layer_inputs"):
-                return candidate
         return None
-    def _per_layer_kwargs(self, input_ids: torch.Tensor | None) -> dict:
         """Precompute Gemma 4 per-layer embeddings (PLE) from clean `input_ids`.
         Gemma 4 builds a token-identity PLE component by looking `input_ids` up
@@ -1519,6 +1323,11 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             return {}
         return {"per_layer_inputs": text_model.get_per_layer_inputs(input_ids, None)}
     def _compute_encoder_output_lengths(
         self,
         audio_attention_mask: torch.Tensor,
@@ -1626,7 +1435,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         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.
@@ -1673,12 +1482,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         valid_lengths = None
         if attention_mask is not None:
             if attention_mask.shape[-1] != sequence_length:
-                raise ValueError(
                     "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."
                 )
             valid_lengths = attention_mask.to(device).sum(dim=-1)
         if getattr(config, "mask_time_prob", 0.0) > 0:
@@ -1747,7 +1557,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             (mask_prob * lengths.float() / mask_length + 0.5).long(), min=min_masks
         )  # (B,)
-        # Sample start positions independently per sample × 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.
@@ -1780,11 +1590,11 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         use_cache: bool | None = None,
         cache_position: torch.Tensor | None = None,
         audio_token_counts: torch.Tensor | None = None,
-        **kwargs,
     ) -> CausalLMOutputWithPast:
         """Forward pass for training and inference."""
         if inputs_embeds is None:
-            inputs_embeds = self.language_model.get_input_embeddings()(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.
@@ -1869,7 +1679,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         if labels is not None and self._lm_accepts_skip_logits:
             kwargs.setdefault("skip_logits", True)
-        return self.language_model(
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
@@ -1880,8 +1690,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             **ple_kwargs,
             **kwargs,
         )
-    def prepare_inputs_for_generation(self, *args, **kwargs):
         """Prepare inputs for generation, handling audio features for cached decoding."""
         input_features = kwargs.pop("input_features", None)
         cache_position = kwargs.get("cache_position")
@@ -1926,7 +1737,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             return_tensors="pt",
             enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
         )
-        ids = chat_result.input_ids
         return (ids[0] if ids.dim() > 1 else ids).to(torch.long)
     def _left_pad_prompt_rows(
@@ -1934,23 +1748,12 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         rows: list[torch.Tensor],
         device: torch.device,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Stack per-sample prompt rows into a left-padded batch.
-        Left, not right: these feed `generate`, and right padding would sit
-        between the prompt and the first generated token. Pad positions never
-        collide with `audio_token_id`, so the masked_scatter below is unaffected.
         """
-        pad_id = self.tokenizer.pad_token_id
-        if pad_id is None:
-            pad_id = self.tokenizer.eos_token_id or 0
-        input_ids = pad_sequence(
-            rows, batch_first=True, padding_value=int(pad_id), padding_side="left"
-        )
-        # The mask is padded from ones rather than derived from `input_ids !=
-        # pad_id`: a real token may equal `pad_id` when pad falls back to eos.
-        attention_mask = pad_sequence(
-            [torch.ones_like(row) for row in rows], batch_first=True, padding_side="left"
-        )
         return input_ids.to(device), attention_mask.to(device)
     def _prepare_audio_inputs(
@@ -1991,7 +1794,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             input_ids, attention_mask = self._left_pad_prompt_rows(rows, device)
         # Get text embeddings and replace audio tokens with audio embeddings
-        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
         audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
         inputs_embeds = inputs_embeds.masked_scatter(
             audio_token_mask.to(inputs_embeds.device),
@@ -2012,16 +1815,19 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         Batch 1 and uniform-length batches pad nothing and pass. `eager` is
         correct on MPS at any padding, and the bug is Metal-specific, so a CUDA
-        box is unaffected either way. `_resolve_attn_implementation` documents
         the measurements.
         """
         if attention_mask is None or attention_mask.device.type != "mps":
             return
-        if self.language_model.config._attn_implementation != "sdpa":
             return
         if not bool((attention_mask == 0).any()):
             return
-        raise ValueError(
             "Refusing to generate: this is a left-padded batch of "
             f"{attention_mask.shape[0]} on MPS with sdpa attention, where "
             "Metal's kernel returns NaN for the padded rows and every one of "
@@ -2030,26 +1836,40 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             "to decode), batch only equal-length clips, or generate one clip "
             "at a time."
         )
-    @torch.no_grad()
     def generate(
         self,
         input_ids: torch.Tensor | None = None,
         input_features: torch.Tensor | None = None,
         audio_attention_mask: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
-        **generate_kwargs,
-    ):
         """Generate transcription from audio input.
         Can be called in two ways:
         1. With input_ids containing <audio> tokens (from processor)
         2. With just audio, and we build the prompt internally
         """
         if input_features is None:
-            raise ValueError("input_features required for generation")
         if audio_attention_mask is None:
-            raise ValueError("audio_attention_mask required for generation")
         input_ids, attention_mask, inputs_embeds = self._prepare_audio_inputs(
             input_features,
@@ -2064,36 +1884,22 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Pull any score-related flags out of generate_kwargs and apply them to
         # a derived generation_config so they actually take effect.
         gen_cfg = self.generation_config
-        score_flags = {}
         for flag in ("output_scores", "output_logits", "return_dict_in_generate"):
             if flag in generate_kwargs:
                 score_flags[flag] = generate_kwargs.pop(flag)
         if score_flags:
-            from copy import copy as _copy
-            gen_cfg = _copy(self.generation_config)
             for flag, value in score_flags.items():
                 setattr(gen_cfg, flag, value)
             # output_scores requires return_dict_in_generate for HF generate to
             # actually populate .scores on the output object.
-            if gen_cfg.output_scores and not gen_cfg.return_dict_in_generate:
                 gen_cfg.return_dict_in_generate = True
-        # Generate using language model.
-        #
-        # Default path passes both input_ids and inputs_embeds so
-        # repetition_penalty works correctly (it needs 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().
-        ple_kwargs = self._per_layer_kwargs(input_ids)
-        lm_inputs = (
-            {"inputs_embeds": inputs_embeds, **ple_kwargs}
-            if ple_kwargs
-            else {"input_ids": input_ids, "inputs_embeds": inputs_embeds}
-        )
         output = self.language_model.generate(
             **lm_inputs,
             attention_mask=attention_mask,
@@ -2110,15 +1916,35 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # compatibility with existing callers.
         prompt_len = input_ids.shape[1] if "input_ids" in lm_inputs else 0
         if isinstance(output, torch.Tensor):
-            return output[:, prompt_len:]
-        output.sequences = output.sequences[:, prompt_len:]
         return output
     def generate_streaming(
         self,
         input_features: torch.Tensor,
         audio_attention_mask: torch.Tensor,
-        **generate_kwargs,
     ) -> Iterator[str]:
         """Generate transcription with streaming token output.
@@ -2145,13 +1971,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
         # Prepare generation kwargs
-        gen_kwargs = {
-            "inputs_embeds": inputs_embeds,
             "attention_mask": attention_mask,
             "generation_config": self.generation_config,
             "streamer": streamer,
-            # Gemma 4 needs PLE precomputed from the prompt ids; {} elsewhere.
-            **self._per_layer_kwargs(input_ids),
             **generate_kwargs,
         }
@@ -2192,10 +2018,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         thread.join()
-    def save_pretrained(self, save_directory: str | Path, **kwargs) -> None:
-        """Save model, tokenizer, and processor."""
-        import shutil
         save_dir = Path(save_directory)
         save_dir.mkdir(parents=True, exist_ok=True)
@@ -2212,15 +2041,23 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             self.config.vocab_size = vocab_size
             self.config.text_config.vocab_size = vocab_size
-        if hasattr(self.audio_tower.config, "num_mel_bins"):
-            self.config.audio_config.num_mel_bins = self.audio_tower.config.num_mel_bins
         # Save model (temporarily remove non-serializable attributes)
         tokenizer = self.tokenizer
         del self.tokenizer
         try:
-            super().save_pretrained(save_dir, **kwargs)
         finally:
             self.tokenizer = tokenizer
@@ -2259,6 +2096,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Add processor auto_map to preprocessor_config.json
         config_path = save_dir / "preprocessor_config.json"
         if config_path.exists():
             with config_path.open() as f:
                 processor_config = json.load(f)
@@ -2308,7 +2146,22 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         if card_fn is not None:
             card_fn(str(output_dir))
-    def push_to_hub(self, repo_id: str, **kwargs) -> str:
         """Push model to HuggingFace Hub, ensuring adapter_config points to repo.
         IMPORTANT: Sets base_model_name_or_path in adapter_config.json to repo_id
@@ -2318,7 +2171,18 @@ 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
-        return super().push_to_hub(repo_id, **kwargs)
 # Register with transformers Auto classes

 import functools
+import importlib.util
 import inspect
 import json
 import logging
+import shutil
+from collections.abc import Callable, Iterator
+from copy import copy
+from os import PathLike
 from pathlib import Path
 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 safetensors.torch import load_file
+from torch.nn import functional
 from transformers import (
+    AutoFeatureExtractor,
     AutoModel,
     AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
     AutoTokenizer,
+    BatchEncoding,
     PreTrainedModel,
+    PreTrainedTokenizerBase,
+    SequenceFeatureExtractor,
     TextIteratorStreamer,
+    WhisperModel,
 )
+from transformers import __version__ as transformers_version
+from transformers.generation.utils import GenerateOutput, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils.hub import (
+    cached_file,
+)
+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_processing import ASRProcessor, left_pad_prompt_rows
+    from .asr_types import (
+        DecoderLoadKwargs,
+        GenerativeDecoder,
+        HubFileKwargs,
+        LoadKwargs,
+        LoadStateDictResult,
+        OutputLengthProjector,
+        StateDictT,
+    )
+    from .projectors import PROJECTOR_CLASSES, MLPAudioProjector
+    # 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_processing import ASRProcessor, left_pad_prompt_rows
+        from .asr_types import LoadStateDictResult, OutputLengthProjector, StateDictT
+        from .projectors import PROJECTOR_CLASSES
+    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_processing import ASRProcessor, left_pad_prompt_rows
+        from asr_types import LoadStateDictResult, OutputLengthProjector, StateDictT
+        from projectors import PROJECTOR_CLASSES
+    # 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.
+    try:
+        from transformers import GraniteSpeech5Encoder
+    except ImportError:  # transformers < 5.16: no native granite_speech5
+        GraniteSpeech5Encoder = None
+    try:
+        from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+        from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM
+    except ImportError:  # transformers without the gemma4 architecture
+        Gemma4TextConfig = None
+        Gemma4ForCausalLM = None
 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
 # tensor cores on sm_75+ and is what Qwen already pads its published tables to.
 VOCAB_PAD_MULTIPLE = 128
+def _resolve_dtype(name: object, fallback: torch.dtype) -> torch.dtype:
     """Resolve a config dtype name to a `torch.dtype`, falling back when unusable.
     Only a genuine `str` naming a real floating-point dtype is honoured. Stand-in
     )
 def _gather_audio_embeds(
     audio_embeds: torch.Tensor,
     token_counts: torch.Tensor,
     _, max_len, _ = audio_embeds.shape
     needed = int(token_counts.max().item()) if max_tokens is None else max_tokens
     if needed > max_len:
+        audio_embeds = functional.pad(audio_embeds, (0, 0, 0, needed - max_len))
         max_len = needed
     indices = torch.arange(max_len, device=audio_embeds.device).unsqueeze(0)
     mask = indices < token_counts.unsqueeze(1)
 def _assert_audio_token_counts(
     audio_embeds: torch.Tensor,
     token_counts: torch.Tensor,
+    projector: OutputLengthProjector,
     encoder_valid_lengths: torch.Tensor | None = None,
     max_tokens: int | None = None,
 ) -> None:
         actual = torch.as_tensor(actual, device=token_counts.device).to(torch.long).reshape(-1)
         if actual.shape == token_counts.shape and not torch.equal(actual, token_counts):
             rows = (actual != token_counts).nonzero().flatten()[:8].tolist()
+            msg = (
                 "Audio token count mismatch between prompt and projector. Rows "
                 f"{rows}: prompt expects {token_counts[rows].tolist()}, encoder+projector "
                 f"produced {actual[rows].tolist()}. A wrong `encoder_conv_layers` for this "
                 "encoder is the usual cause."
             )
+            raise ValueError(msg)
     available = audio_embeds.shape[1]
     if max_tokens is not None:
     else:
         needed = int(token_counts.max().item()) if token_counts.numel() else 0
     if needed > available:
+        msg = (
             f"Projector produced {available} audio frames but the prompt expects up to "
             f"{needed}. Without this check `_gather_audio_embeds` would zero-pad the "
             "deficit, silently feeding the decoder zero vectors in place of audio."
         )
+        raise ValueError(msg)
+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.
     # 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.prepare_inputs_for_generation = prepare_inputs_for_generation
+def _assert_projector_loaded(incompatible_keys: LoadStateDictResult, projector_type: str) -> None:
     """Fail loudly when a checkpoint's projector doesn't match the built one.
     `save_pretrained` serializes `projector.*` (plus the language model when
     if not missing and not unexpected:
         return
+    detail: list[str] = []
     if missing:
         detail.append(f"missing from the checkpoint: {missing}")
     if unexpected:
     touch these functions at all. `level` is INFO for an inference load
     (`from_pretrained`), where the slower path is not something to act on.
     """
     found = {
         pkg: importlib.util.find_spec(mod) is not None
         for pkg, mod in (("causal-conv1d", "causal_conv1d"), ("flash-linear-attention", "fla"))
     TRANSCRIBE_PROMPT = "Transcribe the speech to text"
     @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | PathLike[str] | None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> "ASRModel":
         """Load model from pretrained, handling device placement correctly."""
+        if pretrained_model_name_or_path is None:
+            msg = "ASRModel.from_pretrained needs a checkpoint directory or Hub repo id"
+            raise ValueError(msg)
         config = kwargs.pop("config", None)
         if config is None:
             # Load projector weights from safetensors
             subfolder = kwargs.get("subfolder")
             revision = kwargs.get("revision")
+            cache_kwargs: HubFileKwargs = {}
             if subfolder:
                 cache_kwargs["subfolder"] = subfolder
             if revision:
                 )
                 if adapter_config_file is not None:
                     # Load saved adapter weights using the original repo_id/path
+                    # PEFT handles Hub downloads and caching internally.
+                    #
                     # Adapter weights load onto the base model's device. Left to
                     # PEFT they go straight to "cuda" whenever CUDA reports
                     # available, which ZeroGPU Spaces do at startup without a
                     # GPU attached ("No CUDA GPUs are available").
                     base_device = next(model.language_model.parameters()).device
+                    peft = PeftModel.from_pretrained(
                         model.language_model,
                         pretrained_model_name_or_path,
                         is_trainable=True,
                         torch_device=str(base_device),
                         **cache_kwargs,
                     )
+                    # See `language_model` in __init__: PEFT forwards the
+                    # decoder interface, so callers keep using it unchanged.
+                    model.language_model = cast("GenerativeDecoder", peft)
                 else:
                     # No saved adapters - initialize fresh LLM LoRA for training.
                     # __init__ skips _setup_lora while loading, so call it here.
         finally:
             cls._is_loading_from_pretrained = False
+    def __init__(self, config: ASRConfig, **kwargs: Any) -> None:
         """Build encoder, projector, decoder and tokenizer from `config`.
         `**kwargs` are the loader arguments `from_pretrained` forwards (e.g.
         # Shadows the class attribute when the config names one, so a run that
         # trained under a specific instruction decodes under the same one.
+        prompt = getattr(config, "transcribe_prompt", None)
+        if prompt is not None:
+            self.TRANSCRIBE_PROMPT = prompt  # pyright: ignore[reportConstantRedefinition]
         target_dtype = getattr(torch, config.model_dtype)
         # Audio encoder (frozen)
         self.audio_tower = 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
+        # decoder it wraps, so the decoder's interface is the right static view.
+        self.language_model: GenerativeDecoder = self._load_language_model(config, target_dtype)
         # Does the decoder's forward take `skip_logits`? Only liger's patched
         # `lce_forward` declares it; the stock transformers forward does not.
         # this line 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
+        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
         # For model parallelism
         self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
+    def _create_feature_extractor(self, config: ASRConfig) -> SequenceFeatureExtractor:
         """Create the appropriate feature extractor for the audio encoder."""
+        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
         from its pretrained features). See `ASRTrainer.create_optimizer` for the
         LR routing.
         """
+        encoder_kwargs: LoadKwargs = {
+            "attn_implementation": resolve_attn_implementation(
                 config.attn_implementation, config.audio_model_id
             ),
             "low_cpu_mem_usage": True,
             "dtype": dtype,
         }
+        encoder: nn.Module
         if "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
             # checkpoint ships modeling .py files but its config.json has no
             # `auto_map`, so Auto* classes can't resolve them on older
             # versions -- they fail with KeyError('granite_speech5_ctc').
+            if GraniteSpeech5Encoder is None:
+                msg = (
+                    "Granite Speech 5.0 encoders need transformers >= 5.16 for the "
+                    f"native granite_speech5 architecture; found {transformers_version}."
+                )
+                raise ImportError(msg)
             # Granite's block-attention implementation has no FA2 kernel and
             # raises ValueError on attn_implementation="flash_attention_2".
+            # resolve_attn_implementation only downgrades FA2 when CUDA is
             # absent, so on a CUDA box the inherited default would hard-fail
             # at load. Pin sdpa regardless of what the config asks for.
+            granite_kwargs: LoadKwargs = {**encoder_kwargs, "attn_implementation": "sdpa"}
+            granite = GraniteSpeech5Encoder
+            encoder = granite.from_pretrained(config.audio_model_id, **granite_kwargs)
         elif "glm" in config.audio_model_id.lower():
             # GLM-ASR stores its encoder at audio_tower (GlmAsrEncoder), but
             # which object owns that attribute depends on the transformers
             # flat checkpoints hung them off the top-level model. Resolve the
             # owner instead of assuming, so neither layout AttributeErrors
             # at load.
+            glm_model = AutoModelForSeq2SeqLM.from_pretrained(
                 config.audio_model_id, trust_remote_code=True, **encoder_kwargs
             )
+            inner = getattr(glm_model, "model", None)
+            holder: Any = (
+                inner if inner is not None and hasattr(inner, "audio_tower") else glm_model
+            )
             if not hasattr(holder, "audio_tower"):
+                msg = (
+                    f"{type(glm_model).__name__} exposes no audio_tower at "
                     "`.audio_tower` or `.model.audio_tower`; GLM-ASR encoder "
                     f"extraction needs updating for transformers "
                     f"{transformers_version}."
                 )
+                raise AttributeError(msg)
             encoder = holder.audio_tower
             # Drop the LLM decoder and projector to free their VRAM. These must
             # go through the same owner as audio_tower: assigning None to a name
             # a plain attribute and frees nothing.
             holder.language_model = None
             holder.multi_modal_projector = None
+            del glm_model
         else:
             encoder = AutoModel.from_pretrained(config.audio_model_id, **encoder_kwargs)
         # `encoder_dtype` overrides the stack dtype when the encoder has
         # 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]
         if getattr(config, "freeze_audio_encoder", True):
             encoder.requires_grad_(False)
             encoder.train(False)  # equivalent to .eval(); avoids a security hook false-positive
         return encoder
     @classmethod
+    def _load_language_model(cls, config: ASRConfig, dtype: torch.dtype) -> "GenerativeDecoder":
         """Load and freeze the language model."""
+        decoder_kwargs: DecoderLoadKwargs = {
+            "attn_implementation": resolve_decoder_attn_implementation(
+                config.attn_implementation, config.text_model_id
+            ),
             "trust_remote_code": True,
             "low_cpu_mem_usage": True,
             "dtype": dtype,
                 logging.INFO if cls._is_loading_from_pretrained else logging.WARNING
             )
+        decoder = cast(
+            "GenerativeDecoder",
+            AutoModelForCausalLM.from_pretrained(config.text_model_id, **decoder_kwargs),
+        )
         # Gemma 4 checkpoints are natively multimodal: AutoModelForCausalLM maps
         # `gemma4` to Gemma4ForConditionalGeneration, which is both wasteful
         # 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]
         decoder.config.use_cache = getattr(config, "use_cache", True)
         if getattr(config, "freeze_language_model", True):
             decoder.requires_grad_(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
         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()
         # Patched on the instance rather than via a subclass: subclassing
         # Gemma4ForCausalLM trips transformers' experts-implementation check
         # ("does not support setting experts implementation") during __init__.
+        # 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
         if config.encoder_dim is None:
                 enc_cfg, "d_model", None
             )
             if config.encoder_dim is None:
+                msg = "Could not auto-detect encoder_dim. Please specify in config."
+                raise ValueError(msg)
         if config.llm_dim is None:
             dec_cfg = self.language_model.config
                 dec_cfg, "d_model", None
             )
             if config.llm_dim is None:
+                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)
         # Move projector to same device as language model (important when using quantization)
         self._projector_dtype = proj_dtype
         return projector.to(device=device, dtype=proj_dtype)
+    def _setup_lora(self, config: ASRConfig) -> None:
         """Apply LoRA adapters to the language model for Stage 2 fine-tuning."""
         lora_config = LoraConfig(
             r=config.lora_rank,
             lora_alpha=config.lora_alpha,
             bias="none",
             task_type="CAUSAL_LM",
         )
+        # See `language_model` in __init__ for why the PEFT wrapper is
+        # typed as the decoder it wraps.
+        self.language_model = cast(
+            "GenerativeDecoder", get_peft_model(self.language_model, lora_config)
+        )
+    def _init_tokenizer(self, config: ASRConfig) -> None:
         """Initialize tokenizer with audio token."""
+        self.tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(
+            config.text_model_id, trust_remote_code=True
+        )
         # Set pad token. Prefer a dedicated pad token if the tokenizer has one
         # (e.g. Qwen's <|finetune_right_pad_id|>); otherwise fall back to
         already_in_vocab = existing_id is not None and existing_id != self.tokenizer.unk_token_id
         if not already_in_vocab:
+            existing_special: list[str] = (
+                getattr(self.tokenizer, "additional_special_tokens", None) or []
+            )
             self.tokenizer.add_special_tokens(
                 {"additional_special_tokens": [*existing_special, self.audio_token]}
             )
                 mean_resizing=True,
             )
+        audio_token_id = self.tokenizer.convert_tokens_to_ids(self.audio_token)
+        # The token is in the vocab by now, so a single name maps to one id.
+        if not isinstance(audio_token_id, int):
+            msg = f"Tokenizer has no id for the audio token {self.audio_token!r}"
+            raise TypeError(msg)
+        self.audio_token_id = audio_token_id
         self.tokenizer.padding_side = "right"
         # Sync token IDs to configs
                 cfg.eos_token_id = self.tokenizer.eos_token_id
                 cfg.bos_token_id = self.tokenizer.bos_token_id
+    def _resolve_eos_token_ids(self) -> list[int]:
+        """Collect the ids generation stops on: turn closer, known end tokens, tokenizer EOS."""
+        # Candidates that don't exist in the tokenizer are filtered out.
+        # `convert_tokens_to_ids` reports "not in vocab" inconsistently: Qwen-
+        # style tokenizers return None, while Gemma's returns unk_token_id.
+        # Filtering on None alone left Gemma with eos_token_id=[unk, unk], so
+        # generation never stopped and every sample ran to max_new_tokens.
+        #
+        # The membership test is a name round-trip, NOT `id == unk_token_id`:
+        # in GPT-2-lineage tokenizers (SmolLM2, Qwen) unk_token *is*
+        # "<|endoftext|>", so comparing ids would discard a legitimate stop
+        # token here while fixing Gemma. A token that fell back to unk comes
+        # back under unk's name; one that is really in the vocab comes back
+        # as itself.
+        #
+        # The name list alone is not enough, and silently produced garbage on
+        # Gemma 4: that template dropped Gemma 2/3's "<end_of_turn>" for a new
+        # "<|turn>role\n ... <turn|>\n" scheme, so all three probes missed and
+        # eos collapsed to the bare "<eos>" -- a token the template never
+        # emits. Training taught the model to close its turn with "<turn|>"
+        # (id 106) and generation then ignored it, running every sample to
+        # max_new_tokens and burying the transcript under 250 copies of
+        # "<turn|>" that `skip_special_tokens` silently swallowed.
+        #
+        # So derive the real terminator from the template instead of guessing
+        # its name: render an assistant turn holding a sentinel and take the
+        # first token that follows it. That token is the turn closer by
+        # construction, for any decoder, including ones released after this
+        # code was written. The name probes stay as a fallback for tokenizers
+        # whose template renders no trailing special token.
+        eos_ids: list[int] = []
+        derived = self._derive_turn_end_token_id()
+        if derived is not None:
+            eos_ids.append(derived)
+        for token in ("<|im_end|>", "<|endoftext|>", "<end_of_turn>"):
+            token_id = self.tokenizer.convert_tokens_to_ids(token)
+            # A single token name maps to a single id (None when unknown).
+            if (
+                not isinstance(token_id, int)
+                or self.tokenizer.convert_ids_to_tokens(token_id) != token
+            ):
+                continue
+            if token_id not in eos_ids:
+                eos_ids.append(token_id)
+        # The chat-template stop tokens above are additions, not replacements:
+        # the tokenizer's own EOS is always a valid stop and is the only one
+        # left for a decoder using none of those three templates.
+        tokenizer_eos = self.tokenizer.eos_token_id
+        if isinstance(tokenizer_eos, int) and tokenizer_eos not in eos_ids:
+            eos_ids.append(tokenizer_eos)
+        return eos_ids
     def _derive_turn_end_token_id(self) -> int | None:
         """Return the token id the chat template uses to close an assistant turn.
             return None
         if not isinstance(rendered, str) or sentinel not in rendered:
             return None
+        tail_ids: list[int] = self.tokenizer(
+            rendered.split(sentinel)[-1], add_special_tokens=False
+        )["input_ids"]
         if not tail_ids:
             return None
         # Require a special token: a template that ends the turn with plain
         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
         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
+    def train(self, mode: bool = True) -> Self:
         """Set train/eval mode, but keep frozen submodules out of train mode.
         HF Trainer calls `model.train()` at the top of every training step, which
             # "Norm"). This mirrors what the partial-unfreeze path already gets
             # for free via the `freeze_audio_encoder: true` branch above.
             for module in self.audio_tower.modules():
+                if isinstance(
+                    module,
+                    (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm),
+                ):
                     module.eval()
         if getattr(self.config, "freeze_language_model", True):
             self.language_model.train(False)
     def _set_gradient_checkpointing(
         self,
         enable: bool = True,
+        gradient_checkpointing_func: Callable[..., Any] | None = None,
         every_n_layers: int = 1,
+        **kwargs: Any,
+    ) -> None:
         """Enable/disable gradient checkpointing on the trainable submodules.
         Routes the request to whichever components are actually trainable in
         Upstream sniffs for it to detect the pre-4.35 checkpointing format and
         would silently take the legacy `self.apply(...)` path instead.
         """
+        forwardable: dict[str, Any] = dict(kwargs)
         forwardable["every_n_layers"] = every_n_layers
         if gradient_checkpointing_func is not None:
             forwardable["gradient_checkpointing_func"] = gradient_checkpointing_func
         # Gradient checkpointing trades compute for memory by recomputing activations.
         for submodule in self._gradient_checkpointing_targets():
             setter = getattr(submodule, "_set_gradient_checkpointing", None)
+            enable_fn = getattr(submodule, "gradient_checkpointing_enable", None)
+            disable_fn = getattr(submodule, "gradient_checkpointing_disable", None)
             if setter is not None:
                 accepted = inspect.signature(setter).parameters
                 passthrough = {k: v for k, v in forwardable.items() if k in accepted}
                 setter(enable=enable, **passthrough)
+            elif enable_fn is not None and enable:
+                enable_fn(gradient_checkpointing_kwargs={"use_reentrant": False})
+            elif disable_fn is not None and not enable:
+                disable_fn()
     def _gradient_checkpointing_targets(self) -> list[nn.Module]:
         """Return the submodules that should respond to gradient_checkpointing
         """Replace the decoder's token embedding module."""
         self.language_model.set_input_embeddings(value)
+    def get_output_embeddings(self) -> nn.Module | None:
+        """Return the decoder's LM head (None for a decoder without one)."""
         return self.language_model.get_output_embeddings()
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
         """Replace the decoder's LM head."""
+        self.language_model.set_output_embeddings(new_embeddings)
+    def get_processor(self) -> ASRProcessor:
         """Get the processor for this model."""
         return ASRProcessor(
             feature_extractor=self.feature_extractor,
             tokenizer=self.tokenizer,
             audio_token=self.audio_token,
         )
+    @overload
+    def state_dict(
+        self, *, destination: StateDictT, prefix: str = ..., keep_vars: bool = ...
+    ) -> StateDictT: ...
+    @overload
+    def state_dict(self, *, prefix: str = ..., keep_vars: bool = ...) -> dict[str, Any]: ...
+    def state_dict(
+        self,
+        *,
+        destination: dict[str, Any] | None = None,
+        prefix: str = "",
+        keep_vars: bool = False,
+    ) -> dict[str, Any]:
         """Save trainable weights: projector, plus the encoder/LM when unfrozen.
         Every module this returns is gated on the same freeze flag that decides
         `freeze_text_embed_tokens` freezes an embedding that must still
         round-trip with the rest of the LM.
         """
+        sd = {
+            f"projector.{k}": v for k, v in self.projector.state_dict(keep_vars=keep_vars).items()
+        }
         # Keyed on autograd state, not `freeze_audio_encoder`. Under the
         # documented partial-unfreeze recipe that flag stays True while the top
         # N blocks train, so gating on it meant a trained encoder was never
             sd.update(
                 {
                     f"audio_tower.{k}": v
+                    for k, v in self.audio_tower.state_dict(keep_vars=keep_vars).items()
                     if fully or k in enc_trainable
                 }
             )
         if not getattr(self.config, "freeze_language_model", True):
             lm = self.language_model
             if hasattr(lm, "peft_config"):
+                for name, v in lm.state_dict(keep_vars=keep_vars).items():
                     if "lora_" in name:
                         continue
                     k = name.removeprefix("base_model.model.")
                     k = k.replace(".base_layer.", ".")
                     sd[f"language_model.{k}"] = v
             else:
+                sd.update(
+                    {
+                        f"language_model.{k}": v
+                        for k, v in lm.state_dict(keep_vars=keep_vars).items()
+                    }
+                )
+        if destination is None:
+            destination = {}
+        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
         )
         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
             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)
+        return embeds
     def _compute_encoder_output_lengths(
         self,
         audio_attention_mask: torch.Tensor,
         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.
         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_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.
         use_cache: bool | None = None,
         cache_position: torch.Tensor | None = None,
         audio_token_counts: torch.Tensor | None = None,
+        **kwargs: Any,
     ) -> CausalLMOutputWithPast:
         """Forward pass for training and inference."""
         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.
         if labels is not None and self._lm_accepts_skip_logits:
             kwargs.setdefault("skip_logits", True)
+        outputs: CausalLMOutputWithPast = self.language_model(
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
             **ple_kwargs,
             **kwargs,
         )
+        return outputs
+    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)
         cache_position = kwargs.get("cache_position")
             return_tensors="pt",
             enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
         )
+        if not isinstance(chat_result, BatchEncoding):
+            msg = f"apply_chat_template returned {type(chat_result).__name__}, not BatchEncoding"
+            raise TypeError(msg)
+        ids = chat_result["input_ids"]
         return (ids[0] if ids.dim() > 1 else ids).to(torch.long)
     def _left_pad_prompt_rows(
         rows: list[torch.Tensor],
         device: torch.device,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Stack per-sample prompt rows into a left-padded batch on `device`.
+        See `left_pad_prompt_rows`; pad positions never collide with
+        `audio_token_id`, so the masked_scatter below is unaffected.
         """
+        input_ids, attention_mask = left_pad_prompt_rows(rows, self.tokenizer)
         return input_ids.to(device), attention_mask.to(device)
     def _prepare_audio_inputs(
             input_ids, attention_mask = self._left_pad_prompt_rows(rows, device)
         # Get text embeddings and replace audio tokens with audio embeddings
+        inputs_embeds = self._embed_tokens(input_ids)
         audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
         inputs_embeds = inputs_embeds.masked_scatter(
             audio_token_mask.to(inputs_embeds.device),
         Batch 1 and uniform-length batches pad nothing and pass. `eager` is
         correct on MPS at any padding, and the bug is Metal-specific, so a CUDA
+        box is unaffected either way. `resolve_attn_implementation` documents
         the measurements.
         """
         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"
+        ):
             return
         if not bool((attention_mask == 0).any()):
             return
+        msg = (
             "Refusing to generate: this is a left-padded batch of "
             f"{attention_mask.shape[0]} on MPS with sdpa attention, where "
             "Metal's kernel returns NaN for the padded rows and every one of "
             "to decode), batch only equal-length clips, or generate one clip "
             "at a time."
         )
+        raise ValueError(msg)
+    @torch.no_grad()  # pyright: ignore[reportUntypedFunctionDecorator]
     def generate(
         self,
+        *args: Any,
         input_ids: torch.Tensor | None = None,
         input_features: torch.Tensor | None = None,
         audio_attention_mask: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
+        **generate_kwargs: Any,
+    ) -> GenerateOutput | torch.LongTensor:
         """Generate transcription from audio input.
         Can be called in two ways:
         1. With input_ids containing <audio> tokens (from processor)
         2. With just audio, and we build the prompt internally
+        The audio arguments are keyword-only; `*args` exists only so the
+        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)
         input_ids, attention_mask, inputs_embeds = self._prepare_audio_inputs(
             input_features,
         # Pull any score-related flags out of generate_kwargs and apply them to
         # a derived generation_config so they actually take effect.
         gen_cfg = self.generation_config
+        score_flags: dict[str, Any] = {}
         for flag in ("output_scores", "output_logits", "return_dict_in_generate"):
             if flag in generate_kwargs:
                 score_flags[flag] = generate_kwargs.pop(flag)
         if score_flags:
+            gen_cfg = copy(self.generation_config)
             for flag, value in score_flags.items():
                 setattr(gen_cfg, flag, value)
             # output_scores requires return_dict_in_generate for HF generate to
             # actually populate .scores on the output object.
+            if getattr(gen_cfg, "output_scores", None) and not getattr(
+                gen_cfg, "return_dict_in_generate", None
+            ):
                 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,
         # 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().
+        """
+        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(
         self,
         input_features: torch.Tensor,
         audio_attention_mask: torch.Tensor,
+        **generate_kwargs: Any,
     ) -> Iterator[str]:
         """Generate transcription with streaming token output.
         )
         # Prepare generation kwargs
+        # Same prompt arguments as generate(); with input_ids present the
+        # streamer's skip_prompt drops the echoed prompt tokens.
+        gen_kwargs: dict[str, Any] = {
+            **self._lm_generate_inputs(input_ids, inputs_embeds),
             "attention_mask": attention_mask,
             "generation_config": self.generation_config,
             "streamer": streamer,
             **generate_kwargs,
         }
         thread.join()
+    def save_pretrained(
+        self, save_directory: str | PathLike[str], *args: Any, **kwargs: Any
+    ) -> None:
+        """Save model, tokenizer, and processor.
+        Extra positional arguments are forwarded to `PreTrainedModel.save_pretrained`.
+        """
         save_dir = Path(save_directory)
         save_dir.mkdir(parents=True, exist_ok=True)
             self.config.vocab_size = vocab_size
             self.config.text_config.vocab_size = vocab_size
+        encoder_config = self.audio_tower.config
+        if hasattr(encoder_config, "num_mel_bins"):
+            # The default is never used: hasattr just confirmed the attribute.
+            # Read through getattr because nn.Module types `.config` as
+            # Tensor | Module, which hasattr does not narrow.
+            self.config.audio_config.num_mel_bins = getattr(encoder_config, "num_mel_bins", None)
         # Save model (temporarily remove non-serializable attributes)
         tokenizer = self.tokenizer
         del self.tokenizer
         try:
+            super().save_pretrained(
+                save_dir,
+                *args,
+                **kwargs,
+            )
         finally:
             self.tokenizer = tokenizer
         # Add processor auto_map to preprocessor_config.json
         config_path = save_dir / "preprocessor_config.json"
+        processor_config: dict[str, Any]
         if config_path.exists():
             with config_path.open() as f:
                 processor_config = json.load(f)
         if card_fn is not None:
             card_fn(str(output_dir))
+    # typeshed's functools._Wrapped (how PreTrainedModel declares push_to_hub) is
+    # not a descriptor, so no method can match it; the signature is the
+    # PushToHubMixin one that PreTrainedModel wraps.
+    def push_to_hub(  # pyright: ignore[reportIncompatibleVariableOverride]
+        self,
+        repo_id: str,
+        *,
+        commit_message: str | None = None,
+        commit_description: str | None = None,
+        private: bool | None = None,
+        token: bool | str | None = None,
+        revision: str | None = None,
+        create_pr: bool = False,
+        max_shard_size: int | str | None = "50GB",
+        tags: list[str] | None = None,
+    ) -> str:
         """Push model to HuggingFace Hub, ensuring adapter_config points to repo.
         IMPORTANT: Sets base_model_name_or_path in adapter_config.json to repo_id
         # Store repo_id in config so save_pretrained can access it
         self.config.pretrained_model_path = repo_id
         # Call parent's push_to_hub
+        return PreTrainedModel.push_to_hub(
+            self,
+            repo_id,
+            commit_message=commit_message,
+            commit_description=commit_description,
+            private=private,
+            token=token,
+            revision=revision,
+            create_pr=create_pr,
+            max_shard_size=max_shard_size,
+            tags=tags,
+        )
 # Register with transformers Auto classes

asr_pipeline.py CHANGED Viewed

@@ -1,32 +1,41 @@
 """ASR pipeline for audio-to-text transcription with optional timestamps and diarization."""
 import re
 from pathlib import Path
-from typing import Any, cast
 import numpy as np
 import torch
 import transformers
 from transformers.pipelines.audio_utils import ffmpeg_read
-try:
-    from .alignment import ForcedAligner, QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_processing import prepend_lead_in
-    from .diarization import NemotronDiarizer, SpeakerDiarizer
-except ImportError:
-    from alignment import ForcedAligner, QwenForcedAligner  # type: ignore[no-redef]
-    from asr_modeling import ASRModel  # type: ignore[no-redef]
-    from asr_processing import prepend_lead_in  # type: ignore[no-redef]
-    from diarization import NemotronDiarizer, SpeakerDiarizer  # type: ignore[no-redef]
 # Re-export for backwards compatibility
 __all__ = [
     "ASRPipeline",
-    "ForcedAligner",
     "NemotronDiarizer",
     "QwenForcedAligner",
-    "SpeakerDiarizer",
 ]
 # Audio is transcribed in chunks cut at the quietest point between
@@ -37,11 +46,15 @@ CHUNK_MIN_S = 8.0
 def chunk_bounds(
-    audio: np.ndarray, sample_rate: int, max_s: float = CHUNK_MAX_S, min_s: float = CHUNK_MIN_S
 ) -> list[tuple[int, int]]:
     """Sample ranges of at most `max_s`, each cut at the quietest 100 ms frame after `min_s`."""
     frame = int(0.1 * sample_rate)
-    bounds, start, n = [], 0, len(audio)
     while n - start > max_s * sample_rate:
         lo = start + int(min_s * sample_rate)
         hi = start + int(max_s * sample_rate)
@@ -62,7 +75,7 @@ def chunk_bounds(
 SILENCE_RMS = 1e-5
-def is_silent(audio: np.ndarray) -> bool:
     """True for digital silence (or an empty array): nothing for the model to hear."""
     return audio.size == 0 or float(np.sqrt(np.mean(np.square(audio)))) < SILENCE_RMS
@@ -73,32 +86,61 @@ _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)
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
     """ASR Pipeline for audio-to-text transcription."""
     model: ASRModel
-    def __init__(self, model: ASRModel, **kwargs):
         """Initialize ASR pipeline.
         Args:
             model: ASRModel instance for transcription
             **kwargs: Additional arguments (feature_extractor, tokenizer, device)
         """
-        feature_extractor = kwargs.pop("feature_extractor", None)
-        tokenizer = kwargs.pop("tokenizer", model.tokenizer)
         if feature_extractor is None:
             feature_extractor = model.get_processor().feature_extractor
         super().__init__(
-            model=model, feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs
         )
-    def _sanitize_parameters(self, **kwargs):
         """Intercept our custom parameters before parent class validates them."""
-        # Remove our custom parameters so parent doesn't see them
-        kwargs.pop("return_timestamps", None)
         kwargs.pop("return_speakers", None)
         kwargs.pop("num_speakers", None)
         kwargs.pop("min_speakers", None)
@@ -106,13 +148,25 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         kwargs.pop("hf_token", None)
         kwargs.pop("user_prompt", None)
-        return super()._sanitize_parameters(**kwargs)
     def __call__(
         self,
-        inputs,
-        **kwargs,
-    ):
         """Transcribe audio with optional word-level timestamps and speaker diarization.
         Args:
@@ -128,16 +182,25 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             Dict with 'text' key, 'words' key if return_timestamps=True,
             and speaker labels on words plus 'speaker_segments' if return_speakers=True
         """
         # Extract our params before super().__call__ (which will also call _sanitize_parameters)
-        return_timestamps = kwargs.pop("return_timestamps", False)
-        return_speakers = kwargs.pop("return_speakers", False)
-        user_prompt = kwargs.pop("user_prompt", None)
         if kwargs.pop("min_speakers", None) is not None:
-            raise ValueError(
                 "min_speakers is not supported: Nemotron-3-Diarization decides how many "
                 "speakers it hears. Pass num_speakers (exact) or max_speakers instead."
             )
-        diarization_params = {
             "num_speakers": kwargs.pop("num_speakers", None),
             "max_speakers": kwargs.pop("max_speakers", None),
         }
@@ -150,7 +213,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         # exception anywhere in the call leaves the custom prompt in place for
         # every later call on this pipeline, and the eval harness swallows that
         # exception and keeps scoring against the wrong prompt.
-        original_prompt = None
         if user_prompt:
             original_prompt = self.model.TRANSCRIBE_PROMPT
             self.model.TRANSCRIBE_PROMPT = user_prompt
@@ -168,7 +231,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             if original_prompt is not None:
                 self.model.TRANSCRIBE_PROMPT = original_prompt
-    def _transcribe_plain(self, inputs, **kwargs):
         """Transcribe in the same chunks as `_transcribe_timed`, without timing words.
         One call over a whole recording returns only its first
@@ -187,7 +252,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         if len(bounds) == 1:
             return self._transcribe_chunk(array, sr, **kwargs)
         result: dict[str, Any] = {}
-        texts = []
         for s, e in bounds:
             out = self._transcribe_chunk(array[s:e], sr, **kwargs)
             texts.append(out["text"])
@@ -198,23 +263,26 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         result["text"] = " ".join(t for t in texts if t)
         return result
-    def _transcribe_chunk(self, chunk: np.ndarray, sample_rate: int, **kwargs) -> dict:
         """One model call on one chunk; digital silence is "" without calling the model."""
         if is_silent(chunk):
             return {"text": ""}
-        # One input in, one dict out; the parent is typed for batches as well.
         return cast(
-            "dict", super().__call__({"raw": chunk, "sampling_rate": sample_rate}, **kwargs)
         )
     def _transcribe_timed(
         self,
-        inputs,
         *,
         return_speakers: bool,
-        diarization_params: dict,
-        **kwargs,
-    ):
         """Transcribe in chunks, time every word, and label speakers if requested.
         Chunks of 8-18 s cut at quiet points are transcribed one by one and each
@@ -226,21 +294,25 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         """
         audio = self._extract_audio(inputs)
         if audio is None:
-            raise ValueError(f"Cannot read audio from {type(inputs).__name__} for timestamps")
         array = np.asarray(audio["array"], dtype=np.float32)
         sr = audio.get("sampling_rate", 16000)
         bounds = chunk_bounds(array, sr)
-        texts = [self._transcribe_chunk(array[s:e], sr, **kwargs)["text"] for s, e in bounds]
         result: dict[str, Any] = {"text": " ".join(t for t in texts if t)}
         try:
             aligned = QwenForcedAligner.align_chunks(
-                [(array[s:e], text) for (s, e), text in zip(bounds, texts)], sample_rate=sr
             )
             result["words"] = [
                 {**w, "start": w["start"] + s / sr, "end": w["end"] + s / sr}
-                for (s, _), chunk_words in zip(bounds, aligned)
                 for w in chunk_words
             ]
         except Exception as e:
@@ -261,18 +333,19 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         return result
-    def _extract_audio(self, inputs) -> dict | None:
         """Extract audio array from various input formats using HF utilities."""
         if isinstance(inputs, dict):
-            if "array" in inputs:
                 return {
-                    "array": inputs["array"],
-                    "sampling_rate": inputs.get("sampling_rate", 16000),
                 }
-            if "raw" in inputs:
                 return {
-                    "array": inputs["raw"],
-                    "sampling_rate": inputs.get("sampling_rate", 16000),
                 }
         elif isinstance(inputs, str):
             # File path - load audio using ffmpeg (same as HF pipeline)
@@ -287,21 +360,25 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         return None
-    def preprocess(self, inputs, **preprocess_params):
         """Preprocess audio inputs for the model.
         Args:
-            inputs: Audio input (dict with array, file path, etc.)
             **preprocess_params: Additional preprocessing parameters
         Yields:
             Model input dicts with input_features and attention_mask
         """
         # Handle dict with "array" key (from datasets)
         if isinstance(inputs, dict) and "array" in inputs:
             inputs = {
-                "raw": inputs["array"],
-                "sampling_rate": inputs.get("sampling_rate", self.feature_extractor.sampling_rate),
             }
         # Inference-only lead-in silence. See prepend_lead_in for the
@@ -319,23 +396,26 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             else 0.0
         )
         if lead_in > 0 and isinstance(inputs, dict) and "raw" in inputs:
-            inputs = dict(inputs)
-            inputs["raw"] = prepend_lead_in(
-                inputs["raw"],
-                inputs.get("sampling_rate", self.feature_extractor.sampling_rate),
                 lead_in,
             )
         for item in super().preprocess(inputs, **preprocess_params):
             if "is_last" not in item:
                 item["is_last"] = True
             yield item
-    def _forward(self, model_inputs, **generate_kwargs) -> dict[str, Any]:
         """Run model forward pass to generate transcription.
         Args:
-            model_inputs: Dict with input_features and attention_mask
             **generate_kwargs: Generation parameters. Pass ``output_scores=True``
                 (and ``return_dict_in_generate=True``, which is then implied) to
                 also return per-step top-1 and top-2 log-probabilities — used by
@@ -347,6 +427,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             ``top1_logprob`` / ``top2_logprob`` tensors when scores were
             requested.
         """
         # Extract audio features and is_last flag
         is_last = model_inputs.pop("is_last", True) if isinstance(model_inputs, dict) else True
@@ -383,11 +466,12 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         # annotated with sample 0's confidence and nothing downstream could
         # tell. Fail loudly instead of returning a plausible wrong number.
         if sequences.shape[0] > 1:
-            raise ValueError(
                 f"ASRPipeline received a batch of {sequences.shape[0]} but only "
                 "returns one transcript; per-item confidence would be wrong. "
                 "Call it with one clip per invocation."
             )
         top1_logprobs: list[float] = []
         top2_logprobs: list[float] = []
         if scores:
@@ -397,7 +481,8 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 torch.topk(torch.log_softmax(step_logits[0].float(), dim=-1), k=2).values
                 for step_logits in scores
             ]
-            for top1, top2 in torch.stack(per_step).tolist():
                 top1_logprobs.append(top1)
                 top2_logprobs.append(top2)
         return {
@@ -407,7 +492,14 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             "is_last": is_last,
         }
-    def postprocess(self, model_outputs, **kwargs) -> dict[str, str]:
         """Convert model output tokens to text.
         Args:
@@ -423,7 +515,13 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         tokens = model_outputs.get("tokens")
         if tokens is None:
-            return super().postprocess(model_outputs, **kwargs)
         if torch.is_tensor(tokens):
             tokens = tokens.cpu()
@@ -432,12 +530,18 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         # Filter out eos tokens that the tokenizer doesn't recognize as special
         # (generation_config.eos_token_id may differ from tokenizer.eos_token_id)
-        eos_ids = self.model.generation_config.eos_token_id
         if eos_ids is not None:
             eos_set = set(eos_ids) if isinstance(eos_ids, list) else {eos_ids}
-            tokens = [t for t in tokens.tolist() if t not in eos_set]
-        text = self.tokenizer.decode(tokens, skip_special_tokens=True).strip()
         # Strip <think>...</think> tags (Qwen3 doesn't respect /no_think prompt)
         if "<think>" in text:
             text = _THINK_TAG_RE.sub("", text).strip()
@@ -453,6 +557,22 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         return out
 def _truncate_repetitions(text: str) -> str:
     """Truncate repeated words/phrases/characters at end of text.

 """ASR pipeline for audio-to-text transcription with optional timestamps and diarization."""
 import re
+from collections.abc import Iterator, MutableMapping
 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
+if TYPE_CHECKING:
+    from .alignment import QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_processing import prepend_lead_in
+    from .diarization import NemotronDiarizer
+else:
+    try:
+        from .alignment import QwenForcedAligner
+        from .asr_modeling import ASRModel
+        from .asr_processing import prepend_lead_in
+        from .diarization import NemotronDiarizer
+    except ImportError:  # flat layout on the Hub: sibling modules, no package
+        from alignment import QwenForcedAligner
+        from asr_modeling import ASRModel
+        from asr_processing import prepend_lead_in
+        from diarization import NemotronDiarizer
 # Re-export for backwards compatibility
 __all__ = [
     "ASRPipeline",
     "NemotronDiarizer",
     "QwenForcedAligner",
 ]
 # Audio is transcribed in chunks cut at the quietest point between
 def chunk_bounds(
+    audio: npt.NDArray[np.float32],
+    sample_rate: int,
+    max_s: float = CHUNK_MAX_S,
+    min_s: float = CHUNK_MIN_S,
 ) -> list[tuple[int, int]]:
     """Sample ranges of at most `max_s`, each cut at the quietest 100 ms frame after `min_s`."""
     frame = int(0.1 * sample_rate)
+    bounds: list[tuple[int, int]] = []
+    start, n = 0, len(audio)
     while n - start > max_s * sample_rate:
         lo = start + int(min_s * sample_rate)
         hi = start + int(max_s * sample_rate)
 SILENCE_RMS = 1e-5
+def is_silent(audio: npt.NDArray[np.float32]) -> bool:
     """True for digital silence (or an empty array): nothing for the model to hear."""
     return audio.size == 0 or float(np.sqrt(np.mean(np.square(audio)))) < SILENCE_RMS
 _TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{_MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
+class _RawAudio(TypedDict):
+    """A waveform and its sampling rate, as `_extract_audio` reads them."""
+    array: npt.ArrayLike
+    sampling_rate: int
+class _DiarizationParams(TypedDict):
+    """Speaker-count hints forwarded to `NemotronDiarizer.top_speakers`."""
+    num_speakers: int | None
+    max_speakers: int | None
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
     """ASR Pipeline for audio-to-text transcription."""
     model: ASRModel
+    def __init__(self, model: ASRModel, **kwargs: Any) -> None:
         """Initialize ASR pipeline.
         Args:
             model: ASRModel instance for transcription
             **kwargs: Additional arguments (feature_extractor, tokenizer, device)
         """
+        feature_extractor: SequenceFeatureExtractor | None = kwargs.pop("feature_extractor", None)
+        tokenizer: PreTrainedTokenizerBase | None = kwargs.pop("tokenizer", model.tokenizer)
         if feature_extractor is None:
             feature_extractor = model.get_processor().feature_extractor
         super().__init__(
+            model=model,
+            feature_extractor=feature_extractor,
+            # Annotated as the slow `PreTrainedTokenizer`, but any tokenizer works.
+            tokenizer=tokenizer,  # pyright: ignore[reportArgumentType]
+            **kwargs,
         )
+    def _sanitize_parameters(
+        self,
+        chunk_length_s: float | None = None,
+        stride_length_s: float | None = None,
+        ignore_warning: bool | None = None,
+        decoder_kwargs: dict[str, Any] | None = None,
+        return_timestamps: bool | str | None = None,
+        return_language: bool | None = None,
+        **kwargs: Any,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
         """Intercept our custom parameters before parent class validates them."""
+        # Remove our custom parameters so parent doesn't see them.
+        # `return_timestamps` means word timestamps here (handled in __call__),
+        # not the parent's CTC/Whisper timestamps, so it is dropped too.
+        del return_timestamps
         kwargs.pop("return_speakers", None)
         kwargs.pop("num_speakers", None)
         kwargs.pop("min_speakers", None)
         kwargs.pop("hf_token", None)
         kwargs.pop("user_prompt", None)
+        return super()._sanitize_parameters(
+            chunk_length_s=chunk_length_s,
+            stride_length_s=stride_length_s,
+            ignore_warning=ignore_warning,
+            decoder_kwargs=decoder_kwargs,
+            return_language=return_language,
+            **kwargs,
+        )
+    # The parent annotates `list[dict]`, but one input yields one dict (here and
+    # in the parent); a batch the parent handles yields a list. Hence `Any`.
     def __call__(
         self,
+        inputs: Any,
+        *args: Any,
+        num_workers: int | None = None,
+        batch_size: int | None = None,
+        **kwargs: Any,
+    ) -> Any:
         """Transcribe audio with optional word-level timestamps and speaker diarization.
         Args:
             Dict with 'text' key, 'words' key if return_timestamps=True,
             and speaker labels on words plus 'speaker_segments' if return_speakers=True
         """
+        if args:
+            msg = f"ASRPipeline.__call__ takes 1 positional input, got {1 + len(args)}"
+            raise TypeError(msg)
+        # Not ours: handed back to the parent pipeline's __call__ untouched.
+        if num_workers is not None:
+            kwargs["num_workers"] = num_workers
+        if batch_size is not None:
+            kwargs["batch_size"] = batch_size
         # Extract our params before super().__call__ (which will also call _sanitize_parameters)
+        return_timestamps: bool = kwargs.pop("return_timestamps", False)
+        return_speakers: bool = kwargs.pop("return_speakers", False)
+        user_prompt: str | None = kwargs.pop("user_prompt", None)
         if kwargs.pop("min_speakers", None) is not None:
+            msg = (
                 "min_speakers is not supported: Nemotron-3-Diarization decides how many "
                 "speakers it hears. Pass num_speakers (exact) or max_speakers instead."
             )
+            raise ValueError(msg)
+        diarization_params: _DiarizationParams = {
             "num_speakers": kwargs.pop("num_speakers", None),
             "max_speakers": kwargs.pop("max_speakers", None),
         }
         # exception anywhere in the call leaves the custom prompt in place for
         # every later call on this pipeline, and the eval harness swallows that
         # exception and keeps scoring against the wrong prompt.
+        original_prompt: str | None = None
         if user_prompt:
             original_prompt = self.model.TRANSCRIBE_PROMPT
             self.model.TRANSCRIBE_PROMPT = user_prompt
             if original_prompt is not None:
                 self.model.TRANSCRIBE_PROMPT = original_prompt
+    def _transcribe_plain(
+        self, inputs: Any, **kwargs: Any
+    ) -> dict[str, Any] | list[dict[str, Any]]:
         """Transcribe in the same chunks as `_transcribe_timed`, without timing words.
         One call over a whole recording returns only its first
         if len(bounds) == 1:
             return self._transcribe_chunk(array, sr, **kwargs)
         result: dict[str, Any] = {}
+        texts: list[str] = []
         for s, e in bounds:
             out = self._transcribe_chunk(array[s:e], sr, **kwargs)
             texts.append(out["text"])
         result["text"] = " ".join(t for t in texts if t)
         return result
+    def _transcribe_chunk(
+        self, chunk: npt.NDArray[np.float32], sample_rate: int, **kwargs: Any
+    ) -> dict[str, Any]:
         """One model call on one chunk; digital silence is "" without calling the model."""
         if is_silent(chunk):
             return {"text": ""}
+        # One input in, one dict out; the parent is annotated as always returning a list.
         return cast(
+            "dict[str, Any]",
+            super().__call__({"raw": chunk, "sampling_rate": sample_rate}, **kwargs),
         )
     def _transcribe_timed(
         self,
+        inputs: Any,
         *,
         return_speakers: bool,
+        diarization_params: _DiarizationParams,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Transcribe in chunks, time every word, and label speakers if requested.
         Chunks of 8-18 s cut at quiet points are transcribed one by one and each
         """
         audio = self._extract_audio(inputs)
         if audio is None:
+            msg = f"Cannot read audio from {type(inputs).__name__} for timestamps"
+            raise ValueError(msg)
         array = np.asarray(audio["array"], dtype=np.float32)
         sr = audio.get("sampling_rate", 16000)
         bounds = chunk_bounds(array, sr)
+        texts: list[str] = [
+            self._transcribe_chunk(array[s:e], sr, **kwargs)["text"] for s, e in bounds
+        ]
         result: dict[str, Any] = {"text": " ".join(t for t in texts if t)}
         try:
             aligned = QwenForcedAligner.align_chunks(
+                [(array[s:e], text) for (s, e), text in zip(bounds, texts, strict=True)],
+                sample_rate=sr,
             )
             result["words"] = [
                 {**w, "start": w["start"] + s / sr, "end": w["end"] + s / sr}
+                for (s, _), chunk_words in zip(bounds, aligned, strict=True)
                 for w in chunk_words
             ]
         except Exception as e:
         return result
+    def _extract_audio(self, inputs: object) -> _RawAudio | None:
         """Extract audio array from various input formats using HF utilities."""
         if isinstance(inputs, dict):
+            sample = inputs
+            if "array" in sample:
                 return {
+                    "array": sample["array"],
+                    "sampling_rate": sample.get("sampling_rate", 16000),
                 }
+            if "raw" in sample:
                 return {
+                    "array": sample["raw"],
+                    "sampling_rate": sample.get("sampling_rate", 16000),
                 }
         elif isinstance(inputs, str):
             # File path - load audio using ffmpeg (same as HF pipeline)
         return None
+    def preprocess(self, *args: Any, **preprocess_params: Any) -> Iterator[dict[str, Any]]:
         """Preprocess audio inputs for the model.
         Args:
+            inputs: Audio input (dict with array, file path, etc.), the one
+                positional argument (see `_single_input`)
             **preprocess_params: Additional preprocessing parameters
         Yields:
             Model input dicts with input_features and attention_mask
         """
+        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),
             }
         # Inference-only lead-in silence. See prepend_lead_in for the
             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),
                 lead_in,
             )
+            inputs = padded
         for item in super().preprocess(inputs, **preprocess_params):
             if "is_last" not in item:
                 item["is_last"] = True
             yield item
+    def _forward(self, *args: Any, **generate_kwargs: Any) -> dict[str, Any]:
         """Run model forward pass to generate transcription.
         Args:
+            model_inputs: Dict with input_features and attention_mask, the one
+                positional argument (see `_single_input`)
             **generate_kwargs: Generation parameters. Pass ``output_scores=True``
                 (and ``return_dict_in_generate=True``, which is then implied) to
                 also return per-step top-1 and top-2 log-probabilities — used by
             ``top1_logprob`` / ``top2_logprob`` tensors when scores were
             requested.
         """
+        model_inputs: MutableMapping[str, Any] = _single_input(
+            args, generate_kwargs, "model_inputs"
+        )
         # Extract audio features and is_last flag
         is_last = model_inputs.pop("is_last", True) if isinstance(model_inputs, dict) else True
         # annotated with sample 0's confidence and nothing downstream could
         # tell. Fail loudly instead of returning a plausible wrong number.
         if sequences.shape[0] > 1:
+            msg = (
                 f"ASRPipeline received a batch of {sequences.shape[0]} but only "
                 "returns one transcript; per-item confidence would be wrong. "
                 "Call it with one clip per invocation."
             )
+            raise ValueError(msg)
         top1_logprobs: list[float] = []
         top2_logprobs: list[float] = []
         if scores:
                 torch.topk(torch.log_softmax(step_logits[0].float(), dim=-1), k=2).values
                 for step_logits in scores
             ]
+            steps: list[list[float]] = np.asarray(torch.stack(per_step).cpu()).tolist()
+            for top1, top2 in steps:
                 top1_logprobs.append(top1)
                 top2_logprobs.append(top2)
         return {
             "is_last": is_last,
         }
+    def postprocess(
+        self,
+        model_outputs: dict[str, Any] | list[dict[str, Any]],
+        decoder_kwargs: dict[str, Any] | None = None,
+        return_timestamps: bool | str | None = None,
+        return_language: bool | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Convert model output tokens to text.
         Args:
         tokens = model_outputs.get("tokens")
         if tokens is None:
+            return super().postprocess(
+                model_outputs,
+                decoder_kwargs=decoder_kwargs,
+                return_timestamps=return_timestamps,
+                return_language=return_language,
+                **kwargs,
+            )
         if torch.is_tensor(tokens):
             tokens = tokens.cpu()
         # Filter out eos tokens that the tokenizer doesn't recognize as special
         # (generation_config.eos_token_id may differ from tokenizer.eos_token_id)
+        eos_ids: int | list[int] | None = getattr(
+            self.model.generation_config, "eos_token_id", None
+        )
         if eos_ids is not None:
             eos_set = set(eos_ids) if isinstance(eos_ids, list) else {eos_ids}
+            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
+        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()
         return out
+def _single_input(args: tuple[Any, ...], kwargs: dict[str, Any], name: str) -> Any:
+    """The one input of a pipeline stage, given positionally or as `name=`.
+    `preprocess` and `_forward` take `*args` because their two transformers bases
+    name that input differently (`inputs`/`model_inputs` in the ASR pipeline,
+    `input_`/`input_tensors` in `Pipeline`), and no single name overrides both.
+    `Pipeline` always passes it positionally.
+    """
+    if len(args) == 1 and name not in kwargs:
+        return args[0]
+    if not args and name in kwargs:
+        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.

asr_processing.py CHANGED Viewed

@@ -1,24 +1,82 @@
 """Processor that turns raw audio (and optional text) into model inputs."""
-from typing import ClassVar, Union
 import numpy as np
 import torch
 import transformers
 from torch.nn.utils.rnn import pad_sequence
-from transformers import ProcessorMixin
-try:
-    from .asr_config import DEFAULT_ENCODER_CONV_LAYERS, ASRConfig, compute_encoder_output_length
-except ImportError:
-    from asr_config import (  # type: ignore[no-redef]
         DEFAULT_ENCODER_CONV_LAYERS,
         ASRConfig,
         compute_encoder_output_length,
     )
-def prepend_lead_in(audio, sampling_rate: int, seconds: float):
     """Prepend `seconds` of silence to a waveform (or each waveform in a list).
     Peoples ships fixed ~15s grid cuts rather than sentence-aligned segments,
@@ -38,19 +96,24 @@ def prepend_lead_in(audio, sampling_rate: int, seconds: float):
         return audio
     if isinstance(audio, (list, tuple)) and audio and not isinstance(audio[0], (int, float)):
-        return [prepend_lead_in(a, sampling_rate, seconds) for a in audio]
     pad = round(sampling_rate * seconds)
     if pad <= 0:
-        return audio
-    arr = np.asarray(audio)
-    return np.concatenate([np.zeros(pad, dtype=arr.dtype), arr])
 class ASRProcessor(ProcessorMixin):
     """Processor for Whisper-based ASR models."""
     attributes: ClassVar[list[str]] = ["feature_extractor", "tokenizer"]
     feature_extractor_class = "AutoFeatureExtractor"
     tokenizer_class = "AutoTokenizer"
     # Fallback only. The real value comes from `ASRConfig.audio_token`, which
@@ -64,10 +127,10 @@ class ASRProcessor(ProcessorMixin):
     def __init__(
         self,
-        feature_extractor,
-        tokenizer,
-        projector=None,
-        encoder_conv_layers: list | None = None,
         audio_token: str | None = None,
         lead_in_seconds: float = 0.0,
     ):
@@ -107,12 +170,16 @@ class ASRProcessor(ProcessorMixin):
         if text is not None:
             messages.append({"role": "assistant", "content": text})
-        tokenized = self.tokenizer.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=(text is None),
-            return_tensors="pt",
-            enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
         )
         # apply_chat_template returns a bare tensor or a BatchEncoding/mapping.
@@ -120,33 +187,25 @@ class ASRProcessor(ProcessorMixin):
         return (ids[0] if ids.dim() > 1 else ids).to(torch.long)
     def _stack_prompt_rows(self, rows: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
-        """Stack per-sample prompt rows into a batch, left-padding if ragged.
-        Left, not right: these feed `generate`, so padding must not sit between
-        the prompt and the first generated token. Mirrors
-        `ASRModel._left_pad_prompt_rows`; pad positions never carry
-        `audio_token_id`, so the model's masked_scatter is unaffected.
         """
-        pad_id = self.tokenizer.pad_token_id
-        if pad_id is None:
-            pad_id = self.tokenizer.eos_token_id or 0
-        input_ids = pad_sequence(
-            rows, batch_first=True, padding_value=int(pad_id), padding_side="left"
-        )
-        # Padded from ones rather than `input_ids != pad_id`: a real token may
-        # equal `pad_id` when pad falls back to eos.
-        attention_mask = pad_sequence(
-            [torch.ones_like(row) for row in rows], batch_first=True, padding_side="left"
-        )
-        return input_ids, attention_mask
-    def __call__(
         self,
-        audio: Union[list, "torch.Tensor"] | None = None,
         text: str | None = None,
         return_tensors: str = "pt",
-        **kwargs,
-    ) -> dict:
         """Process audio and text inputs for inference.
         Args:
@@ -157,15 +216,16 @@ class ASRProcessor(ProcessorMixin):
         Returns:
             Dict with input_features, input_ids, attention_mask
         """
-        result = {}
         token_counts = [0]
         # Process audio
         if audio is not None:
             sr = getattr(self.feature_extractor, "sampling_rate", 16000)
-            audio = prepend_lead_in(audio, sr, self.lead_in_seconds)
-            audio_inputs = self.feature_extractor(
-                audio,
                 sampling_rate=getattr(self.feature_extractor, "sampling_rate", 16000),
                 return_attention_mask=True,
                 return_tensors=return_tensors,
@@ -175,10 +235,11 @@ class ASRProcessor(ProcessorMixin):
             result["audio_attention_mask"] = audio_inputs["attention_mask"]
             if self.projector is None:
-                raise ValueError(
                     "ASRProcessor needs a projector to size the audio prompt. Build it "
                     "with ASRModel.get_processor() instead of constructing it directly."
                 )
             # One count per sample, from that sample's own mel length. Sizing a
             # single shared prompt from the batch max -- which this used to do --

 """Processor that turns raw audio (and optional text) into model inputs."""
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar, cast, overload
 import numpy as np
+import numpy.typing as npt
 import torch
 import transformers
 from torch.nn.utils.rnn import pad_sequence
+from transformers import (
+    BatchFeature,
+    PreTrainedTokenizerBase,
+    ProcessorMixin,
+    SequenceFeatureExtractor,
+)
+if TYPE_CHECKING:
+    from .asr_config import (
         DEFAULT_ENCODER_CONV_LAYERS,
         ASRConfig,
+        ConvLayerSpec,
         compute_encoder_output_length,
     )
+    from .asr_types import AudioFeatureExtractor, AudioInput, Waveform
+    from .projectors import MLPAudioProjector
+else:
+    try:
+        from .asr_config import (
+            DEFAULT_ENCODER_CONV_LAYERS,
+            ASRConfig,
+            ConvLayerSpec,
+            compute_encoder_output_length,
+        )
+        from .asr_types import AudioInput
+    except ImportError:  # flat layout on the Hub: sibling modules, no package
+        from asr_config import (
+            DEFAULT_ENCODER_CONV_LAYERS,
+            ASRConfig,
+            ConvLayerSpec,
+            compute_encoder_output_length,
+        )
+        from asr_types import AudioInput
+def left_pad_prompt_rows(
+    rows: list[torch.Tensor], tokenizer: PreTrainedTokenizerBase
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Stack per-sample prompt rows into a left-padded batch: `(input_ids, attention_mask)`.
+    Left, not right: these feed `generate`, so padding must not sit between
+    the prompt and the first generated token. Pads with the tokenizer's pad
+    token, falling back to eos, then 0. Pad positions never carry
+    `audio_token_id`, so the model's masked_scatter is unaffected.
+    """
+    # transformers types special-token ids as any token value; a single id is an int.
+    pad_id = cast("int | None", tokenizer.pad_token_id)
+    if pad_id is None:
+        pad_id = cast("int | None", tokenizer.eos_token_id) or 0
+    input_ids = pad_sequence(rows, batch_first=True, padding_value=int(pad_id), padding_side="left")
+    # Padded from ones rather than `input_ids != pad_id`: a real token may
+    # equal `pad_id` when pad falls back to eos.
+    attention_mask = pad_sequence(
+        [torch.ones_like(row) for row in rows], batch_first=True, padding_side="left"
+    )
+    return input_ids, attention_mask
+@overload
+def prepend_lead_in[ScalarT: np.generic](
+    audio: npt.NDArray[ScalarT], sampling_rate: int, seconds: float | None
+) -> npt.NDArray[ScalarT]: ...
+@overload
+def prepend_lead_in[ScalarT: np.generic](
+    audio: list[npt.NDArray[ScalarT]], sampling_rate: int, seconds: float | None
+) -> list[npt.NDArray[ScalarT]]: ...
+@overload
+def prepend_lead_in(audio: AudioInput, sampling_rate: int, seconds: float | None) -> AudioInput: ...
+def prepend_lead_in(audio: AudioInput, sampling_rate: int, seconds: float | None) -> AudioInput:
     """Prepend `seconds` of silence to a waveform (or each waveform in a list).
     Peoples ships fixed ~15s grid cuts rather than sentence-aligned segments,
         return audio
     if isinstance(audio, (list, tuple)) and audio and not isinstance(audio[0], (int, float)):
+        batch = cast("Sequence[Waveform]", audio)
+        return [cast("Waveform", prepend_lead_in(a, sampling_rate, seconds)) for a in batch]
+    waveform = cast("Waveform", audio)
     pad = round(sampling_rate * seconds)
     if pad <= 0:
+        return waveform
+    arr: npt.NDArray[Any] = np.asarray(waveform)
+    padded: npt.NDArray[Any] = np.pad(arr, (pad, 0))
+    return padded
 class ASRProcessor(ProcessorMixin):
     """Processor for Whisper-based ASR models."""
     attributes: ClassVar[list[str]] = ["feature_extractor", "tokenizer"]
+    feature_extractor: SequenceFeatureExtractor
+    tokenizer: PreTrainedTokenizerBase
     feature_extractor_class = "AutoFeatureExtractor"
     tokenizer_class = "AutoTokenizer"
     # Fallback only. The real value comes from `ASRConfig.audio_token`, which
     def __init__(
         self,
+        feature_extractor: SequenceFeatureExtractor,
+        tokenizer: PreTrainedTokenizerBase,
+        projector: "MLPAudioProjector | None" = None,
+        encoder_conv_layers: list[ConvLayerSpec] | None = None,
         audio_token: str | None = None,
         lead_in_seconds: float = 0.0,
     ):
         if text is not None:
             messages.append({"role": "assistant", "content": text})
+        # With `tokenize=True, return_tensors="pt"` the ids come back as tensors.
+        tokenized = cast(
+            "torch.Tensor | Mapping[str, torch.Tensor]",
+            self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=(text is None),
+                return_tensors="pt",
+                enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
+            ),
         )
         # apply_chat_template returns a bare tensor or a BatchEncoding/mapping.
         return (ids[0] if ids.dim() > 1 else ids).to(torch.long)
     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.
+        `ProcessorMixin.__call__` takes `(images, text, videos, audio, ...)`; this
+        processor takes audio first, so the arguments are forwarded unchanged to
+        `_process`, which carries the real signature.
         """
+        return BatchFeature(data=self._process(*args, **kwargs))
+    def _process(
         self,
+        audio: AudioInput | None = None,
         text: str | None = None,
         return_tensors: str = "pt",
+        **kwargs: Any,
+    ) -> dict[str, torch.Tensor]:
         """Process audio and text inputs for inference.
         Args:
         Returns:
             Dict with input_features, input_ids, attention_mask
         """
+        result: dict[str, torch.Tensor] = {}
         token_counts = [0]
         # Process audio
         if audio is not None:
             sr = getattr(self.feature_extractor, "sampling_rate", 16000)
+            padded_audio = prepend_lead_in(audio, sr, self.lead_in_seconds)
+            extract = cast("AudioFeatureExtractor", self.feature_extractor)
+            audio_inputs = extract(
+                padded_audio,
                 sampling_rate=getattr(self.feature_extractor, "sampling_rate", 16000),
                 return_attention_mask=True,
                 return_tensors=return_tensors,
             result["audio_attention_mask"] = audio_inputs["attention_mask"]
             if self.projector is None:
+                msg = (
                     "ASRProcessor needs a projector to size the audio prompt. Build it "
                     "with ASRModel.get_processor() instead of constructing it directly."
                 )
+                raise ValueError(msg)
             # One count per sample, from that sample's own mel length. Sizing a
             # single shared prompt from the batch max -- which this used to do --

asr_types.py ADDED Viewed

	@@ -0,0 +1,108 @@

+"""Static types shared by the asr_* modules: audio aliases, loader kwargs, protocols.
+Typing scaffolding with no behaviour of its own: the `TypedDict`s describe
+the keyword arguments ASRModel forwards to `from_pretrained` / `cached_file`,
+and the protocol classes describe surfaces with several implementations (or
+test fakes) and no common library base that declares them.
+"""
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Protocol, TypedDict, TypeVar
+import numpy.typing as npt
+import torch
+from transformers import BatchFeature, PreTrainedModel
+from transformers.generation.utils import GenerateOutput, GenerationMixin
+# One waveform (array, tensor or list of samples), or a batch of them.
+Waveform = npt.ArrayLike | torch.Tensor
+AudioInput = Waveform | Sequence[Waveform]
+# A `state_dict(destination=...)` mapping, returned as the caller's own type.
+StateDictT = TypeVar("StateDictT", bound=dict[str, Any])
+if TYPE_CHECKING:
+    class GenerativeDecoder(PreTrainedModel, GenerationMixin):
+        """Static view of the decoder: a PreTrainedModel that can generate.
+        transformers annotates GenerationMixin's `self` with a protocol
+        (`GenerativePreTrainedModel`) that no PreTrainedModel satisfies -- it
+        asks for a mutable `device` where the model has a read-only property --
+        so neither checker can bind `generate` on a real model. This restates
+        the two generation entry points ASRModel calls without that bound.
+        Type-checking only; the runtime object is whatever the loader returned.
+        """
+        def generate(self, *args: Any, **kwargs: Any) -> GenerateOutput | torch.LongTensor:
+            """`GenerationMixin.generate`."""
+            ...
+        def prepare_inputs_for_generation(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+            """`GenerationMixin.prepare_inputs_for_generation`."""
+            ...
+class LoadKwargs(TypedDict):
+    """Loader arguments shared by the encoder and decoder `from_pretrained` calls."""
+    attn_implementation: str | None
+    low_cpu_mem_usage: bool
+    dtype: torch.dtype
+class DecoderLoadKwargs(LoadKwargs):
+    """Decoder loader arguments: the shared ones plus remote-code opt-in."""
+    trust_remote_code: bool
+class HubFileKwargs(TypedDict, total=False):
+    """Hub location arguments forwarded to `cached_file` and PEFT's loader."""
+    subfolder: str
+    revision: str
+class AudioFeatureExtractor(Protocol):
+    """A concrete audio feature extractor's `__call__` (e.g. Whisper's).
+    `SequenceFeatureExtractor` declares no `__call__`, though every concrete
+    extractor has one; cast an extractor to this to call it.
+    """
+    def __call__(
+        self,
+        raw_speech: AudioInput,
+        /,
+        *,
+        sampling_rate: int,
+        return_attention_mask: bool,
+        return_tensors: str,
+        padding: bool | str = ...,
+        **kwargs: Any,
+    ) -> BatchFeature:
+        """Featurize raw audio into `input_features` (plus `attention_mask`)."""
+        ...
+class OutputLengthProjector(Protocol):
+    """The projector surface the token-count check needs."""
+    def get_output_length(self, input_length: torch.Tensor) -> torch.Tensor:
+        """Projector output length for encoder output length `input_length`."""
+        ...
+class LoadStateDictResult(Protocol):
+    """What `load_state_dict` reports back (torch's `_IncompatibleKeys`)."""
+    @property
+    def missing_keys(self) -> Sequence[str]:
+        """Model keys the state dict did not provide."""
+        ...
+    @property
+    def unexpected_keys(self) -> Sequence[str]:
+        """State-dict keys the model has no slot for."""
+        ...

diarization.py CHANGED Viewed

@@ -1,26 +1,32 @@
-"""Speaker diarization: Nemotron-3-Diarization, and TEN-VAD + ECAPA-TDNN + spectral clustering.
-Spectral clustering implementation adapted from FunASR/3D-Speaker:
-https://github.com/alibaba-damo-academy/FunASR
-MIT License (https://opensource.org/licenses/MIT)
-"""
-import warnings
 import numpy as np
 import torch
-# The ECAPA SpeakerDiarizer's packages, imported inside try blocks: transformers
-# scans every bundled file's imports before loading the model and refuses to load
-# if one is missing, and it skips imports inside `try`. The pipeline diarizes with
-# NemotronDiarizer, so these stay optional.
-_ECAPA_EXTRAS = (
-    "SpeakerDiarizer (ECAPA) needs extra packages: "
-    "pip install speechbrain ten-vad scikit-learn scipy librosa"
-)
-def _get_device() -> torch.device:
     """Get best available device for inference."""
     if torch.cuda.is_available():
         return torch.device("cuda")
@@ -29,13 +35,13 @@ def _get_device() -> torch.device:
     return torch.device("cpu")
-def _label_runs(labels: np.ndarray) -> list[tuple[int, int, int]]:
     """Split a 1-D label array into maximal constant runs.
     Returns `(value, start, end)` triples with `end` exclusive, in order. This
-    is the frames-to-segments step for both the VAD decisions and the voting
-    grid: a boundary is wherever the label changes, and the last run always
-    closes at `len(labels)`, so a trailing run needs no special case.
     """
     labels = np.asarray(labels)
     if labels.size == 0:
@@ -43,671 +49,7 @@ def _label_runs(labels: np.ndarray) -> list[tuple[int, int, int]]:
     change = np.flatnonzero(labels[1:] != labels[:-1]) + 1
     starts = np.concatenate(([0], change))
     ends = np.concatenate((change, [labels.size]))
-    return [(labels[s].item(), int(s), int(e)) for s, e in zip(starts, ends)]
-class SpectralCluster:
-    """Spectral clustering using unnormalized Laplacian of affinity matrix.
-    Adapted from FunASR/3D-Speaker and SpeechBrain implementations.
-    Uses eigenvalue gap to automatically determine number of speakers.
-    """
-    def __init__(self, min_num_spks: int = 1, max_num_spks: int = 15, pval: float = 0.06):
-        """Configure the speaker-count search range and the affinity pruning fraction."""
-        self.min_num_spks = min_num_spks
-        self.max_num_spks = max_num_spks
-        self.pval = pval
-    def __call__(self, embeddings: np.ndarray, oracle_num: int | None = None) -> np.ndarray:
-        """Run spectral clustering on embeddings.
-        Args:
-            embeddings: Speaker embeddings of shape [N, D]
-            oracle_num: Optional known number of speakers
-        Returns:
-            Cluster labels of shape [N]
-        """
-        # Similarity matrix computation
-        sim_mat = self.get_sim_mat(embeddings)
-        # Refining similarity matrix with pval
-        prunned_sim_mat = self.p_pruning(sim_mat)
-        # Symmetrization
-        sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)
-        # Laplacian calculation
-        laplacian = self.get_laplacian(sym_prund_sim_mat)
-        # Get Spectral Embeddings
-        emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)
-        # Perform clustering
-        return self.cluster_embs(emb, num_of_spk)
-    def get_sim_mat(self, embeddings: np.ndarray) -> np.ndarray:
-        """Compute cosine similarity matrix."""
-        try:
-            from sklearn.metrics.pairwise import cosine_similarity
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        return cosine_similarity(embeddings, embeddings)
-    def p_pruning(self, affinity: np.ndarray) -> np.ndarray:
-        """Prune low similarity values in affinity matrix (keep top pval fraction)."""
-        n = affinity.shape[0]
-        pval = max(self.pval, 6.0 / n)
-        k_keep = max(1, int(pval * n))
-        # Vectorized: find top-k indices per row and zero out the rest
-        top_k_idx = np.argpartition(affinity, -k_keep, axis=1)[:, -k_keep:]
-        mask = np.zeros_like(affinity, dtype=bool)
-        np.put_along_axis(mask, top_k_idx, True, axis=1)
-        affinity[~mask] = 0
-        return affinity
-    def get_laplacian(self, sim_mat: np.ndarray) -> np.ndarray:
-        """Compute unnormalized Laplacian matrix."""
-        try:
-            from scipy.sparse.csgraph import laplacian
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        np.fill_diagonal(sim_mat, 0)
-        return laplacian(sim_mat, normed=False)
-    def get_spec_embs(
-        self, laplacian: np.ndarray, k_oracle: int | None = None
-    ) -> tuple[np.ndarray, int]:
-        """Extract spectral embeddings from Laplacian."""
-        try:
-            from scipy.linalg import eigh
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        lambdas, eig_vecs = eigh(laplacian)
-        if k_oracle is not None:
-            num_of_spk = k_oracle
-        else:
-            lambda_gap_list = self.get_eigen_gaps(
-                lambdas[self.min_num_spks - 1 : self.max_num_spks + 1]
-            )
-            num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks
-        emb = eig_vecs[:, :num_of_spk]
-        return emb, num_of_spk
-    def cluster_embs(self, emb: np.ndarray, k: int) -> np.ndarray:
-        """Cluster spectral embeddings using k-means."""
-        try:
-            from sklearn.cluster._kmeans import k_means
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        _, labels, _ = k_means(emb, k, n_init=10)
-        return labels
-    def get_eigen_gaps(self, eig_vals: np.ndarray) -> np.ndarray:
-        """Compute gaps between consecutive eigenvalues."""
-        return np.diff(eig_vals)
-class SpeakerClusterer:
-    """Speaker clustering backend using spectral clustering with speaker merging.
-    Features:
-    - Spectral clustering with eigenvalue gap for auto speaker count detection
-    - P-pruning for affinity matrix refinement
-    - Post-clustering speaker merging by cosine similarity
-    """
-    def __init__(
-        self,
-        min_num_spks: int = 2,
-        max_num_spks: int = 10,
-        merge_thr: float = 0.90,  # Moderate merging
-    ):
-        """Configure the speaker-count range and the centroid-merge cosine threshold."""
-        self.min_num_spks = min_num_spks
-        self.max_num_spks = max_num_spks
-        self.merge_thr = merge_thr
-        self._spectral_cluster: SpectralCluster | None = None
-    def _get_spectral_cluster(self) -> SpectralCluster:
-        """Lazy-load spectral clusterer."""
-        if self._spectral_cluster is None:
-            self._spectral_cluster = SpectralCluster(
-                min_num_spks=self.min_num_spks,
-                max_num_spks=self.max_num_spks,
-            )
-        return self._spectral_cluster
-    def __call__(self, embeddings: np.ndarray, num_speakers: int | None = None) -> np.ndarray:
-        """Cluster speaker embeddings and return labels.
-        Args:
-            embeddings: Speaker embeddings of shape [N, D]
-            num_speakers: Optional oracle number of speakers
-        Returns:
-            Cluster labels of shape [N]
-        """
-        if len(embeddings.shape) != 2:
-            raise ValueError(f"Expected 2D array, got shape {embeddings.shape}")
-        # Handle edge cases
-        if embeddings.shape[0] == 0:
-            return np.array([], dtype=int)
-        if embeddings.shape[0] == 1:
-            return np.array([0], dtype=int)
-        if embeddings.shape[0] < 6:
-            return np.zeros(embeddings.shape[0], dtype=int)
-        # Normalize embeddings and replace NaN/inf
-        try:
-            from sklearn.preprocessing import normalize
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0)
-        embeddings = normalize(embeddings)
-        # Run spectral clustering (suppress numerical warnings)
-        spectral = self._get_spectral_cluster()
-        # Update min/max for oracle case
-        if num_speakers is not None:
-            spectral.min_num_spks = num_speakers
-            spectral.max_num_spks = num_speakers
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=RuntimeWarning)
-            labels = spectral(embeddings, oracle_num=num_speakers)
-        # Reset min/max
-        if num_speakers is not None:
-            spectral.min_num_spks = self.min_num_spks
-            spectral.max_num_spks = self.max_num_spks
-        # Merge similar speakers if no oracle
-        if num_speakers is None:
-            labels = self._merge_by_cos(labels, embeddings, self.merge_thr)
-        # Re-index labels sequentially
-        _, labels = np.unique(labels, return_inverse=True)
-        return labels
-    def _merge_by_cos(self, labels: np.ndarray, embs: np.ndarray, cos_thr: float) -> np.ndarray:
-        """Merge similar speakers by cosine similarity of centroids."""
-        try:
-            from sklearn.cluster import AgglomerativeClustering
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        try:
-            from sklearn.preprocessing import normalize
-        except ImportError as e:
-            raise ImportError(_ECAPA_EXTRAS) from e
-        unique_labels, inverse = np.unique(labels, return_inverse=True)
-        if len(unique_labels) <= 1:
-            return labels
-        # Average-linkage agglomeration over the normalized speaker centroids,
-        # cut where centroids are closer than `1 - cos_thr` in cosine distance.
-        centroids = normalize(np.array([embs[labels == lbl].mean(0) for lbl in unique_labels]))
-        merged = AgglomerativeClustering(
-            n_clusters=None,
-            distance_threshold=1.0 - cos_thr,
-            metric="cosine",
-            linkage="average",
-        ).fit_predict(centroids)
-        return merged[inverse]
-class SpeakerDiarizer:
-    """Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.
-    Example:
-        >>> segments = SpeakerDiarizer.diarize(audio_array)
-        >>> for seg in segments:
-        ...     print(f"{seg['speaker']}: {seg['start']:.2f} - {seg['end']:.2f}")
-    Pipeline:
-    1. TEN-VAD detects speech segments
-    2. Sliding window (1.0s, 75% overlap) for uniform embedding extraction
-    3. ECAPA-TDNN extracts speaker embeddings per window
-    4. Spectral clustering with eigenvalue gap for auto speaker detection
-    5. Frame-level consensus voting for segment reconstruction
-    6. Post-processing merges short segments to reduce flicker
-    Tunable Parameters (class attributes):
-    - WINDOW_SIZE: Embedding extraction window size in seconds
-    - STEP_SIZE: Sliding window step size (overlap = WINDOW_SIZE - STEP_SIZE)
-    - VAD_THRESHOLD: Speech detection threshold (lower = more sensitive)
-    - VAD_MIN_DURATION: Minimum speech segment duration
-    - VAD_MAX_GAP: Maximum gap to bridge between segments
-    - VAD_PAD_ONSET/OFFSET: Padding added to speech segments
-    - VOTING_RATE: Frame resolution for consensus voting
-    - MIN_SEGMENT_DURATION: Minimum final segment duration
-    - SAME_SPEAKER_GAP: Maximum gap to merge same-speaker segments
-    - TAIL_COVERAGE_RATIO: Minimum tail coverage to add extra window
-    """
-    _ten_vad_model = None
-    _ecapa_model = None
-    _device = None
-    # ==================== TUNABLE PARAMETERS ====================
-    # Sliding window for embedding extraction
-    WINDOW_SIZE = 0.75  # seconds - shorter window for finer resolution
-    STEP_SIZE = 0.15  # seconds (80% overlap for more votes)
-    TAIL_COVERAGE_RATIO = 0.1  # Add extra window if tail > this ratio of window
-    # VAD hysteresis parameters
-    VAD_THRESHOLD = 0.25  # Balanced threshold
-    VAD_MIN_DURATION = 0.05  # Minimum speech segment duration (seconds)
-    VAD_MAX_GAP = 0.50  # Bridge gaps shorter than this (seconds)
-    VAD_PAD_ONSET = 0.05  # Padding at segment start (seconds)
-    VAD_PAD_OFFSET = 0.05  # Padding at segment end (seconds)
-    # Frame-level voting
-    VOTING_RATE = 0.01  # 10ms resolution for consensus voting
-    # Post-processing
-    MIN_SEGMENT_DURATION = 0.15  # Minimum final segment duration (seconds)
-    SHORT_SEGMENT_GAP = 0.1  # Gap threshold for merging short segments
-    SAME_SPEAKER_GAP = 0.5  # Gap threshold for merging same-speaker segments
-    # ===========================================================
-    @classmethod
-    def _get_ten_vad_model(cls):
-        """Lazy-load TEN-VAD model (singleton)."""
-        if cls._ten_vad_model is None:
-            try:
-                from ten_vad import TenVad
-            except ImportError as e:
-                raise ImportError(_ECAPA_EXTRAS) from e
-            cls._ten_vad_model = TenVad(hop_size=256, threshold=cls.VAD_THRESHOLD)
-        return cls._ten_vad_model
-    @classmethod
-    def _get_device(cls) -> torch.device:
-        """Get the best available device."""
-        if cls._device is None:
-            cls._device = _get_device()
-        return cls._device
-    @classmethod
-    def _get_ecapa_model(cls):
-        """Lazy-load ECAPA-TDNN speaker embedding model (singleton)."""
-        if cls._ecapa_model is None:
-            # Suppress torchaudio deprecation warning from SpeechBrain
-            with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", message="torchaudio._backend")
-                try:
-                    from speechbrain.inference.speaker import EncoderClassifier
-                except ImportError as e:
-                    raise ImportError(_ECAPA_EXTRAS) from e
-                device = cls._get_device()
-                cls._ecapa_model = EncoderClassifier.from_hparams(
-                    source="speechbrain/spkrec-ecapa-voxceleb",
-                    run_opts={"device": str(device)},
-                )
-        return cls._ecapa_model
-    @classmethod
-    def diarize(
-        cls,
-        audio: np.ndarray | str,
-        sample_rate: int = 16000,
-        num_speakers: int | None = None,
-        min_speakers: int = 2,
-        max_speakers: int = 10,
-        **_kwargs,
-    ) -> list[dict]:
-        """Run speaker diarization on audio.
-        Args:
-            audio: Audio waveform as numpy array or path to audio file
-            sample_rate: Audio sample rate (default 16000)
-            num_speakers: Exact number of speakers (if known)
-            min_speakers: Minimum number of speakers
-            max_speakers: Maximum number of speakers
-        Returns:
-            List of dicts with 'speaker', 'start', 'end' keys
-        """
-        # Handle file path input
-        if isinstance(audio, str):
-            try:
-                import librosa
-            except ImportError as e:
-                raise ImportError(_ECAPA_EXTRAS) from e
-            audio, sample_rate = librosa.load(audio, sr=16000)
-        # Ensure correct sample rate
-        if sample_rate != 16000:
-            try:
-                import librosa
-            except ImportError as e:
-                raise ImportError(_ECAPA_EXTRAS) from e
-            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
-            sample_rate = 16000
-        audio = audio.astype(np.float32)
-        total_duration = len(audio) / sample_rate
-        # Step 1: VAD (returns segments and raw frame-level decisions)
-        segments, vad_frames = cls._get_speech_segments(audio, sample_rate)
-        if not segments:
-            return []
-        # Step 2: Extract embeddings
-        embeddings, window_segments = cls._extract_embeddings(audio, segments, sample_rate)
-        if len(embeddings) == 0:
-            return []
-        # Step 3: Cluster
-        clusterer = SpeakerClusterer(min_num_spks=min_speakers, max_num_spks=max_speakers)
-        labels = clusterer(embeddings, num_speakers)
-        # Step 4: Post-process with consensus voting (VAD-aware)
-        return cls._postprocess_segments(window_segments, labels, total_duration, vad_frames)
-    @classmethod
-    def _get_speech_segments(
-        cls, audio_array: np.ndarray, sample_rate: int = 16000
-    ) -> tuple[list[dict], list[bool]]:
-        """Get speech segments using TEN-VAD.
-        Returns:
-            Tuple of (segments list, vad_frames list of per-frame speech decisions)
-        """
-        vad_model = cls._get_ten_vad_model()
-        # Convert to int16 as required by TEN-VAD
-        # Clip to prevent integer overflow
-        if audio_array.dtype != np.int16:
-            audio_int16 = (np.clip(audio_array, -1.0, 1.0) * 32767).astype(np.int16)
-        else:
-            audio_int16 = audio_array
-        # Process frame by frame
-        hop_size = 256
-        frame_duration = hop_size / sample_rate
-        speech_frames: list[bool] = []
-        # `- hop_size + 1`, not `- hop_size`: the latter stops one frame
-        # early and drops the last COMPLETE frame (3 frames for 1024 samples
-        # instead of 4), truncating the trailing speech segment by one hop and
-        # forcing a silence/speaker cut at the end of every clip.
-        for i in range(0, len(audio_int16) - hop_size + 1, hop_size):
-            frame = audio_int16[i : i + hop_size]
-            _, is_speech = vad_model.process(frame)
-            speech_frames.append(is_speech)
-        # Convert frame-level decisions to segments: one per run of speech frames.
-        segments = []
-        for is_speech, start_idx, end_idx in _label_runs(np.array(speech_frames, dtype=bool)):
-            if not is_speech:
-                continue
-            start_time = start_idx * frame_duration
-            end_time = end_idx * frame_duration
-            segments.append(
-                {
-                    "start": start_time,
-                    "end": end_time,
-                    "start_sample": int(start_time * sample_rate),
-                    "end_sample": int(end_time * sample_rate),
-                }
-            )
-        return cls._apply_vad_hysteresis(segments, sample_rate), speech_frames
-    @classmethod
-    def _apply_vad_hysteresis(cls, segments: list[dict], sample_rate: int = 16000) -> list[dict]:
-        """Apply hysteresis-like post-processing to VAD segments."""
-        if not segments:
-            return segments
-        segments = sorted(segments, key=lambda x: x["start"])
-        # Fill short gaps
-        merged = [segments[0].copy()]
-        for seg in segments[1:]:
-            gap = seg["start"] - merged[-1]["end"]
-            if gap <= cls.VAD_MAX_GAP:
-                merged[-1]["end"] = seg["end"]
-                merged[-1]["end_sample"] = seg["end_sample"]
-            else:
-                merged.append(seg.copy())
-        # Remove short segments
-        filtered = [seg for seg in merged if (seg["end"] - seg["start"]) >= cls.VAD_MIN_DURATION]
-        # Dilate segments (add padding)
-        for seg in filtered:
-            seg["start"] = max(0.0, seg["start"] - cls.VAD_PAD_ONSET)
-            seg["end"] = seg["end"] + cls.VAD_PAD_OFFSET
-            seg["start_sample"] = int(seg["start"] * sample_rate)
-            seg["end_sample"] = int(seg["end"] * sample_rate)
-        return filtered
-    @classmethod
-    def _extract_embeddings(
-        cls, audio_array: np.ndarray, segments: list[dict], sample_rate: int
-    ) -> tuple[np.ndarray, list[dict]]:
-        """Extract speaker embeddings using sliding windows."""
-        speaker_model = cls._get_ecapa_model()
-        window_samples = int(cls.WINDOW_SIZE * sample_rate)
-        step_samples = int(cls.STEP_SIZE * sample_rate)
-        embeddings = []
-        window_segments = []
-        with torch.no_grad():
-            for seg in segments:
-                seg_start = seg["start_sample"]
-                seg_end = seg["end_sample"]
-                chunk = audio_array[seg_start:seg_end]
-                if len(chunk) <= window_samples:
-                    # Pad short segments with reflection to a full window.
-                    starts = [seg_start]
-                    padded = np.pad(chunk, (0, window_samples - len(chunk)), mode="reflect")
-                    windows = torch.from_numpy(padded).float().unsqueeze(0)
-                else:
-                    chunk = torch.from_numpy(chunk).float()
-                    # `unfold` is the sliding window: row i covers
-                    # [i * step, i * step + window), as far as a full window fits.
-                    windows = chunk.unfold(0, window_samples, step_samples)
-                    starts = [seg_start + i * step_samples for i in range(windows.shape[0])]
-                    # Cover tail if > TAIL_COVERAGE_RATIO of window remains
-                    remainder = seg_end - (starts[-1] + window_samples)
-                    if remainder > (window_samples * cls.TAIL_COVERAGE_RATIO):
-                        windows = torch.cat([windows, chunk[-window_samples:].unsqueeze(0)])
-                        starts.append(seg_end - window_samples)
-                # One batched forward per segment: every window is exactly
-                # `window_samples` long, so no length masking is needed.
-                batch = speaker_model.encode_batch(windows).reshape(len(starts), -1).cpu().numpy()
-                for c_start, embedding in zip(starts, batch):
-                    # Validate embedding
-                    if np.isfinite(embedding).all() and np.linalg.norm(embedding) > 1e-8:
-                        embeddings.append(embedding)
-                        window_segments.append(
-                            {
-                                "start": c_start / sample_rate,
-                                "end": min(c_start + window_samples, seg_end) / sample_rate,
-                            }
-                        )
-        # Normalize all embeddings at once
-        if embeddings:
-            try:
-                from sklearn.preprocessing import normalize
-            except ImportError as e:
-                raise ImportError(_ECAPA_EXTRAS) from e
-            return normalize(np.array(embeddings)), window_segments
-        return np.array([]), []
-    @classmethod
-    def _resample_vad(cls, vad_frames: list[bool], num_frames: int) -> np.ndarray:
-        """Resample VAD frame decisions to match voting grid resolution.
-        VAD operates at 256 samples / 16000 Hz = 16ms per frame.
-        Voting operates at VOTING_RATE (default 10ms) per frame.
-        This maps VAD decisions to the finer voting grid.
-        """
-        if not vad_frames:
-            return np.zeros(num_frames, dtype=bool)
-        vad_rate = 256 / 16000  # 16ms per VAD frame
-        vad_arr = np.array(vad_frames)
-        # Vectorized: compute VAD frame indices for each voting frame
-        voting_times = np.arange(num_frames) * cls.VOTING_RATE
-        vad_indices = np.clip((voting_times / vad_rate).astype(int), 0, len(vad_arr) - 1)
-        return vad_arr[vad_indices]
-    @classmethod
-    def _postprocess_segments(
-        cls,
-        window_segments: list[dict],
-        labels: np.ndarray,
-        total_duration: float,
-        vad_frames: list[bool],
-    ) -> list[dict]:
-        """Post-process using frame-level consensus voting with VAD-aware silence."""
-        if not window_segments or len(labels) == 0:
-            return []
-        # Correct labels to be contiguous
-        unique_labels, clean_labels = np.unique(labels, return_inverse=True)
-        num_speakers = len(unique_labels)
-        if num_speakers == 0:
-            return []
-        # Create voting grid
-        num_frames = int(np.ceil(total_duration / cls.VOTING_RATE)) + 1
-        votes = np.zeros((num_frames, num_speakers), dtype=np.float32)
-        # Accumulate votes
-        for win, label in zip(window_segments, clean_labels):
-            start_frame = int(win["start"] / cls.VOTING_RATE)
-            end_frame = int(win["end"] / cls.VOTING_RATE)
-            end_frame = min(end_frame, num_frames)
-            if start_frame < end_frame:
-                votes[start_frame:end_frame, label] += 1.0
-        # Determine winner per frame
-        frame_speakers = np.argmax(votes, axis=1)
-        max_votes = np.max(votes, axis=1)
-        # Resample VAD to voting grid resolution for silence-aware voting
-        vad_resampled = cls._resample_vad(vad_frames, num_frames)
-        # Force silence (-1) where VAD says no speech or no window voted.
-        frame_speakers = np.where((max_votes > 0) & vad_resampled, frame_speakers, -1)
-        # Convert frames to segments: one per run of the same speaker.
-        final_segments = [
-            {
-                "speaker": f"SPEAKER_{speaker}",
-                "start": start * cls.VOTING_RATE,
-                "end": end * cls.VOTING_RATE,
-            }
-            for speaker, start, end in _label_runs(frame_speakers)
-            if speaker != -1
-        ]
-        return cls._merge_short_segments(final_segments)
-    @classmethod
-    def _merge_short_segments(cls, segments: list[dict]) -> list[dict]:
-        """Merge short segments to reduce flicker."""
-        if not segments:
-            return []
-        clean: list[dict] = []
-        for seg in segments:
-            dur = seg["end"] - seg["start"]
-            if dur < cls.MIN_SEGMENT_DURATION:
-                if (
-                    clean
-                    and clean[-1]["speaker"] == seg["speaker"]
-                    and seg["start"] - clean[-1]["end"] < cls.SHORT_SEGMENT_GAP
-                ):
-                    clean[-1]["end"] = seg["end"]
-                continue
-            if (
-                clean
-                and clean[-1]["speaker"] == seg["speaker"]
-                and seg["start"] - clean[-1]["end"] < cls.SAME_SPEAKER_GAP
-            ):
-                clean[-1]["end"] = seg["end"]
-            else:
-                clean.append(seg)
-        return clean
-    @classmethod
-    def assign_speakers_to_words(
-        cls,
-        words: list[dict],
-        speaker_segments: list[dict],
-    ) -> list[dict]:
-        """Assign speaker labels to words based on timestamp overlap.
-        Args:
-            words: List of word dicts with 'word', 'start', 'end' keys
-            speaker_segments: List of speaker dicts with 'speaker', 'start', 'end' keys
-        Returns:
-            Words list with 'speaker' key added to each word
-        """
-        for word in words:
-            word_mid = (word["start"] + word["end"]) / 2
-            # Find the speaker segment that contains this word's midpoint
-            best_speaker = None
-            for seg in speaker_segments:
-                if seg["start"] <= word_mid <= seg["end"]:
-                    best_speaker = seg["speaker"]
-                    break
-            # If no exact match, find closest segment
-            if best_speaker is None and speaker_segments:
-                min_dist = float("inf")
-                for seg in speaker_segments:
-                    seg_mid = (seg["start"] + seg["end"]) / 2
-                    dist = abs(word_mid - seg_mid)
-                    if dist < min_dist:
-                        min_dist = dist
-                        best_speaker = seg["speaker"]
-            word["speaker"] = best_speaker
-        return words
 class NemotronDiarizer:
@@ -733,46 +75,54 @@ class NemotronDiarizer:
     FRAME_S = 0.01
     SEGMENT_THRESHOLD = 0.5  # a speaker's segment = frames above this
     MIN_WORD_ACTIVITY = 0.1  # below this nobody is heard: the word inherits its neighbour
-    _model = None
-    _processor = None
     @classmethod
-    def get_instance(cls):
-        if cls._model is None:
-            import importlib.util
             if importlib.util.find_spec("transformers.models.nemotron3_diarization") is None:
-                raise ImportError(
                     "Speaker diarization uses Nemotron-3-Diarization, which needs a transformers "
                     "build with `nemotron3_diarization`: "
                     "pip install git+https://github.com/huggingface/transformers"
                 )
-            from transformers import AutoModelForAudioFrameClassification, AutoProcessor
-            model = AutoModelForAudioFrameClassification.from_pretrained(cls.MODEL_ID)
-            cls._model = model.to(_get_device()).eval()
-            cls._processor = AutoProcessor.from_pretrained(cls.MODEL_ID)
         return cls._model, cls._processor
     @classmethod
     @torch.inference_mode()
-    def activity(cls, audio: np.ndarray, sample_rate: int = 16000) -> np.ndarray:
         """(frames, 8) speech probabilities at 10 ms; columns in arrival order."""
         if sample_rate != 16000:
-            import torchaudio
-            audio = torchaudio.functional.resample(
-                torch.as_tensor(audio, dtype=torch.float32), sample_rate, 16000
-            ).numpy()
         model, processor = cls.get_instance()
         inputs = processor(np.asarray(audio, dtype=np.float32), sampling_rate=16000)
         inputs = inputs.to(model.device, dtype=model.dtype)
-        return model(**inputs).logits[0].float().sigmoid().cpu().numpy()
     @classmethod
     def top_speakers(
         cls,
-        activity: np.ndarray,
         num_speakers: int | None = None,
         max_speakers: int | None = None,
     ) -> list[int]:
@@ -793,10 +143,10 @@ class NemotronDiarizer:
         return active or [0]
     @classmethod
-    def segments(cls, activity: np.ndarray, keep: list[int]) -> list[dict]:
         """Speaker turns `{"speaker", "start", "end"}` by start time; overlaps allowed."""
         names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
-        out = []
         for c in keep:
             for on, s, e in _label_runs(activity[:, c] > cls.SEGMENT_THRESHOLD):
                 if on:
@@ -809,7 +159,7 @@ class NemotronDiarizer:
     def speaker_columns(
         cls,
         spans: list[tuple[float | None, float | None]],
-        activity: np.ndarray,
         keep: list[int],
     ) -> list[int]:
         """The kept activity column that said each `(start_s, end_s)` word.
@@ -832,7 +182,7 @@ class NemotronDiarizer:
             picked.append(int(cols[best]) if mean[best] >= cls.MIN_WORD_ACTIVITY else None)
         known = [p for p in picked if p is not None]
         last = known[0] if known else keep[0]
-        out = []
         for p in picked:
             last = p if p is not None else last
             out.append(last)
@@ -840,11 +190,11 @@ class NemotronDiarizer:
     @classmethod
     def assign_speakers_to_words(
-        cls, words: list[dict], activity: np.ndarray, keep: list[int]
-    ) -> list[dict]:
         """Label each `{"start", "end"}` word `SPEAKER_i` by `speaker_columns`."""
         names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
         spans = [(w["start"], w["end"]) for w in words]
-        for word, col in zip(words, cls.speaker_columns(spans, activity, keep)):
             word["speaker"] = names[col]
         return words

+"""Speaker diarization with NVIDIA Nemotron-3-Diarization."""
+import importlib.util
+from typing import TYPE_CHECKING, Any, TypedDict
 import numpy as np
+import numpy.typing as npt
 import torch
+import torchaudio
+from transformers import AutoModelForAudioFrameClassification, AutoProcessor
+if TYPE_CHECKING:
+    # In transformers main after 5.17; loaded through the Auto classes at runtime
+    # so an older build reaches get_instance's install hint instead of failing here.
+    from transformers import (
+        Nemotron3DiarizationForAudioFrameClassification,
+        Nemotron3DiarizationProcessor,
+    )
+class SpeakerSegment(TypedDict):
+    """One speaker turn, in seconds."""
+    speaker: str
+    start: float
+    end: float
+def get_device() -> torch.device:
     """Get best available device for inference."""
     if torch.cuda.is_available():
         return torch.device("cuda")
     return torch.device("cpu")
+def _label_runs(labels: npt.NDArray[np.generic]) -> list[tuple[int, int, int]]:
     """Split a 1-D label array into maximal constant runs.
     Returns `(value, start, end)` triples with `end` exclusive, in order. This
+    is the frames-to-segments step for speaker activity: a boundary is wherever
+    the label changes, and the last run always closes at `len(labels)`, so a
+    trailing run needs no special case.
     """
     labels = np.asarray(labels)
     if labels.size == 0:
     change = np.flatnonzero(labels[1:] != labels[:-1]) + 1
     starts = np.concatenate(([0], change))
     ends = np.concatenate((change, [labels.size]))
+    return [(labels[s].item(), int(s), int(e)) for s, e in zip(starts, ends, strict=True)]
 class NemotronDiarizer:
     FRAME_S = 0.01
     SEGMENT_THRESHOLD = 0.5  # a speaker's segment = frames above this
     MIN_WORD_ACTIVITY = 0.1  # below this nobody is heard: the word inherits its neighbour
+    _model: "Nemotron3DiarizationForAudioFrameClassification | None" = None
+    _processor: "Nemotron3DiarizationProcessor | None" = None
     @classmethod
+    def get_instance(
+        cls,
+    ) -> "tuple[Nemotron3DiarizationForAudioFrameClassification, Nemotron3DiarizationProcessor]":
+        """Load the diarization model and processor once, then return the cached pair."""
+        if cls._model is None or cls._processor is None:
             if importlib.util.find_spec("transformers.models.nemotron3_diarization") is None:
+                msg = (
                     "Speaker diarization uses Nemotron-3-Diarization, which needs a transformers "
                     "build with `nemotron3_diarization`: "
                     "pip install git+https://github.com/huggingface/transformers"
                 )
+                raise ImportError(msg)
+            model: Nemotron3DiarizationForAudioFrameClassification = (
+                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]
+            model.eval()
+            cls._model = model
+            processor: Nemotron3DiarizationProcessor = AutoProcessor.from_pretrained(cls.MODEL_ID)
+            cls._processor = processor
         return cls._model, cls._processor
     @classmethod
     @torch.inference_mode()
+    def activity(cls, audio: npt.ArrayLike, sample_rate: int = 16000) -> npt.NDArray[np.float32]:
         """(frames, 8) speech probabilities at 10 ms; columns in arrival order."""
         if sample_rate != 16000:
+            audio = np.asarray(
+                torchaudio.functional.resample(
+                    torch.as_tensor(audio, dtype=torch.float32), sample_rate, 16000
+                ),
+                dtype=np.float32,
+            )
         model, processor = cls.get_instance()
         inputs = processor(np.asarray(audio, dtype=np.float32), sampling_rate=16000)
         inputs = inputs.to(model.device, dtype=model.dtype)
+        probs: npt.NDArray[np.float32] = model(**inputs).logits[0].float().sigmoid().cpu().numpy()
+        return probs
     @classmethod
     def top_speakers(
         cls,
+        activity: npt.NDArray[np.float32],
         num_speakers: int | None = None,
         max_speakers: int | None = None,
     ) -> list[int]:
         return active or [0]
     @classmethod
+    def segments(cls, activity: npt.NDArray[np.float32], keep: list[int]) -> list[SpeakerSegment]:
         """Speaker turns `{"speaker", "start", "end"}` by start time; overlaps allowed."""
         names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+        out: list[SpeakerSegment] = []
         for c in keep:
             for on, s, e in _label_runs(activity[:, c] > cls.SEGMENT_THRESHOLD):
                 if on:
     def speaker_columns(
         cls,
         spans: list[tuple[float | None, float | None]],
+        activity: npt.NDArray[np.float32],
         keep: list[int],
     ) -> list[int]:
         """The kept activity column that said each `(start_s, end_s)` word.
             picked.append(int(cols[best]) if mean[best] >= cls.MIN_WORD_ACTIVITY else None)
         known = [p for p in picked if p is not None]
         last = known[0] if known else keep[0]
+        out: list[int] = []
         for p in picked:
             last = p if p is not None else last
             out.append(last)
     @classmethod
     def assign_speakers_to_words(
+        cls, words: list[dict[str, Any]], activity: npt.NDArray[np.float32], keep: list[int]
+    ) -> list[dict[str, Any]]:
         """Label each `{"start", "end"}` word `SPEAKER_i` by `speaker_columns`."""
         names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
         spans = [(w["start"], w["end"]) for w in words]
+        for word, col in zip(words, cls.speaker_columns(spans, activity, keep), strict=True):
             word["speaker"] = names[col]
         return words

handler.py CHANGED Viewed

@@ -1,26 +1,25 @@
 """Custom inference handler for HuggingFace Inference Endpoints."""
-from typing import Any
-import torch
-try:
-    # For remote execution, imports are relative
     from .asr_modeling import ASRModel
     from .asr_pipeline import ASRPipeline
-except ImportError:
-    # For local execution, imports are not relative
-    from asr_modeling import ASRModel  # type: ignore[no-redef]
-    from asr_pipeline import ASRPipeline  # type: ignore[no-redef]
-def _best_device() -> torch.device:
-    """Best available inference device (cuda > mps > cpu)."""
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    if torch.backends.mps.is_available():
-        return torch.device("mps")
-    return torch.device("cpu")
 class EndpointHandler:
@@ -36,11 +35,12 @@ class EndpointHandler:
         Args:
             path: Path to model directory or HuggingFace model ID
         """
-        import os
-        import nltk
-        nltk.download("punkt_tab", quiet=True)
         os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
@@ -56,7 +56,8 @@ class EndpointHandler:
         # flash_attn is missing.
         self.model = ASRModel.from_pretrained(path)
         self.device = _best_device()
-        self.model.to(self.device)
         self.model.eval()
         self.pipe = ASRPipeline(
@@ -77,9 +78,11 @@ class EndpointHandler:
         """
         inputs = data.get("inputs")
         if inputs is None:
-            raise ValueError("Missing 'inputs' in request data")
         # Pass through any parameters from request, let model config provide defaults
         params = data.get("parameters", {})
-        return self.pipe(inputs, **params)

 """Custom inference handler for HuggingFace Inference Endpoints."""
+import os
+from typing import TYPE_CHECKING, Any
+import nltk
+if TYPE_CHECKING:
     from .asr_modeling import ASRModel
     from .asr_pipeline import ASRPipeline
+    from .diarization import get_device as _best_device
+else:
+    try:
+        # For remote execution, imports are relative
+        from .asr_modeling import ASRModel
+        from .asr_pipeline import ASRPipeline
+        from .diarization import get_device as _best_device
+    except ImportError:
+        # For local execution, imports are not relative
+        from asr_modeling import ASRModel
+        from asr_pipeline import ASRPipeline
+        from diarization import get_device as _best_device
 class EndpointHandler:
         Args:
             path: Path to model directory or HuggingFace model ID
         """
+        # Only fetch when absent: nltk.download re-checks the remote index
+        # every call, i.e. a network round-trip on every endpoint start.
+        try:
+            nltk.data.find("tokenizers/punkt_tab/english/")
+        except LookupError:
+            nltk.download("punkt_tab", quiet=True)
         os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
         # 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(
         """
         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(inputs, **params)
+        return result

projectors.py CHANGED Viewed

@@ -4,6 +4,8 @@ This module contains all projector architectures:
 - MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling
 """
 import torch
 import torch.nn as nn
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
@@ -90,7 +92,7 @@ class MLPAudioProjector(nn.Module):
     path from encoder to decoder and a frozen encoder that cannot adapt to it.
     """
-    def __init__(self, config):
         """Initialize MLP projector.
         Args:
@@ -98,20 +100,24 @@ class MLPAudioProjector(nn.Module):
         """
         super().__init__()
-        encoder_dim = getattr(config, "encoder_dim", 768)
-        llm_dim = getattr(config, "llm_dim", 2048)
-        self.k = getattr(config, "projector_pool_stride", 4)
         # Frame stacking: concat k adjacent frames then project
         in_dim = encoder_dim * self.k
         # Hidden dim defaults to llm_dim, can be overridden via config
-        hidden_dim = 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)
-    def get_output_length(self, input_length: int) -> int:
         """Calculate output sequence length given input length (matches GLM-ASR)."""
         return _frame_stack_length(input_length, self.k)
@@ -136,7 +142,11 @@ class MLPAudioProjector(nn.Module):
 # =============================================================================
-def _frame_stack_length(length, k: int):
     """Frames left after stacking k adjacent frames (GLM-ASR's rule).
     Trailing frames that don't fill a complete k-frame window are dropped.

 - MLPAudioProjector: Simple 2-layer MLP with frame stacking downsampling
 """
+from typing import overload
 import torch
 import torch.nn as nn
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
     path from encoder to decoder and a frozen encoder that cannot adapt to it.
     """
+    def __init__(self, config: object) -> None:
         """Initialize MLP projector.
         Args:
         """
         super().__init__()
+        encoder_dim: int = getattr(config, "encoder_dim", 768)
+        llm_dim: int = getattr(config, "llm_dim", 2048)
+        self.k: int = getattr(config, "projector_pool_stride", 4)
         # Frame stacking: concat k adjacent frames then project
         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)
+    @overload
+    def get_output_length(self, input_length: int) -> int: ...
+    @overload
+    def get_output_length(self, input_length: torch.Tensor) -> torch.Tensor: ...
+    def get_output_length(self, input_length: int | torch.Tensor) -> int | torch.Tensor:
         """Calculate output sequence length given input length (matches GLM-ASR)."""
         return _frame_stack_length(input_length, self.k)
 # =============================================================================
+@overload
+def _frame_stack_length(length: int, k: int) -> int: ...
+@overload
+def _frame_stack_length(length: torch.Tensor, k: int) -> torch.Tensor: ...
+def _frame_stack_length(length: int | torch.Tensor, k: int) -> int | torch.Tensor:
     """Frames left after stacking k adjacent frames (GLM-ASR's rule).
     Trailing frames that don't fill a complete k-frame window are dropped.