Update custom model files, README, and requirements

alignment.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import itertools
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, TypedDict
 
@@ -126,9 +127,7 @@ class QwenForcedAligner:
         if not todo:  # nothing to time: don't load the model
             return results
         model, processor = cls.get_instance()
-        # Slices, not itertools.batched (3.12+): the Hub copy runs on Python 3.11.
-        for start in range(0, len(todo), cls.BATCH_SIZE):
-            batch = todo[start : start + cls.BATCH_SIZE]
+        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],

asr_config.py

@@ -76,9 +76,10 @@ def compute_encoder_output_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)
+    # Serialized configs carry their own `model_type`, which would collide with
+    # for_model's positional argument; the config class already sets it.
+    kwargs = {key: value for key, value in values.items() if key != "model_type"}
+    return transformers.AutoConfig.for_model(model_type, **kwargs)
 
 
 class ASRConfig(transformers.PretrainedConfig):
@@ -396,10 +397,12 @@ class ASRConfig(transformers.PretrainedConfig):
         # The pipeline looks for config.encoder._name_or_path for feature extractor
         self.encoder = self.audio_config
 
+        # The Hub infers the model's task from the first model class listed, so
+        # AutoModelForSpeechSeq2Seq must precede AutoModel (else: feature-extraction).
         self.auto_map = {
             "AutoConfig": "asr_config.ASRConfig",
-            "AutoModel": "asr_modeling.ASRModel",
             "AutoModelForSpeechSeq2Seq": "asr_modeling.ASRModel",
+            "AutoModel": "asr_modeling.ASRModel",
             "AutoProcessor": "asr_processing.ASRProcessor",
         }
         self.custom_pipelines = {

asr_layers.py

@@ -115,11 +115,10 @@ def find_encoder_layer_stack(encoder: nn.Module) -> tuple[str, nn.ModuleList] |
     """
     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
+        try:
+            obj = encoder.get_submodule(path)
+        except AttributeError:
+            continue
         if isinstance(obj, nn.ModuleList) and len(obj) > 0:
             return path, obj
     return None

asr_modeling.py

@@ -6,7 +6,7 @@ import json
 import logging
 import shutil
 from collections.abc import Callable, Iterator
-from copy import copy
+from copy import deepcopy
 from os import PathLike
 from pathlib import Path
 from threading import Thread
@@ -1852,25 +1852,18 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
         self._assert_sdpa_safe_on_mps(attention_mask)
 
-        # transformers v5 deprecates passing generation flags as kwargs when a
-        # `generation_config` is also passed — the kwargs get silently dropped.
-        # 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
+        # Fold generation flags (max_new_tokens, num_beams, output_scores, ...)
+        # into a private copy of generation_config instead of passing them
+        # alongside it: transformers v5 still applies such kwargs but warns that
+        # combining them with `generation_config` is deprecated. update() returns
+        # the kwargs that aren't config fields (streamer, logits_processor, ...),
+        # which still go to generate() directly.
+        gen_cfg = deepcopy(self.generation_config)
+        generate_kwargs = gen_cfg.update(**generate_kwargs)
+        # output_scores requires return_dict_in_generate for HF generate to
+        # actually populate .scores on the output object.
+        if gen_cfg.output_scores:
+            gen_cfg.return_dict_in_generate = True
 
         lm_inputs = self._lm_generate_inputs(input_ids, inputs_embeds)
         output = self.language_model.generate(
@@ -1946,10 +1939,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Prepare generation kwargs
         # Same prompt arguments as generate(); with input_ids present the
         # streamer's skip_prompt drops the echoed prompt tokens.
+        # Generation flags go into a copy of generation_config, as in generate().
+        gen_cfg = deepcopy(self.generation_config)
+        generate_kwargs = gen_cfg.update(**generate_kwargs)
         gen_kwargs: dict[str, Any] = {
             **self._lm_generate_inputs(input_ids, inputs_embeds),
             "attention_mask": attention_mask,
-            "generation_config": self.generation_config,
+            "generation_config": gen_cfg,
             "streamer": streamer,
             **generate_kwargs,
         }
@@ -2042,30 +2038,31 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Don't save embedding layers - the <audio> token embedding is never used
         # (it's replaced with projected audio embeddings before the LLM sees it)
         if hasattr(self.language_model, "peft_config"):
-            self.language_model.save_pretrained(save_dir, save_embedding_layers=False)
-
-            # Clear base_model_name_or_path in adapter_config.json to prevent HF pipeline
-            # from redirecting to the base LLM repo (like Qwen) which breaks feature
-            # extractor loading for multimodal models. If a repo_id is provided, use that
-            # so the model can be loaded directly from the Hub.
-            adapter_config_path = save_dir / "adapter_config.json"
-            if adapter_config_path.exists():
-                with adapter_config_path.open() as f:
-                    adapter_config = json.load(f)
-
-                # Use repo_id if available, otherwise clear to prevent redirect.
-                # Use empty string instead of None to avoid str(None) -> "None" bug
-                # in some transformers/PEFT versions.
-                repo_id = (
-                    kwargs.get("repo_id")
-                    or kwargs.get("push_to_hub_model_id")
-                    or getattr(self.config, "pretrained_model_path", None)
-                    or ""  # Use empty string instead of None
-                )
-                adapter_config["base_model_name_or_path"] = repo_id
-
-                with adapter_config_path.open("w") as f:
-                    json.dump(adapter_config, f, indent=2)
+            # Override base_model_name_or_path in the saved adapter_config.json to
+            # prevent HF pipeline from redirecting to the base LLM repo (like Qwen),
+            # which breaks feature extractor loading for multimodal models. If a
+            # repo_id is provided, use that so the model can be loaded directly
+            # from the Hub. Use empty string instead of None: PeftModel.save_pretrained
+            # fills the field with the base LLM's name only when it is None, and an
+            # empty string also avoids the str(None) -> "None" bug in some
+            # transformers/PEFT versions.
+            repo_id = (
+                kwargs.get("repo_id")
+                or kwargs.get("push_to_hub_model_id")
+                or getattr(self.config, "pretrained_model_path", None)
+                or ""
+            )
+            peft_configs = list(self.language_model.peft_config.values())
+            # Set on the live configs only for the save, then restore, so the
+            # in-memory model keeps its real base model name.
+            original_names = [cfg.base_model_name_or_path for cfg in peft_configs]
+            try:
+                for cfg in peft_configs:
+                    cfg.base_model_name_or_path = repo_id
+                self.language_model.save_pretrained(save_dir, save_embedding_layers=False)
+            finally:
+                for cfg, name in zip(peft_configs, original_names, strict=True):
+                    cfg.base_model_name_or_path = name
 
         # Add processor auto_map to preprocessor_config.json
         config_path = save_dir / "preprocessor_config.json"

asr_pipeline.py

@@ -1,7 +1,7 @@
 """ASR pipeline for audio-to-text transcription with optional timestamps and diarization."""
 
 import re
-from collections.abc import Iterator, MutableMapping
+from collections.abc import Callable, Iterator, MutableMapping, Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, TypedDict, cast
 
@@ -16,86 +16,57 @@ if TYPE_CHECKING:
 
     from .alignment import QwenForcedAligner
     from .asr_modeling import ASRModel
-    from .asr_processing import prepend_lead_in
-    from .diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
+    from .asr_processing import (
+        audible_chunks,
+        chunk_bounds,
+        collate_chunks,
+        is_silent,
+        prepend_lead_in,
+    )
+    from .asr_types import PreparedChunk
+    from .diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
 else:
     try:
         from .alignment import QwenForcedAligner
         from .asr_modeling import ASRModel
-        from .asr_processing import prepend_lead_in
-        from .diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
+        from .asr_processing import (
+            audible_chunks,
+            chunk_bounds,
+            collate_chunks,
+            is_silent,
+            prepend_lead_in,
+        )
+        from .asr_types import PreparedChunk
+        from .diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
     except ImportError:  # flat layout on the Hub: sibling modules, no package
         from alignment import QwenForcedAligner
         from asr_modeling import ASRModel
-        from asr_processing import prepend_lead_in
-        from diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
+        from asr_processing import (
+            audible_chunks,
+            chunk_bounds,
+            collate_chunks,
+            is_silent,
+            prepend_lead_in,
+        )
+        from asr_types import PreparedChunk
+        from diarization import NemotronDiarizer, StreamChunk, masked_audio, stream_chunks
 
 # Re-export for backwards compatibility
 __all__ = [
     "ASRPipeline",
     "NemotronDiarizer",
+    "PreparedChunk",
     "QwenForcedAligner",
+    "collate_chunks",
 ]
 
-# Audio is transcribed in chunks cut at the quietest point between
-# these lengths. The model trained on clips of at most 19 s; 18 leaves room for
-# the inference lead-in. Short clips are one chunk, so their text is unchanged.
-CHUNK_MAX_S = 18.0
-CHUNK_MIN_S = 8.0
-
-
-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)
-        k = (hi - lo) // frame
-        rms = np.sqrt(np.mean(np.square(audio[lo : lo + k * frame].reshape(k, frame)), axis=1))
-        cut = lo + int(np.argmin(rms)) * frame + frame // 2
-        bounds.append((start, cut))
-        start = cut
-    bounds.append((start, n))
-    return bounds
-
-
-def stream_chunks(
-    audio: npt.NDArray[np.float32],
-    spans: list[tuple[int, int]],
-    sample_rate: int,
-    max_s: float = CHUNK_MAX_S,
-) -> list[tuple[int, int]]:
-    """Ranges of at most `max_s` over one speaker's `spans`; long turns cut by `chunk_bounds`."""
-    pieces: list[tuple[int, int]] = []
-    for s, e in spans:
-        pieces.extend((s + a, s + b) for a, b in chunk_bounds(audio[s:e], sample_rate, max_s))
-    return pack_spans(pieces, int(max_s * sample_rate))
-
-
-# Below this RMS (-100 dBFS) a chunk is digital silence: exact zeros, as in
-# edited or remixed recordings. Given one, the model answers with a memorized
-# training sentence ("The film was directed by the same director who directed
-# 'The Man with the Moustache'") -- ten such chunks cost 2.3 WER on one AMI
-# meeting -- so it is skipped. Quiet real speech sits near -60 dBFS.
-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
-
 
 _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", flags=re.DOTALL)
 _MIN_REPEATS = 3
+_LOOSE_MIN_REPEATS = 6
 _TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{_MIN_REPEATS - 1},}}$")
 _TRAILING_WORD_RE = re.compile(rf"\b(\w+)(?:\s+\1){{{_MIN_REPEATS - 1},}}\s*$", re.IGNORECASE)
+_NON_WORD_RE = re.compile(r"\W+")
 
 
 class _RawAudio(TypedDict):
@@ -128,6 +99,13 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
     model: ASRModel
 
+    #: Chunks per batched `generate` call when this pipeline runs them itself.
+    max_batch_size: int = 16
+    #: Runs prepared chunks through the model, returning one text each. None
+    #: means `generate_prepared` in batches of `max_batch_size`; a server
+    #: swaps in a queue that batches chunks across concurrent requests.
+    chunk_runner: "Callable[[list[PreparedChunk]], list[str]] | None" = None
+
     def __init__(self, model: ASRModel, **kwargs: Any) -> None:
         """Initialize ASR pipeline.
 
@@ -269,12 +247,17 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         array, sr = audio["array"], audio["sampling_rate"]
 
         bounds = chunk_bounds(array, sr)
+        if not kwargs:
+            texts = self._transcribe_chunks(audible_chunks(array, bounds, sr), sr)
+            return {"text": " ".join(t for t in texts if t)}
+        # Generate kwargs (e.g. output_scores for eval confidence) need a
+        # one-chunk batch: `_forward` refuses per-item scores for more.
         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)
+        for chunk in audible_chunks(array, bounds, sr):
+            out = self._transcribe_chunk(chunk, sr, **kwargs)
             texts.append(out["text"])
             # Per-step logprobs (output_scores=True) run on across chunks.
             for key in ("top1_logprob", "top2_logprob"):
@@ -283,6 +266,61 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         result["text"] = " ".join(t for t in texts if t)
         return result
 
+    def _transcribe_chunks(
+        self, chunks: Sequence[npt.NDArray[np.float32]], sample_rate: int, **kwargs: Any
+    ) -> list[str]:
+        """One text per chunk, batched; digital silence is "" without the model.
+
+        All of a request's chunks go to the model together (`chunk_runner`)
+        rather than one call each, so a long recording, or diarization's
+        per-speaker streams, costs a few batched `generate` calls instead of
+        one per 18 s. Generate kwargs fall back to one call per chunk.
+        """
+        texts = [""] * len(chunks)
+        todo = [i for i, chunk in enumerate(chunks) if not is_silent(chunk)]
+        if kwargs:
+            for i in todo:
+                texts[i] = self._transcribe_chunk(chunks[i], sample_rate, **kwargs)["text"]
+            return texts
+        prepared = [self.prepare_chunk(chunks[i], sample_rate) for i in todo]
+        run = self.chunk_runner or self._generate_in_batches
+        for i, text in zip(todo, run(prepared) if prepared else [], strict=True):
+            texts[i] = text
+        return texts
+
+    def prepare_chunk(self, chunk: npt.NDArray[np.float32], sample_rate: int) -> PreparedChunk:
+        """A chunk's encoder inputs on CPU: lead-in, then the feature extractor.
+
+        The same `preprocess` a one-chunk call runs, so batched and one-at-a-time
+        decoding see identical features.
+        """
+        item = next(self.preprocess({"raw": chunk, "sampling_rate": sample_rate}))
+        return {"input_features": item["input_features"], "attention_mask": item["attention_mask"]}
+
+    def _generate_in_batches(self, prepared: list[PreparedChunk]) -> list[str]:
+        """`generate_prepared` over `prepared` in batches of `max_batch_size`."""
+        size = max(1, self.max_batch_size)
+        batches = (prepared[start : start + size] for start in range(0, len(prepared), size))
+        return [text for batch in batches for text in self.generate_prepared(batch)]
+
+    @torch.inference_mode()
+    def generate_prepared(
+        self, prepared: Sequence[PreparedChunk], **generate_kwargs: Any
+    ) -> list[str]:
+        """One batched `generate` (kwargs passed on) over prepared chunks; one text each.
+
+        Rows are left-padded prompts sized to each chunk's own audio tokens, so a
+        ragged batch decodes like its rows alone, up to padding numerics.
+        """
+        batch = collate_chunks(prepared)
+        tokens = self.model.generate(
+            input_features=batch["input_features"].to(self.model.device),
+            audio_attention_mask=batch["attention_mask"].to(self.model.device),
+            **generate_kwargs,
+        )
+        assert torch.is_tensor(tokens)  # no output_scores: plain token ids
+        return [self.postprocess({"tokens": row})["text"] for row in tokens.cpu()]
+
     def _transcribe_chunk(
         self, chunk: npt.NDArray[np.float32], sample_rate: int, **kwargs: Any
     ) -> dict[str, Any]:
@@ -305,7 +343,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
     ) -> dict[str, Any]:
         """Transcribe in chunks and time every word; per-speaker streams if requested.
 
-        Chunks of 8-18 s cut at quiet points are transcribed one by one and each
+        Chunks of 8-18 s cut at quiet points are transcribed in batches and each
         is aligned against its own transcript (batched), then offset onto the
         recording's timeline -- so audio of any length gets every word timed,
         and the model never sees a clip longer than it trained on. Speakers
@@ -325,9 +363,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
     ) -> dict[str, Any]:
         """`{"text", "words"}` for the whole recording: chunk, transcribe, align, offset."""
         bounds = chunk_bounds(array, sr)
-        texts: list[str] = [
-            self._transcribe_chunk(array[s:e], sr, **kwargs)["text"] for s, e in bounds
-        ]
+        texts = self._transcribe_chunks(audible_chunks(array, bounds, sr), sr, **kwargs)
         result: dict[str, Any] = {"text": " ".join(t for t in texts if t)}
 
         try:
@@ -386,7 +422,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                     {"col": c, "start": s, "audio": masked_audio(array, spans, s, e)}
                     for s, e in stream_chunks(array, spans, sr)
                 )
-            texts = [self._transcribe_chunk(ch["audio"], sr, **kwargs)["text"] for ch in chunks]
+            texts = self._transcribe_chunks([ch["audio"] for ch in chunks], sr, **kwargs)
             result = self._align_streams(chunks, texts, sr, masks, keep)
             if result["words"]:
                 result["words"] = NemotronDiarizer.dedupe_words(result["words"], activity, keep)
@@ -655,6 +691,11 @@ def _truncate_repetitions(text: str) -> str:
     - Repeated phrases: "i am sorry i am sorry i am sorry" -> "i am sorry"
     - Repeated characters: "444444" -> "4"
 
+    Words match ignoring case. A run of `_LOOSE_MIN_REPEATS` also matches
+    ignoring punctuation, since the model punctuates inside a loop ("Yeah Yeah
+    Yeah. Yeah. Yeah,"); shorter punctuated runs are real speech ("No, no,
+    no,"), and collapsing those cost 40 errors over 47 rows of the served eval.
+
     Args:
         text: Input text to process
 
@@ -678,7 +719,7 @@ def _truncate_repetitions(text: str) -> str:
     # to be possible. set(window) == window means all unique → no repetition.
     window = words[-_MIN_REPEATS * 2 :]
     if len(set(window)) == len(window):
-        return text
+        return _truncate_punctuated_loop(text)
 
     for phrase_len in range(2, min(21, len(words) // _MIN_REPEATS + 1)):
         phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
@@ -691,4 +732,20 @@ def _truncate_repetitions(text: str) -> str:
             text = (match.group(1) + match.group(2)).strip()
             break
 
+    return _truncate_punctuated_loop(text)
+
+
+def _truncate_punctuated_loop(text: str) -> str:
+    """Collapse a trailing 1-20 word phrase run `_LOOSE_MIN_REPEATS`+ times, punctuation ignored."""
+    words = text.split()
+    keys = [_NON_WORD_RE.sub("", w).lower() for w in words]
+    end = len(words)
+    # Shortest phrase first, so "Yeah Yeah. Yeah," stops at "Yeah".
+    for n in range(1, min(21, end // _LOOSE_MIN_REPEATS + 1)):
+        phrase = keys[end - n : end]
+        reps = 1
+        while end - (reps + 1) * n >= 0 and keys[end - (reps + 1) * n : end - reps * n] == phrase:
+            reps += 1
+        if reps >= _LOOSE_MIN_REPEATS:
+            return " ".join(words[: end - (reps - 1) * n])
     return text

asr_processing.py

@@ -1,7 +1,7 @@
 """Processor that turns raw audio (and optional text) into model inputs."""
 
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any, ClassVar, cast, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -22,7 +22,7 @@ if TYPE_CHECKING:
         ConvLayerSpec,
         compute_encoder_output_length,
     )
-    from .asr_types import AudioFeatureExtractor, AudioInput, Waveform
+    from .asr_types import AudioFeatureExtractor, AudioInput, PreparedChunk, Waveform
     from .projectors import MLPAudioProjector
 else:
     try:
@@ -32,7 +32,7 @@ else:
             ConvLayerSpec,
             compute_encoder_output_length,
         )
-        from .asr_types import AudioInput
+        from .asr_types import AudioInput, PreparedChunk
     except ImportError:  # flat layout on the Hub: sibling modules, no package
         from asr_config import (
             DEFAULT_ENCODER_CONV_LAYERS,
@@ -40,7 +40,29 @@ else:
             ConvLayerSpec,
             compute_encoder_output_length,
         )
-        from asr_types import AudioInput
+        from asr_types import AudioInput, PreparedChunk
+
+
+def collate_chunks(prepared: Sequence[PreparedChunk]) -> PreparedChunk:
+    """Pad prepared chunks to the longest and stack them into one batch.
+
+    The time axis is whichever feature axis matches the mask's length (Granite's
+    features are `(1, T, D)`, Whisper's `(1, D, T)`); padded frames are zeros
+    with a 0 in the mask, which the encoder honours (`encoder_attention_mask`).
+    """
+    longest = max(int(p["attention_mask"].shape[-1]) for p in prepared)
+    features: list[torch.Tensor] = []
+    masks: list[torch.Tensor] = []
+    for p in prepared:
+        feats, mask = p["input_features"], p["attention_mask"]
+        length = int(mask.shape[-1])
+        time_axis = 1 if feats.shape[1] == length else feats.dim() - 1
+        pad = longest - length
+        # F.pad lists (left, right) pairs from the LAST axis backwards.
+        spec = [0, 0] * (feats.dim() - 1 - time_axis) + [0, pad]
+        features.append(torch.nn.functional.pad(feats, spec))
+        masks.append(torch.nn.functional.pad(mask, (0, pad)))
+    return {"input_features": torch.cat(features), "attention_mask": torch.cat(masks)}
 
 
 # The instruction the model trained on (scripts/train_collator.py); the model
@@ -111,15 +133,12 @@ def left_pad_prompt_rows(
     return input_ids, attention_mask
 
 
-ScalarT = TypeVar("ScalarT", bound=np.generic)
-
-
 @overload
-def prepend_lead_in(
+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(
+def prepend_lead_in[ScalarT: np.generic](
     audio: list[npt.NDArray[ScalarT]], sampling_rate: int, seconds: float | None
 ) -> list[npt.NDArray[ScalarT]]: ...
 @overload
@@ -156,6 +175,81 @@ def prepend_lead_in(audio: AudioInput, sampling_rate: int, seconds: float | None
     return padded
 
 
+# Audio is transcribed in chunks cut at the quietest point between
+# these lengths. The model trained on clips of at most 19 s; 18 leaves room for
+# the inference lead-in. Short clips are one chunk, so their text is unchanged.
+CHUNK_MAX_S = 18.0
+CHUNK_MIN_S = 8.0
+
+
+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)
+        cut = lo + int(np.argmin(_frame_rms(audio[lo:hi], frame))) * frame + frame // 2
+        bounds.append((start, cut))
+        start = cut
+    bounds.append((start, n))
+    return bounds
+
+
+def _frame_rms(audio: npt.NDArray[np.float32], frame: int) -> npt.NDArray[np.float32]:
+    """RMS of each whole `frame`-sample frame of `audio` (a trailing partial frame is dropped)."""
+    k = len(audio) // frame
+    return np.sqrt(np.mean(np.square(audio[: k * frame].reshape(k, frame)), axis=1))
+
+
+# A chunk whose loudest 100 ms frame is this far below the recording's speech
+# level (its 95th-percentile frame) holds no speech, only the room tone after
+# the talker stopped. Decoded, such a tail comes back as a memorized sentence
+# ("The film was directed by the director of the same name.", 0.7 WER on
+# CommonVoice) or a stray "the"/"ok". On the cached eval clips over 18 s every
+# noise-only chunk sat at -36 dB or below and every chunk with speech at
+# -17.5 dB or above; -30 keeps the wider margin on the speech side.
+QUIET_CHUNK_DB = -30.0
+
+
+def audible_chunks(
+    audio: npt.NDArray[np.float32], bounds: list[tuple[int, int]], sample_rate: int
+) -> list[npt.NDArray[np.float32]]:
+    """The chunks of `audio` at `bounds`, those quieter than `QUIET_CHUNK_DB` emptied.
+
+    An empty chunk `is_silent`, so it transcribes as "" without the model. A
+    one-chunk recording is never emptied: its loudest frame is its own level.
+    """
+    chunks = [audio[s:e] for s, e in bounds]
+    if len(chunks) < 2:
+        return chunks
+    frame = int(0.1 * sample_rate)
+    floor = np.percentile(_frame_rms(audio, frame), 95) * 10 ** (QUIET_CHUNK_DB / 20)
+    return [
+        chunk if len(chunk) >= frame and _frame_rms(chunk, frame).max() >= floor else chunk[:0]
+        for chunk in chunks
+    ]
+
+
+# Below this RMS (-100 dBFS) a chunk is digital silence: exact zeros, as in
+# edited or remixed recordings. Given one, the model answers with a memorized
+# training sentence ("The film was directed by the same director who directed
+# 'The Man with the Moustache'") -- ten such chunks cost 2.3 WER on one AMI
+# meeting -- so it is skipped. Quiet real speech sits near -60 dBFS.
+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
+
+
 class ASRProcessor(ProcessorMixin):
     """Processor for Whisper-based ASR models."""
 

asr_types.py

@@ -21,6 +21,14 @@ AudioInput = Waveform | Sequence[Waveform]
 # A `state_dict(destination=...)` mapping, returned as the caller's own type.
 StateDictT = TypeVar("StateDictT", bound=dict[str, Any])
 
+
+class PreparedChunk(TypedDict):
+    """One chunk's encoder inputs, as `ASRPipeline.prepare_chunk` builds them (batch of 1)."""
+
+    input_features: torch.Tensor
+    attention_mask: torch.Tensor
+
+
 if TYPE_CHECKING:
 
     class GenerativeDecoder(PreTrainedModel, GenerationMixin):

diarization.py

@@ -18,11 +18,14 @@ if TYPE_CHECKING:
     )
 
     from .alignment import AlignedWord, get_device, to_16k
+    from .asr_processing import CHUNK_MAX_S, chunk_bounds
 else:
     try:
         from .alignment import get_device, to_16k
+        from .asr_processing import CHUNK_MAX_S, chunk_bounds
     except ImportError:  # flat layout on the Hub: sibling modules, no package
         from alignment import get_device, to_16k
+        from asr_processing import CHUNK_MAX_S, chunk_bounds
 
 
 class SpeakerSegment(TypedDict):
@@ -83,6 +86,19 @@ def pack_spans(spans: list[tuple[int, int]], limit: int) -> list[tuple[int, int]
     return ranges
 
 
+def stream_chunks(
+    audio: npt.NDArray[np.float32],
+    spans: list[tuple[int, int]],
+    sample_rate: int,
+    max_s: float = CHUNK_MAX_S,
+) -> list[tuple[int, int]]:
+    """Ranges of at most `max_s` over one speaker's `spans`; long turns cut by `chunk_bounds`."""
+    pieces: list[tuple[int, int]] = []
+    for s, e in spans:
+        pieces.extend((s + a, s + b) for a, b in chunk_bounds(audio[s:e], sample_rate, max_s))
+    return pack_spans(pieces, int(max_s * sample_rate))
+
+
 class StreamChunk(TypedDict):
     """One stretch of one speaker's masked audio: kept column, absolute start, samples."""
 

requirements.txt

@@ -5,19 +5,6 @@
 # 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
-# The Inference Endpoints image ships an older torch, torchvision and torchaudio.
-# Pin all three as one matching set (torch 2.8, as in poetry.lock): a bare
-# "torch" keeps the image's version, which this transformers rejects (< 2.5),
-# and upgrading torch alone leaves the image's torchvision built for the old
-# one, so transformers' torchvision import fails ("operator torchvision::nms
-# does not exist", surfacing as "Could not import module
-# 'WhisperForConditionalGeneration'").
-torch==2.8.0
-torchvision==0.23.0
-torchaudio==2.8.0
+torch
+torchaudio
 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