Update custom model files, README, and requirements

asr_pipeline.py

@@ -2,7 +2,7 @@
 
 import re
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -54,6 +54,19 @@ def chunk_bounds(
     return bounds
 
 
+# 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: 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
+
+
 _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", flags=re.DOTALL)
 _MIN_REPEATS = 3
 _TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{_MIN_REPEATS - 1},}}$")
@@ -172,11 +185,11 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
         bounds = chunk_bounds(array, sr)
         if len(bounds) == 1:
-            return super().__call__({"raw": array, "sampling_rate": sr}, **kwargs)
+            return self._transcribe_chunk(array, sr, **kwargs)
         result: dict[str, Any] = {}
         texts = []
         for s, e in bounds:
-            out = super().__call__({"raw": array[s:e], "sampling_rate": sr}, **kwargs)
+            out = self._transcribe_chunk(array[s:e], sr, **kwargs)
             texts.append(out["text"])
             # Per-step logprobs (output_scores=True) run on across chunks.
             for key in ("top1_logprob", "top2_logprob"):
@@ -185,6 +198,15 @@ 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,
@@ -209,12 +231,7 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         sr = audio.get("sampling_rate", 16000)
 
         bounds = chunk_bounds(array, sr)
-        texts = [
-            super(ASRPipeline, self).__call__({"raw": array[s:e], "sampling_rate": sr}, **kwargs)[
-                "text"
-            ]
-            for s, e in bounds
-        ]
+        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:

diarization.py

@@ -806,27 +806,45 @@ class NemotronDiarizer:
         return sorted(out, key=lambda seg: (seg["start"], seg["end"]))
 
     @classmethod
-    def assign_speakers_to_words(
-        cls, words: list[dict], activity: np.ndarray, keep: list[int]
-    ) -> list[dict]:
-        """Give each word the kept speaker most active over its span.
+    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.
 
-        A word where no kept speaker reaches MIN_WORD_ACTIVITY (aligner put it
-        in a pause) takes the previous word's speaker -- the next word's at
-        the start of the recording.
+        The kept speaker most active over the word's span. A word with no
+        time, or where no kept speaker reaches MIN_WORD_ACTIVITY (it sits in
+        a pause), takes the previous word's speaker -- the next word's at the
+        start of the recording.
         """
-        names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
         cols = np.asarray(keep)
-        picked: list[str | None] = []
-        for word in words:
-            lo = min(int(word["start"] / cls.FRAME_S), len(activity) - 1)
-            hi = max(lo + 1, min(int(np.ceil(word["end"] / cls.FRAME_S)), len(activity)))
+        picked: list[int | None] = []
+        for start, end in spans:
+            if start is None or end is None:
+                picked.append(None)
+                continue
+            lo = min(int(start / cls.FRAME_S), len(activity) - 1)
+            hi = max(lo + 1, min(int(np.ceil(end / cls.FRAME_S)), len(activity)))
             mean = activity[lo:hi, cols].mean(axis=0)
             best = int(mean.argmax())
-            picked.append(names[keep[best]] if mean[best] >= cls.MIN_WORD_ACTIVITY else None)
+            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 names[keep[0]]
-        for word, p in zip(words, picked):
+        last = known[0] if known else keep[0]
+        out = []
+        for p in picked:
             last = p if p is not None else last
-            word["speaker"] = last
+            out.append(last)
+        return out
+
+    @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