Update custom model files, README, and requirements

README.md

@@ -127,7 +127,8 @@ for word in result["words"]:
 
 Pass `return_speakers=True` to label every word with who said it. It also turns on word timestamps.
 It works on recordings of any length, keeps speaker labels consistent across hour-long meetings, and
-handles overlapping speech and up to 8 speakers.
+handles up to 8 speakers. Each speaker is transcribed separately on audio where everyone else is
+silenced, so overlapping speech is only partly handled.
 
 Diarization needs transformers installed from `main` until the next release:
 
@@ -154,8 +155,8 @@ for w in result["words"]:
 ```
 
 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:
+you know it, pass `num_speakers` (exact) or `max_speakers` (an upper bound); voices beyond the cap
+are not transcribed:
 
 ```python
 result = pipe("call.wav", return_speakers=True, num_speakers=2)

asr_pipeline.py

@@ -18,18 +18,18 @@ if TYPE_CHECKING:
     from .alignment import QwenForcedAligner
     from .asr_modeling import ASRModel
     from .asr_processing import prepend_lead_in
-    from .diarization import NemotronDiarizer
+    from .diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
 else:
     try:
         from .alignment import QwenForcedAligner
         from .asr_modeling import ASRModel
         from .asr_processing import prepend_lead_in
-        from .diarization import NemotronDiarizer
+        from .diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
     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
+        from diarization import NemotronDiarizer, StreamChunk, masked_audio, pack_spans
 
 # Re-export for backwards compatibility
 __all__ = [
@@ -67,6 +67,19 @@ def chunk_bounds(
     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
@@ -172,7 +185,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         Args:
             inputs: Audio input (file path, dict with array/sampling_rate, etc.)
             return_timestamps: If True, return word-level timestamps (Qwen3-ForcedAligner)
-            return_speakers: If True, label every word with a speaker (Nemotron-3-Diarization)
+            return_speakers: If True, transcribe each Nemotron-3-Diarization speaker on
+                audio with the others silenced and label words by speaker (see
+                `_transcribe_streams`). Implies return_timestamps.
             user_prompt: Custom transcription prompt (default: "Transcribe: ")
             num_speakers: Exact number of speakers, if known
             max_speakers: Upper bound on the number of speakers
@@ -283,14 +298,13 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         diarization_params: _DiarizationParams,
         **kwargs: Any,
     ) -> dict[str, Any]:
-        """Transcribe in chunks, time every word, and label speakers if requested.
+        """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
         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 come
-        from ONE Nemotron pass over the whole recording, so labels are
-        recording-level: chunking never touches speaker identity.
+        and the model never sees a clip longer than it trained on. Speakers
+        come from `_transcribe_streams`.
         """
         audio = self._extract_audio(inputs)
         if audio is None:
@@ -298,7 +312,14 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             raise ValueError(msg)
         array = np.asarray(audio["array"], dtype=np.float32)
         sr = audio.get("sampling_rate", 16000)
+        if return_speakers:
+            return self._transcribe_streams(array, sr, diarization_params, **kwargs)
+        return self._transcribe_aligned(array, sr, **kwargs)
 
+    def _transcribe_aligned(
+        self, array: npt.NDArray[np.float32], sr: int, **kwargs: Any
+    ) -> dict[str, Any]:
+        """`{"text", "words"}` for the whole recording: chunk, transcribe, align, offset."""
         bounds = chunk_bounds(array, sr)
         texts: list[str] = [
             self._transcribe_chunk(array[s:e], sr, **kwargs)["text"] for s, e in bounds
@@ -318,21 +339,81 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         except Exception as e:
             result["words"] = []
             result["timestamp_error"] = str(e)
+        return result
 
-        if return_speakers:
-            try:
-                activity = NemotronDiarizer.activity(array, sample_rate=sr)
-                keep = NemotronDiarizer.top_speakers(activity, **diarization_params)
-                result["speaker_segments"] = NemotronDiarizer.segments(activity, keep)
-                result["words"] = NemotronDiarizer.assign_speakers_to_words(
-                    result["words"], activity, keep
+    def _transcribe_streams(
+        self,
+        array: npt.NDArray[np.float32],
+        sr: int,
+        diarization_params: _DiarizationParams,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Per-speaker masked ASR, a zero-shot port of NeMo's `masked_asr`.
+
+        Nemotron runs first, one pass over the whole recording, so speaker ids
+        are recording-level. Each speaker is transcribed only where they talk
+        (`stream_chunks`), with everyone else zeroed (`masked_audio`); words take
+        their stream's speaker, and a word two streams both heard at once is
+        kept once (`dedupe_words`). Crosstalk under a speaker's mask is heard.
+        With at most one speaker active nothing is masked (today's exact text);
+        if diarization fails, the unmasked transcript gets `diarization_error`.
+        """
+        try:
+            activity = NemotronDiarizer.activity(array, sample_rate=sr)
+            keep = NemotronDiarizer.top_speakers(activity, **diarization_params)
+        except Exception as e:
+            result = self._transcribe_aligned(array, sr, **kwargs)
+            result["speaker_segments"] = []
+            result["diarization_error"] = str(e)
+            return result
+
+        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:
+            result = self._transcribe_aligned(array, sr, **kwargs)
+            name = f"SPEAKER_{keep.index(talking[0]) if talking else 0}"
+            result["words"] = [{**w, "speaker": name} for w in result["words"]]
+        else:
+            chunks: list[StreamChunk] = []
+            for c in talking:
+                spans = NemotronDiarizer.sample_spans(masks[c], sr, len(array))
+                chunks.extend(
+                    {"col": c, "start": s, "audio": masked_audio(array, spans, s, e)}
+                    for s, e in stream_chunks(array, spans, sr)
                 )
-            except Exception as e:
-                result["speaker_segments"] = []
-                result["diarization_error"] = str(e)
-
+            texts = [self._transcribe_chunk(ch["audio"], sr, **kwargs)["text"] for ch in chunks]
+            result = self._align_streams(chunks, texts, sr, masks, keep)
+            if result["words"]:
+                result["words"] = NemotronDiarizer.dedupe_words(result["words"], activity, keep)
+                result["text"] = " ".join(w["word"] for w in result["words"])
+        result["speaker_segments"] = NemotronDiarizer.segments(activity, keep)
         return result
 
+    @staticmethod
+    def _align_streams(
+        chunks: list[StreamChunk],
+        texts: list[str],
+        sr: int,
+        masks: dict[int, npt.NDArray[np.bool_]],
+        keep: list[int],
+    ) -> dict[str, Any]:
+        """Align every stream chunk in one batch; `{"text", "words"}` merged by time.
+
+        If alignment fails the text is the chunk transcripts by start time and
+        `timestamp_error` is set.
+        """
+        try:
+            aligned = QwenForcedAligner.align_chunks(
+                [(ch["audio"], text) for ch, text in zip(chunks, texts, strict=True)],
+                sample_rate=sr,
+            )
+        except Exception as e:
+            order = sorted(range(len(chunks)), key=lambda k: chunks[k]["start"])
+            text = " ".join(texts[k] for k in order if texts[k])
+            return {"text": text, "words": [], "timestamp_error": str(e)}
+        words = NemotronDiarizer.stream_words(chunks, aligned, sr, masks, keep)
+        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."""
         if isinstance(inputs, dict):

diarization.py

@@ -1,6 +1,7 @@
 """Speaker diarization with NVIDIA Nemotron-3-Diarization."""
 
 import importlib.util
+import re
 from typing import TYPE_CHECKING, Any, TypedDict
 
 import numpy as np
@@ -17,6 +18,8 @@ if TYPE_CHECKING:
         Nemotron3DiarizationProcessor,
     )
 
+    from .alignment import AlignedWord
+
 
 class SpeakerSegment(TypedDict):
     """One speaker turn, in seconds."""
@@ -52,6 +55,55 @@ def _label_runs(labels: npt.NDArray[np.generic]) -> list[tuple[int, int, int]]:
     return [(labels[s].item(), int(s), int(e)) for s, e in zip(starts, ends, strict=True)]
 
 
+def masked_audio(
+    audio: npt.NDArray[np.float32], spans: list[tuple[int, int]], start: int, end: int
+) -> npt.NDArray[np.float32]:
+    """`audio[start:end]` with every sample outside `spans` set to 0.0.
+
+    `spans` are absolute `(start, end)` sample ranges, end exclusive. Only the
+    slice is copied, so a speaker stream of a long recording never holds a
+    full-length masked copy. Zeros are digital silence: the encoder sees its
+    silence floor there, and `is_silent` skips a chunk that is nothing else.
+    """
+    out = np.zeros(end - start, dtype=np.float32)
+    for s, e in spans:
+        lo, hi = max(s, start), min(e, end)
+        if lo < hi:
+            out[lo - start : hi - start] = audio[lo:hi]
+    return out
+
+
+def pack_spans(spans: list[tuple[int, int]], limit: int) -> list[tuple[int, int]]:
+    """Merge consecutive sorted `(start, end)` spans into ranges at most `limit` long.
+
+    Each range starts and ends on a span, so a speaker's chunk never begins or
+    ends in audio where they are silent. A span already over `limit` stays alone.
+    """
+    ranges: list[tuple[int, int]] = []
+    for s, e in spans:
+        if ranges and e - ranges[-1][0] <= limit:
+            ranges[-1] = (ranges[-1][0], e)
+        else:
+            ranges.append((s, e))
+    return ranges
+
+
+class StreamChunk(TypedDict):
+    """One stretch of one speaker's masked audio: kept column, absolute start, samples."""
+
+    col: int
+    start: int
+    audio: npt.NDArray[np.float32]
+
+
+_NON_WORD_RE = re.compile(r"[^\w']+")
+
+
+def _norm_word(word: str) -> str:
+    """Lowercase `word` without punctuation, for matching one word heard in two streams."""
+    return _NON_WORD_RE.sub("", word.lower())
+
+
 class NemotronDiarizer:
     """Speaker activity from NVIDIA Nemotron-3-Diarization, per 10 ms frame.
 
@@ -68,13 +120,16 @@ class NemotronDiarizer:
         >>> activity = NemotronDiarizer.activity(audio)  # (frames, 8) probabilities
         >>> keep = NemotronDiarizer.top_speakers(activity, max_speakers=2)
         >>> segments = NemotronDiarizer.segments(activity, keep)
-        >>> words = NemotronDiarizer.assign_speakers_to_words(words, activity, keep)
+
+    `ASRPipeline(..., return_speakers=True)` turns this into speaker streams:
+    `sample_spans` gives each kept speaker a mask, the ASR runs once per
+    speaker on audio with everyone else silenced, and `span_is_active` /
+    `dedupe_words` clean up the merge.
     """
 
     MODEL_ID = "nvidia/Nemotron-3-Diarization"
     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
 
@@ -130,8 +185,7 @@ class NemotronDiarizer:
 
         Every column that is ever above SEGMENT_THRESHOLD, capped by
         `num_speakers` (exact, when known) or `max_speakers` to the columns
-        with the most speech. Capping is how a known count is applied:
-        a dropped column's words go to the most active remaining speaker.
+        with the most speech. A dropped column gets no stream of its own.
         """
         active = [
             c for c in range(activity.shape[1]) if (activity[:, c] > cls.SEGMENT_THRESHOLD).any()
@@ -156,45 +210,95 @@ class NemotronDiarizer:
         return sorted(out, key=lambda seg: (seg["start"], seg["end"]))
 
     @classmethod
-    def speaker_columns(
+    def _frame_range(cls, start: float, end: float, n_frames: int) -> tuple[int, int]:
+        """Frames `[lo, hi)` under the span `start..end` s: at least one, clamped to the array."""
+        lo = min(int(start / cls.FRAME_S), n_frames - 1)
+        hi = max(lo + 1, min(int(np.ceil(end / cls.FRAME_S)), n_frames))
+        return lo, hi
+
+    @classmethod
+    def sample_spans(
+        cls, frames: npt.NDArray[np.bool_], sample_rate: int, n_samples: int
+    ) -> list[tuple[int, int]]:
+        """The active runs of a frame mask as `(start, end)` sample ranges, end exclusive.
+
+        One frame covers `FRAME_S * sample_rate` samples. A run that reaches the
+        last frame runs to `n_samples`: the model's frame count can fall a few
+        samples short of the audio, and that tail belongs to the run.
+        """
+        hop = cls.FRAME_S * sample_rate
+        spans: list[tuple[int, int]] = []
+        for on, s, e in _label_runs(frames):
+            start = min(round(s * hop), n_samples)
+            end = n_samples if e == len(frames) else min(round(e * hop), n_samples)
+            if on and start < end:
+                spans.append((start, end))
+        return spans
+
+    @classmethod
+    def span_is_active(cls, frames: npt.NDArray[np.bool_], start: float, end: float) -> bool:
+        """True if any frame of the span `start..end` s is in the mask `frames`."""
+        if len(frames) == 0:
+            return False
+        lo, hi = cls._frame_range(start, end, len(frames))
+        return bool(frames[lo:hi].any())
+
+    @classmethod
+    def stream_words(
         cls,
-        spans: list[tuple[float | None, float | None]],
-        activity: npt.NDArray[np.float32],
+        chunks: list[StreamChunk],
+        aligned: "list[list[AlignedWord]]",
+        sample_rate: int,
+        masks: dict[int, npt.NDArray[np.bool_]],
         keep: list[int],
-    ) -> list[int]:
-        """The kept activity column that said each `(start_s, end_s)` word.
+    ) -> list[dict[str, Any]]:
+        """Every stream's aligned words on the recording timeline, labelled, by start.
 
-        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.
+        A word placed wholly in its own speaker's silenced audio (`masks`) is a
+        hallucination on silence and is dropped.
         """
-        cols = np.asarray(keep)
-        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(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)
-        return out
+        names = {c: f"SPEAKER_{i}" for i, c in enumerate(keep)}
+        words: list[dict[str, Any]] = []
+        for ch, chunk_words in zip(chunks, aligned, strict=True):
+            offset = ch["start"] / sample_rate
+            for w in chunk_words:
+                start, end = w["start"] + offset, w["end"] + offset
+                if cls.span_is_active(masks[ch["col"]], start, end):
+                    words.append({**w, "start": start, "end": end, "speaker": names[ch["col"]]})
+        return sorted(words, key=lambda w: (w["start"], w["end"]))
+
+    @classmethod
+    def _span_activity(
+        cls, activity: npt.NDArray[np.float32], col: int, start: float, end: float
+    ) -> float:
+        """Mean activity of column `col` over the span `start..end` s."""
+        lo, hi = cls._frame_range(start, end, len(activity))
+        return float(activity[lo:hi, col].mean())
 
     @classmethod
-    def assign_speakers_to_words(
+    def dedupe_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
+        """Drop one copy of a word two speaker streams both heard.
+
+        `words` are sorted by start and labelled `SPEAKER_i` (position in
+        `keep`). Two words from different speakers with the same text (case
+        and punctuation ignored) and overlapping spans are one word that
+        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)}
+        dropped: set[int] = set()
+        for i, a in enumerate(words):
+            for j in range(i + 1, len(words)):
+                b = words[j]
+                if b["start"] >= a["end"]:
+                    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"]):
+                    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"])
+                dropped.add(j if score_a >= score_b else i)
+        return [w for k, w in enumerate(words) if k not in dropped]