Update custom model files, README, and requirements

alignment.py

@@ -38,13 +38,26 @@ class AlignedWord(TypedDict):
     end: float
 
 
-def _get_device() -> str:
+def get_device() -> torch.device:
     """Get best available device for inference."""
     if torch.cuda.is_available():
-        return "cuda"
+        return torch.device("cuda")
     if torch.backends.mps.is_available():
-        return "mps"
-    return "cpu"
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+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
 
 
 class QwenForcedAligner:
@@ -73,9 +86,9 @@ class QwenForcedAligner:
         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()
+            device = get_device()
             # Batches are padded, and Metal sdpa returns NaN for fully masked rows.
-            attn = "eager" if device == "mps" else "sdpa"
+            attn = "eager" if device.type == "mps" else "sdpa"
             model = Qwen3ASRForTokenClassification.from_pretrained(
                 cls.MODEL_ID, dtype=torch.bfloat16, attn_implementation=attn
             )
@@ -86,19 +99,6 @@ class QwenForcedAligner:
             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(
@@ -116,7 +116,7 @@ class QwenForcedAligner:
             kept = [w for w in text.split() if _clean_tokens([w])]
             if not kept:
                 continue
-            audio = cls._to_16k(raw, sample_rate)
+            audio = 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 "

asr_attention.py

@@ -9,6 +9,7 @@ requested implementation otherwise.
 import functools
 import importlib
 import logging
+from typing import Any
 
 import torch
 from transformers import AutoConfig
@@ -41,6 +42,13 @@ def _max_attention_head_dim(text_config: object) -> int | None:
     return max(usable) if usable else None
 
 
+@functools.lru_cache(maxsize=16)
+def _text_config(model_id: str) -> Any:
+    """`model_id`'s decoder config: the text sub-config of a composite config, else itself."""
+    probe = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    return probe.get_text_config() if hasattr(probe, "get_text_config") else probe
+
+
 @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.
@@ -58,8 +66,7 @@ def _has_sliding_window_attention(model_id: str) -> bool:
     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
+        text_config = _text_config(model_id)
     except Exception:
         logger.warning(
             "Could not read the config for %s to check for sliding-window "
@@ -156,9 +163,7 @@ def resolve_decoder_attn_implementation(requested: str | None, model_id: str) ->
     """
     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)
+        head_dim = _max_attention_head_dim(_text_config(model_id))
         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; "

asr_modeling.py

@@ -1,3 +1,4 @@
+import contextlib
 import functools
 import importlib.util
 import inspect
@@ -22,7 +23,6 @@ from transformers import (
     AutoModelForCausalLM,
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
-    BatchEncoding,
     PreTrainedModel,
     PreTrainedTokenizerBase,
     SequenceFeatureExtractor,
@@ -47,7 +47,12 @@ if TYPE_CHECKING:
     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_processing import (
+        DEFAULT_TRANSCRIBE_PROMPT,
+        ASRProcessor,
+        left_pad_prompt_rows,
+        render_audio_prompt,
+    )
     from .asr_types import (
         DecoderLoadKwargs,
         GenerativeDecoder,
@@ -67,14 +72,24 @@ else:
         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_processing import (
+            DEFAULT_TRANSCRIBE_PROMPT,
+            ASRProcessor,
+            left_pad_prompt_rows,
+            render_audio_prompt,
+        )
         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_processing import (
+            DEFAULT_TRANSCRIBE_PROMPT,
+            ASRProcessor,
+            left_pad_prompt_rows,
+            render_audio_prompt,
+        )
         from asr_types import LoadStateDictResult, OutputLengthProjector, StateDictT
         from projectors import PROJECTOR_CLASSES
 
@@ -86,10 +101,8 @@ else:
     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
 
 
@@ -197,6 +210,16 @@ def _assert_audio_token_counts(
         raise ValueError(msg)
 
 
+def _splice_audio(
+    inputs_embeds: torch.Tensor, is_audio_token: torch.Tensor, audio_embeds: torch.Tensor
+) -> torch.Tensor:
+    """Scatter `audio_embeds` into `inputs_embeds` at the audio-placeholder positions."""
+    dev = inputs_embeds.device
+    return inputs_embeds.masked_scatter(
+        is_audio_token.unsqueeze(-1).to(dev), audio_embeds.to(dev, dtype=inputs_embeds.dtype)
+    )
+
+
 def _patch_gemma_decode_loop(model: "GenerativeDecoder") -> None:
     """Make a Gemma4ForCausalLM drop `per_layer_inputs` after the first step.
 
@@ -321,7 +344,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
     supports_gradient_checkpointing = True
     _is_loading_from_pretrained: bool = False
 
-    TRANSCRIBE_PROMPT = "Transcribe the speech to text"
+    TRANSCRIBE_PROMPT = DEFAULT_TRANSCRIBE_PROMPT
 
     @classmethod
     def from_pretrained(
@@ -481,9 +504,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         # Setup LoRA if enabled (Stage 2 fine-tuning)
         # Skip if loading from pretrained - from_pretrained will handle adapter loading
-        if getattr(config, "use_lora", False) and not getattr(
-            self.__class__, "_is_loading_from_pretrained", False
-        ):
+        if getattr(config, "use_lora", False) and not type(self)._is_loading_from_pretrained:
             self._setup_lora(config)
 
         # Freeze projector if specified (for Stage 2 LoRA-only training)
@@ -757,7 +778,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             raise ImportError(msg)
 
         # Only reached for gemma-4 checkpoints, whose text config is Gemma4TextConfig.
-        text_config = cast(Gemma4TextConfig, loaded.config.get_text_config())
+        text_config = cast("Gemma4TextConfig", loaded.config.get_text_config())
         with torch.device("meta"):
             shell = Gemma4ForCausalLM(text_config)
 
@@ -766,17 +787,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         shell.generation_config = loaded.generation_config
         shell.tie_weights()
 
-        # Gemma4ForCausalLM does not drop `per_layer_inputs` between decode
-        # steps, while the multimodal wrapper does. Without the drop, step 2
-        # of generation passes PLE alongside `input_ids` and forward raises
-        # "You cannot specify per_layer_inputs if input_ids is provided". The
-        # drop is semantically right, not a workaround: PLE describes the
-        # prompt, and from step 2 on generation feeds one freshly sampled
-        # token whose PLE Gemma derives itself from that token's id.
-        #
-        # Patched on the instance rather than via a subclass: subclassing
-        # Gemma4ForCausalLM trips transformers' experts-implementation check
-        # ("does not support setting experts implementation") during __init__.
+        # See _patch_gemma_decode_loop for why PLE is dropped after step 1.
         # Gemma4ForCausalLM is a PreTrainedModel with GenerationMixin, which is
         # all `GenerativeDecoder` describes.
         decoder = cast("GenerativeDecoder", shell)
@@ -1171,12 +1182,23 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         to the projector); includes the encoder only when it's trainable.
         """
         targets: list[nn.Module] = [self.language_model]
-        # Autograd state again, not the flag: a partially unfrozen encoder has
-        # an activation tape worth checkpointing even though the flag is True.
-        if any(p.requires_grad for p in self.audio_tower.parameters()):
+        if self._encoder_has_trainable_params():
             targets.append(self.audio_tower)
         return targets
 
+    def _encoder_has_trainable_params(self) -> bool:
+        """True if any encoder parameter requires grad.
+
+        Keyed on autograd state, NOT on `freeze_audio_encoder`. The documented
+        partial-unfreeze recipe is `freeze_audio_encoder: true` PLUS
+        `encoder_trainable_top_layers: N` (see ASRConfig), which leaves the flag
+        True while `_load_audio_encoder` switches the top N blocks back on.
+        Reading the flag instead once wrapped those blocks in no_grad (so the
+        run paid ~1.6 GiB of AdamW state for 198M params that could never
+        move) and kept them out of the checkpoint.
+        """
+        return any(p.requires_grad for p in self.audio_tower.parameters())
+
     def get_input_embeddings(self) -> nn.Module:
         """Return the decoder's token embedding module."""
         return self.language_model.get_input_embeddings()
@@ -1381,22 +1403,12 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         ):
             tower_kwargs["attention_mask"] = audio_attention_mask.to(audio_features.device)
 
-        # Gate on autograd state, NOT on `freeze_audio_encoder` alone. The
-        # documented partial-unfreeze recipe is `freeze_audio_encoder: true`
-        # PLUS `encoder_trainable_top_layers: N` (see ASRConfig), which leaves
-        # the flag True while `_load_audio_encoder` switches the top N blocks
-        # back on. Reading the flag here then wrapped those blocks in no_grad
-        # and zeroed their gradients -- exactly the failure the comment above
-        # warns about -- so the run paid ~1.6 GiB of AdamW state for 198M
-        # params that could never move, with nothing in the logs saying so.
-        encoder_frozen = not any(p.requires_grad for p in self.audio_tower.parameters())
-        if encoder_frozen:
-            with torch.no_grad():
-                encoder_out = self.audio_tower(**tower_kwargs)
-                hidden_states = encoder_out.last_hidden_state
-        else:
+        grad_ctx = (
+            contextlib.nullcontext() if self._encoder_has_trainable_params() else torch.no_grad()
+        )
+        with grad_ctx:
             encoder_out = self.audio_tower(**tower_kwargs)
-            hidden_states = encoder_out.last_hidden_state
+        hidden_states = encoder_out.last_hidden_state
 
         # Conformer encoders (Granite) return their own validity mask, halved
         # alongside each subsampling block. That is the encoder's arithmetic
@@ -1613,11 +1625,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 input_features, audio_token_counts, audio_attention_mask
             )
 
-            audio_token_mask = is_audio_token.unsqueeze(-1)
-            inputs_embeds = inputs_embeds.masked_scatter(
-                audio_token_mask.to(inputs_embeds.device),
-                audio_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype),
-            )
+            inputs_embeds = _splice_audio(inputs_embeds, is_audio_token, audio_embeds)
 
         # Forward label_smoothing to the LM's loss_function via **kwargs.
         # transformers.loss.loss_utils.ForCausalLMLoss → fixed_cross_entropy
@@ -1705,43 +1713,11 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         return model_inputs
 
-    def _get_num_audio_tokens(
-        self,
-        audio_attention_mask: torch.Tensor,
-    ) -> int:
-        """Audio token count for the LONGEST sample in the batch.
-
-        mel_frames -> encoder_frames (via conv formulas) -> projector tokens.
-
-        Do not build a shared prompt from this. Every row shorter than the
-        batch max needs fewer placeholders than this returns; see
-        `_prepare_audio_inputs`, which renders one prompt per sample.
-        """
-        encoder_lengths = self._compute_encoder_output_lengths(audio_attention_mask)
-        encoder_output_len = int(encoder_lengths.max().item())
-        return int(self.projector.get_output_length(encoder_output_len))
-
     def _render_audio_prompt(self, num_audio_tokens: int) -> torch.Tensor:
         """Tokenize one chat prompt carrying exactly `num_audio_tokens` placeholders."""
-        messages: list[dict[str, str]] = []
-        # Audio tokens only (instruction-free) unless a transcribe prompt is set
-        user_content = self.audio_token * num_audio_tokens
-        if self.TRANSCRIBE_PROMPT:
-            user_content += " " + self.TRANSCRIBE_PROMPT
-        messages.append({"role": "user", "content": user_content})
-
-        chat_result = self.tokenizer.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt",
-            enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
+        return render_audio_prompt(
+            self.tokenizer, self.audio_token, num_audio_tokens, self.TRANSCRIBE_PROMPT
         )
-        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(
         self,
@@ -1794,11 +1770,8 @@ 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._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),
-            audio_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype),
+        inputs_embeds = _splice_audio(
+            self._embed_tokens(input_ids), input_ids == self.audio_token_id, audio_embeds
         )
         return input_ids, attention_mask, inputs_embeds
 
@@ -2115,14 +2088,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         # Copy source files for auto-loading
         src_dir = Path(__file__).parent
-        for asr_file in src_dir.glob("asr_*.py"):
-            shutil.copy(asr_file, save_dir / asr_file.name)
-        # Copy projectors module
-        shutil.copy(src_dir / "projectors.py", save_dir / "projectors.py")
-        # Copy alignment module
-        shutil.copy(src_dir / "alignment.py", save_dir / "alignment.py")
-        # Copy diarization module
-        shutil.copy(src_dir / "diarization.py", save_dir / "diarization.py")
+        siblings = ("projectors.py", "alignment.py", "diarization.py")
+        for src in [*src_dir.glob("asr_*.py"), *(src_dir / name for name in siblings)]:
+            shutil.copy(src, save_dir / src.name)
 
     def create_or_update_model_card(self, output_dir: str | Path) -> None:
         """Re-apply the PEFT card metadata to `output_dir/README.md`.

asr_pipeline.py

@@ -14,7 +14,6 @@ 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
@@ -102,10 +101,21 @@ _TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{_MIN_REPEATS - 1},}}\s*$",
 class _RawAudio(TypedDict):
     """A waveform and its sampling rate, as `_extract_audio` reads them."""
 
-    array: npt.ArrayLike
+    array: npt.NDArray[np.float32]
     sampling_rate: int
 
 
+# `__call__` keywords this pipeline handles itself; the parent never sees them.
+_CUSTOM_PARAMS = (
+    "return_speakers",
+    "num_speakers",
+    "min_speakers",
+    "max_speakers",
+    "hf_token",
+    "user_prompt",
+)
+
+
 class _DiarizationParams(TypedDict):
     """Speaker-count hints forwarded to `NemotronDiarizer.top_speakers`."""
 
@@ -154,12 +164,8 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         # `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("max_speakers", None)
-        kwargs.pop("hf_token", None)
-        kwargs.pop("user_prompt", None)
+        for name in _CUSTOM_PARAMS:
+            kwargs.pop(name, None)
 
         return super()._sanitize_parameters(
             chunk_length_s=chunk_length_s,
@@ -260,8 +266,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         audio = self._extract_audio(inputs)
         if audio is None:
             return super().__call__(inputs, **kwargs)
-        array = np.asarray(audio["array"], dtype=np.float32)
-        sr = audio.get("sampling_rate", 16000)
+        array, sr = audio["array"], audio["sampling_rate"]
 
         bounds = chunk_bounds(array, sr)
         if len(bounds) == 1:
@@ -310,8 +315,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         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)
+        array, sr = audio["array"], audio["sampling_rate"]
         if return_speakers:
             return self._transcribe_streams(array, sr, diarization_params, **kwargs)
         return self._transcribe_aligned(array, sr, **kwargs)
@@ -367,15 +371,16 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             result["diarization_error"] = str(e)
             return result
 
+        # `top_speakers` keeps only columns that are ever active, so every
+        # mask has speech -- except its `[0]` fallback when nobody spoke.
         masks = {c: activity[:, c] > NemotronDiarizer.SEGMENT_THRESHOLD for c in keep}
-        talking = [c for c in keep if masks[c].any()]
-        if len(talking) <= 1:
+        if len(keep) <= 1:
             result = self._transcribe_aligned(array, sr, **kwargs)
-            name = f"SPEAKER_{keep.index(talking[0]) if talking else 0}"
+            name = NemotronDiarizer.speaker_names(keep)[keep[0]]
             result["words"] = [{**w, "speaker": name} for w in result["words"]]
         else:
             chunks: list[StreamChunk] = []
-            for c in talking:
+            for c in keep:
                 spans = NemotronDiarizer.sample_spans(masks[c], sr, len(array))
                 chunks.extend(
                     {"col": c, "start": s, "audio": masked_audio(array, spans, s, e)}
@@ -415,31 +420,21 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         return {"text": " ".join(w["word"] for w in words), "words": words}
 
     def _extract_audio(self, inputs: object) -> _RawAudio | None:
-        """Extract audio array from various input formats using HF utilities."""
+        """Extract a float32 audio array and its rate from various input formats."""
         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)
-            with Path(inputs).open("rb") as f:
-                audio = ffmpeg_read(f.read(), sampling_rate=16000)
-            return {"array": audio, "sampling_rate": 16000}
-        elif isinstance(inputs, bytes):
-            audio = ffmpeg_read(inputs, sampling_rate=16000)
-            return {"array": audio, "sampling_rate": 16000}
+            key = "array" if "array" in inputs else "raw" if "raw" in inputs else None
+            if key is None:
+                return None
+            array, sr = inputs[key], inputs.get("sampling_rate", 16000)
+        elif isinstance(inputs, (str, bytes)):
+            # File path or encoded bytes - decode with ffmpeg (same as HF pipeline)
+            data = Path(inputs).read_bytes() if isinstance(inputs, str) else inputs
+            array, sr = ffmpeg_read(data, sampling_rate=16000), 16000
         elif isinstance(inputs, np.ndarray):
-            return {"array": inputs, "sampling_rate": 16000}
-
-        return None
+            array, sr = inputs, 16000
+        else:
+            return None
+        return {"array": np.asarray(array, dtype=np.float32), "sampling_rate": sr}
 
     def preprocess(self, *args: Any, **preprocess_params: Any) -> Iterator[dict[str, Any]]:
         """Preprocess audio inputs for the model.
@@ -556,13 +551,11 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         top1_logprobs: list[float] = []
         top2_logprobs: list[float] = []
         if scores:
-            # Reduce each step on-device, then transfer every step in one go.
-            # Calling .item() per step would block on the device 2x per token.
-            per_step = [
-                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()
+            # One batched reduction over every step, then one device transfer.
+            logits = torch.stack([step_logits[0] for step_logits in scores]).float()
+            steps: list[list[float]] = (
+                logits.log_softmax(dim=-1).topk(k=2, dim=-1).values.cpu().tolist()
+            )
             for top1, top2 in steps:
                 top1_logprobs.append(top1)
                 top2_logprobs.append(top2)

asr_processing.py

@@ -43,6 +43,51 @@ else:
         from asr_types import AudioInput
 
 
+# The instruction the model trained on (scripts/train_collator.py); the model
+# and processor both default to it.
+DEFAULT_TRANSCRIBE_PROMPT = "Transcribe the speech to text"
+
+
+def render_audio_prompt(
+    tokenizer: PreTrainedTokenizerBase,
+    audio_token: str,
+    num_audio_tokens: int,
+    prompt: str | None,
+    text: str | None = None,
+) -> torch.Tensor:
+    """Tokenize one chat prompt carrying exactly `num_audio_tokens` placeholders.
+
+    The user turn is the placeholders, then `prompt` (if any); `text`, when
+    given, is the assistant's reply, otherwise the generation prompt is added.
+    """
+    if num_audio_tokens > 0:
+        user_content = audio_token * num_audio_tokens
+        if prompt:
+            user_content += " " + prompt
+    else:
+        user_content = prompt or ""
+
+    messages = [{"role": "user", "content": user_content}]
+    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]",
+        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.
+    ids = tokenized if isinstance(tokenized, torch.Tensor) else tokenized["input_ids"]
+    return (ids[0] if ids.dim() > 1 else ids).to(torch.long)
+
+
 def left_pad_prompt_rows(
     rows: list[torch.Tensor], tokenizer: PreTrainedTokenizerBase
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -123,7 +168,7 @@ class ASRProcessor(ProcessorMixin):
     # never added to that vocab, so it tokenizes into ordinary subwords and
     # the prompt ends up with zero scatter positions for N audio embeddings.
     AUDIO_TOKEN = "<audio>"
-    TRANSCRIBE_PROMPT = "Transcribe the speech to text"
+    TRANSCRIBE_PROMPT = DEFAULT_TRANSCRIBE_PROMPT
 
     def __init__(
         self,
@@ -153,39 +198,12 @@ class ASRProcessor(ProcessorMixin):
         self.encoder_conv_layers = encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
         self.lead_in_seconds = float(lead_in_seconds)
 
-    def _compute_encoder_output_length(self, mel_length: int) -> int:
-        """Compute encoder output length using conv layer formulas."""
-        return compute_encoder_output_length(mel_length, self.encoder_conv_layers)
-
     def _render_prompt(self, num_audio_tokens: int, text: str | None) -> torch.Tensor:
         """Tokenize one chat prompt carrying exactly `num_audio_tokens` placeholders."""
-        if num_audio_tokens > 0:
-            user_content = self.audio_token * num_audio_tokens
-            if self.TRANSCRIBE_PROMPT:
-                user_content += " " + self.TRANSCRIBE_PROMPT
-        else:
-            user_content = self.TRANSCRIBE_PROMPT or ""
-
-        messages = [{"role": "user", "content": user_content}]
-        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
-            ),
+        return render_audio_prompt(
+            self.tokenizer, self.audio_token, num_audio_tokens, self.TRANSCRIBE_PROMPT, text
         )
 
-        # apply_chat_template returns a bare tensor or a BatchEncoding/mapping.
-        ids = tokenized if isinstance(tokenized, torch.Tensor) else tokenized["input_ids"]
-        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)
@@ -226,7 +244,7 @@ class ASRProcessor(ProcessorMixin):
             extract = cast("AudioFeatureExtractor", self.feature_extractor)
             audio_inputs = extract(
                 padded_audio,
-                sampling_rate=getattr(self.feature_extractor, "sampling_rate", 16000),
+                sampling_rate=sr,
                 return_attention_mask=True,
                 return_tensors=return_tensors,
                 **kwargs,
@@ -249,11 +267,9 @@ class ASRProcessor(ProcessorMixin):
             # silently. This is the same failure `_prepare_audio_inputs`
             # documents as fixed on the model side, and it only shows up on a
             # ragged batch, so batch-1 eval never sees it.
-            mel_lengths = audio_inputs["attention_mask"].sum(dim=-1).reshape(-1)
-            token_counts = [
-                int(self.projector.get_output_length(self._compute_encoder_output_length(int(m))))
-                for m in mel_lengths
-            ]
+            mel_lengths = audio_inputs["attention_mask"].sum(dim=-1).reshape(-1).long()
+            encoder_lengths = compute_encoder_output_length(mel_lengths, self.encoder_conv_layers)
+            token_counts = self.projector.get_output_length(encoder_lengths).tolist()
 
         rows = [self._render_prompt(n, text) for n in token_counts]
         input_ids, attention_mask = self._stack_prompt_rows(rows)

diarization.py

@@ -7,7 +7,6 @@ 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:
@@ -18,7 +17,12 @@ if TYPE_CHECKING:
         Nemotron3DiarizationProcessor,
     )
 
-    from .alignment import AlignedWord
+    from .alignment import AlignedWord, get_device, to_16k
+else:
+    try:
+        from .alignment import get_device, to_16k
+    except ImportError:  # flat layout on the Hub: sibling modules, no package
+        from alignment import get_device, to_16k
 
 
 class SpeakerSegment(TypedDict):
@@ -29,15 +33,6 @@ class SpeakerSegment(TypedDict):
     end: float
 
 
-def get_device() -> torch.device:
-    """Get best available device for inference."""
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    if torch.backends.mps.is_available():
-        return torch.device("mps")
-    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.
 
@@ -161,15 +156,8 @@ class NemotronDiarizer:
     @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 = processor(to_16k(audio, sample_rate), 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
@@ -196,10 +184,15 @@ class NemotronDiarizer:
             active = sorted(sorted(active, key=lambda c: -mass[c])[:cap])
         return active or [0]
 
+    @staticmethod
+    def speaker_names(keep: list[int]) -> dict[int, str]:
+        """Column -> label: `SPEAKER_i` for the column's position in `keep`."""
+        return {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+
     @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)}
+        names = cls.speaker_names(keep)
         out: list[SpeakerSegment] = []
         for c in keep:
             for on, s, e in _label_runs(activity[:, c] > cls.SEGMENT_THRESHOLD):
@@ -257,7 +250,7 @@ class NemotronDiarizer:
         A word placed wholly in its own speaker's silenced audio (`masks`) is a
         hallucination on silence and is dropped.
         """
-        names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+        names = cls.speaker_names(keep)
         words: list[dict[str, Any]] = []
         for ch, chunk_words in zip(chunks, aligned, strict=True):
             offset = ch["start"] / sample_rate
@@ -287,7 +280,8 @@ class NemotronDiarizer:
         leaked into both masks: the copy whose speaker is more active over its
         span stays, the earlier word on a tie.
         """
-        cols = {f"SPEAKER_{i}": c for i, c in enumerate(keep)}
+        cols = {name: c for c, name in cls.speaker_names(keep).items()}
+        norms = [_norm_word(w["word"]) for w in words]
         dropped: set[int] = set()
         for i, a in enumerate(words):
             for j in range(i + 1, len(words)):
@@ -296,7 +290,7 @@ class NemotronDiarizer:
                     break
                 if i in dropped or j in dropped or a["speaker"] == b["speaker"]:
                     continue
-                if not _norm_word(a["word"]) or _norm_word(a["word"]) != _norm_word(b["word"]):
+                if not norms[i] or norms[i] != norms[j]:
                     continue
                 score_a = cls._span_activity(activity, cols[a["speaker"]], a["start"], a["end"])
                 score_b = cls._span_activity(activity, cols[b["speaker"]], b["start"], b["end"])

handler.py

@@ -1,27 +1,49 @@
 """Custom inference handler for HuggingFace Inference Endpoints."""
 
+import base64
+import binascii
 import os
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-import nltk
-
 if TYPE_CHECKING:
+    from .alignment import QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_pipeline import ASRPipeline
+    from .diarization import NemotronDiarizer
     from .diarization import get_device as _best_device
 else:
     try:
         # For remote execution, imports are relative
+        from .alignment import QwenForcedAligner
         from .asr_modeling import ASRModel
         from .asr_pipeline import ASRPipeline
+        from .diarization import NemotronDiarizer
         from .diarization import get_device as _best_device
     except ImportError:
         # For local execution, imports are not relative
+        from alignment import QwenForcedAligner
         from asr_modeling import ASRModel
         from asr_pipeline import ASRPipeline
+        from diarization import NemotronDiarizer
         from diarization import get_device as _best_device
 
 
+def decode_inputs(inputs: Any) -> Any:
+    """A JSON request's base64 audio as bytes; paths, raw bytes and arrays unchanged.
+
+    The toolkit hands `inputs` over as raw bytes for an `audio/*` body, but a
+    JSON body can only carry audio as a base64 string, which the pipeline
+    would otherwise read as a file path.
+    """
+    if not isinstance(inputs, str) or Path(inputs).exists():
+        return inputs
+    try:
+        return base64.b64decode(inputs, validate=True)
+    except (binascii.Error, ValueError):
+        return inputs
+
+
 class EndpointHandler:
     """HuggingFace Inference Endpoints handler for ASR model.
 
@@ -35,13 +57,6 @@ 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")
 
         # `ASRModel.from_pretrained` constructs its own submodules and forwards
@@ -66,15 +81,24 @@ class EndpointHandler:
             tokenizer=self.model.tokenizer,
             device=self.device,
         )
+        # Load the forced aligner (return_timestamps) and Nemotron diarizer
+        # (return_speakers) at boot, as the Spaces demo does, so the first
+        # request that asks for them doesn't pay the download and load.
+        QwenForcedAligner.get_instance()
+        NemotronDiarizer.get_instance()
 
     def __call__(self, data: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
         """Process an inference request.
 
         Args:
-            data: Request data containing 'inputs' (audio path/bytes) and optional 'parameters'
+            data: Request data containing 'inputs' (audio path, bytes or base64
+                string) and optional 'parameters' for the pipeline, e.g.
+                `return_timestamps`, `return_speakers`, `num_speakers`,
+                `max_speakers` -- the same options as the Spaces demo
 
         Returns:
-            Transcription result with 'text' key
+            Transcription result with 'text' key; 'words' with timestamps, plus
+            speaker labels and 'speaker_segments' when diarizing
         """
         inputs = data.get("inputs")
         if inputs is None:
@@ -84,5 +108,5 @@ class EndpointHandler:
         # 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)
+        result: dict[str, Any] | list[dict[str, Any]] = self.pipe(decode_inputs(inputs), **params)
         return result

requirements.txt

@@ -5,6 +5,14 @@
 # Speaker diarization (return_speakers=True) uses Nemotron-3-Diarization, which is in
 # transformers main after 5.17 -- pinned to the commit the pipeline was verified on.
 transformers @ git+https://github.com/huggingface/transformers@91ea1192a00e0ebb684400434375fbd2345c8449
-torch
-torchaudio
+# Pinned floor, not just "torch": the Inference Endpoints image ships an older
+# torch that satisfies a bare requirement, and transformers disables PyTorch
+# below 2.5 ("Could not import module 'WhisperForConditionalGeneration'").
+torch>=2.5
+torchaudio>=2.5
 peft
+safetensors
+numpy
+# Fast path for Qwen3.5's gated-delta-rule layers (transformers falls back to a
+# slower reference implementation without it). Pure Triton, nothing to compile.
+flash-linear-attention>=0.2.2