Upload whisper_decoder.py with huggingface_hub

whisper_decoder.py

@@ -0,0 +1,366 @@
+"""Run the fine-tuned Whisper over a predicted mel.
+
+Uses transformers' own `generate()` rather than a hand-rolled decoder. An
+earlier ONNX export with custom beam search produced single punctuation marks
+and empty strings where `generate()` produced fluent sentences from the
+identical mel -- the model was being silently crippled by the decoding rather
+than by the input.
+
+The mel contract is the subtle part. The checkpoint keeps Whisper's standard
+**80-channel** input even though it was fine-tuned on 12-bin prosody-predicted
+mels: `num_mel_bins` in the training config truncates incoming mels, it never
+resizes conv1. The training collator reconciles the two by zero-padding, and
+this module reproduces that exactly:
+
+    conformer -> [T, 12] -> transpose -> [12, T] -> zero-pad -> [80, 3000]
+
+Bins 12-79 are zeros. Getting this wrong fails at the first conv layer.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+# Whisper's encoder window: 3000 mel frames at 100 Hz = 30 s.
+ENCODER_FRAMES = 3000
+ENCODER_MEL_BINS = 80
+
+# The decoder prefix the checkpoint was fine-tuned behind.
+#
+# Its own forced_decoder_ids is [[1, None], [2, <|transcribe|>]] -- position 1,
+# the language slot, is left unset. Decoding without pinning it does not make
+# the model choose a language: it skips the slot and starts from
+# <|startoftranscript|> alone, so the prefix is two tokens short of what
+# training used.
+#
+# The published inference runs set language: en, and their recorded
+# snapshot_meta.json confirms the full prefix
+# [50258, 50259, 50359, 50363] = <|sot|> <|en|> <|transcribe|> <|notimestamps|>.
+# Matching it is the point: a different prefix is a different model condition,
+# so any comparison against those runs is confounded without it.
+DEFAULT_LANGUAGE = "en"
+
+# Generation limits the published inference runs used.
+DEFAULT_MAX_NEW_TOKENS = 225
+DEFAULT_NO_REPEAT_NGRAM = 3
+
+# 10, matching the published inference config for this checkpoint.
+#
+# Wider is cheap and was tried: 385 ms at 10 beams against 634 ms at 40 on a
+# healthy phrase, four times the hypotheses for a quarter more latency. It
+# made the transcriptions worse -- a wider beam searches harder for the
+# highest-likelihood sequence, and on a signal as weak as a drawn contour
+# that is the YouTube-caption boilerplate the checkpoint overlearned. More
+# beams find the attractor more reliably rather than finding better words.
+DEFAULT_BEAMS = 10
+
+# Phrases Whisper falls back on when the input carries no usable speech.
+#
+# Trained on YouTube captions, it learned what tends to accompany
+# unintelligible audio: "Thank you." for silence, and shouted boilerplate
+# from auto-generated or burned-in captions. Feeding it pure noise returns
+# "I'M NOT GOING TO SAY ANYTHING" in three of the top four beams; feeding it
+# a real drawing returns it in none. Naming it in the UI keeps it from
+# reading as a transcription of the drawing.
+HALLUCINATION_MARKERS = (
+    "i'm not going to say anything",
+    "thank you for watching",
+    "subscribe",
+    "thanks for watching",
+)
+
+
+def looks_like_fallback(text: str) -> bool:
+    """Whether a hypothesis is the model's no-speech boilerplate."""
+    lowered = text.strip().lower()
+    if not lowered:
+        return False
+    if len(lowered) > 3 and text.strip().isupper():
+        return True
+    return any(marker in lowered for marker in HALLUCINATION_MARKERS)
+
+
+# Measured across the templates: real speech runs about 4.3 words/sec, p90
+# 5.6. The drawing is the only evidence about how long the utterance is, so
+# the token budget is derived from its duration rather than left at a flat
+# 48 -- otherwise a half-second scribble can come back as a paragraph, which
+# is the language model talking rather than a transcription.
+WORDS_PER_SECOND = 5.6  # p90, so an unusually fast phrase is not truncated
+TOKENS_PER_WORD = 2.5  # Whisper splits a word into roughly this many tokens
+MIN_TOKENS = 4
+
+
+def pad_mel_for_whisper(
+    mel: np.ndarray,
+    mel_bins: int = ENCODER_MEL_BINS,
+    frames: int = ENCODER_FRAMES,
+) -> np.ndarray:
+    """[T, bins] from the conformer -> [bins_padded, frames] for Whisper.
+
+    Transposes to Whisper's [bins, time] layout, then zero-pads both axes.
+    This mirrors the training collator's `pad_mel_to_expected` branch, so the
+    model sees the same shape it was fine-tuned on.
+    """
+    mel = np.asarray(mel, dtype=np.float32)
+    if mel.ndim != 2:
+        raise ValueError(f"expected a 2-D mel, got {mel.shape}")
+
+    # The conformer emits [T, bins]; Whisper wants [bins, T].
+    mel = mel.T
+
+    out = np.zeros((mel_bins, frames), dtype=np.float32)
+    bins = min(mel.shape[0], mel_bins)
+    time = min(mel.shape[1], frames)
+    out[:bins, :time] = mel[:bins, :time]
+    return out
+
+
+class Whisper:
+    """Transcribes a predicted mel with the fine-tuned checkpoint."""
+
+    def __init__(
+        self,
+        model_dir: Path | str,
+        device: str | None = None,
+        num_beams: int = DEFAULT_BEAMS,
+        half: bool = False,
+        tokenizer_dir: Path | str | None = None,
+    ) -> None:
+        import torch
+        from transformers import (WhisperForConditionalGeneration,
+                                  WhisperTokenizerFast)
+
+        model_dir = Path(model_dir)
+        if not (model_dir / "config.json").exists():
+            raise FileNotFoundError(
+                f"{model_dir} is not a Whisper checkpoint directory "
+                "(no config.json). Fetch the weights with "
+                "python scripts/fetch_models.py"
+            )
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.num_beams = num_beams
+
+        # fp16 is 1.8x faster (493 ms against 793 ms per phrase at beam 10)
+        # and halves the VRAM, and it transcribed all twelve templates
+        # identically -- the acoustic evidence in a drawn contour is far
+        # coarser than the gap between the precisions. fp32 is the default
+        # because that is what the published runs used.
+        #
+        # CPU stays fp32 regardless: torch has no fast half kernels there, so
+        # it would be slower, not faster.
+        self.dtype = (
+            torch.float16 if half and self.device.type == "cuda" else torch.float32
+        )
+        self.model = WhisperForConditionalGeneration.from_pretrained(
+            model_dir, dtype=self.dtype
+        )
+        self.model.eval().to(self.device)
+        self.mel_bins: int = self.model.config.num_mel_bins
+
+        # The checkpoint carries no tokenizer of its own; it is the stock
+        # whisper-medium one. Loading it from a local directory keeps first
+        # run offline -- otherwise this is a HuggingFace Hub fetch, which on
+        # a fresh Space or an air-gapped machine is a startup failure rather
+        # than a slow import.
+        source = str(tokenizer_dir) if tokenizer_dir else "openai/whisper-medium"
+        self.tokenizer = WhisperTokenizerFast.from_pretrained(source)
+
+    @property
+    def provider(self) -> str:
+        return str(self.device)
+
+    def _digit_token_ids(self) -> list[int]:
+        """Whisper tokens containing a digit.
+
+        Banned under `constrain`: a prosody contour carries no evidence for
+        "15" over "fifteen", and numerals are a large, acoustically
+        meaningless slice of the vocabulary. Spelled-out numbers remain
+        available.
+        """
+        if getattr(self, "_digits", None) is None:
+            # Decode each id rather than inspecting the raw token: Whisper's
+            # timestamp markers are spelled "<|24.08|>" and would otherwise
+            # match on the digits in their name. They decode to nothing and
+            # are suppressed already.
+            self._digits = [
+                index
+                for index in range(len(self.tokenizer))
+                if any(c.isdigit() for c in self.tokenizer.decode([index]))
+            ]
+        return self._digits
+
+    @staticmethod
+    def tokens_for(mel: np.ndarray) -> int:
+        """A token budget proportional to the drawing's length.
+
+        The mel runs at 100 Hz, so its frame count is the duration. Letting
+        generation run to a flat limit regardless lets a short scribble come
+        back as a long sentence, which is the decoder's prior rather than the
+        drawing's content.
+        """
+        seconds = len(mel) / 100.0
+        words = seconds * WORDS_PER_SECOND
+        return max(MIN_TOKENS, int(round(words * TOKENS_PER_WORD)))
+
+    def token_budget(self, requested: int, language: str | None) -> int:
+        """`requested`, clipped to what the decoder can actually emit.
+
+        Whisper's decoder has max_target_positions=448 and counts the prefix
+        against it, so a budget derived from the input length has to leave
+        room for the prefix or generate() raises -- which it did on a 4.8 s
+        clip, where the duration-derived budget came to 478 and the
+        four-token prefix pushed the total to 482.
+        """
+        limit = int(getattr(self.model.config, "max_target_positions", 448))
+        prefix = len(self.decoder_prefix(language))
+        # One spare for EOS, which also occupies a position.
+        return max(1, min(int(requested), limit - prefix - 1))
+
+    def decoder_prefix(self, language: str | None = DEFAULT_LANGUAGE) -> list[int]:
+        """The tokens the decoder is primed with before any content.
+
+        Built from the tokenizer with the language and task set, which is how
+        the published runs got theirs: their snapshot_meta.json records
+        [50258, 50259, 50359, 50363], and this reproduces exactly that.
+        """
+        if language:
+            # prefix_tokens is a property derived from these two, so set them
+            # on the tokenizer rather than hand-assembling the sequence.
+            self.tokenizer.set_prefix_tokens(language=language, task="transcribe")
+            return [int(t) for t in self.tokenizer.prefix_tokens]
+
+        prefix = [self.model.config.decoder_start_token_id]
+        forced = getattr(self.model.generation_config, "forced_decoder_ids", None)
+        for _position, token in forced or []:
+            if token is not None:
+                prefix.append(int(token))
+        return prefix
+
+    def hypotheses(
+        self,
+        mel: np.ndarray,
+        max_tokens: int | None = None,
+        num_beams: int | None = None,
+        top_k: int | None = None,
+        language: str | None = DEFAULT_LANGUAGE,
+        no_repeat_ngram_size: int | None = None,
+        constrain: bool = True,
+    ) -> list[tuple[str, float]]:
+        """The beam's ranked candidates as (text, score), best first.
+
+        `language` pins the decoder prefix; see DEFAULT_LANGUAGE. Pass None
+        to leave the checkpoint's own forced_decoder_ids alone, which gives a
+        two-token-shorter prefix than training used.
+
+        On input this uncertain the top hypothesis is often a narrow winner,
+        so the runners-up say more about what the model heard than the single
+        best string does.
+
+        Scores are HuggingFace's length-normalised sequence log-probabilities
+        where available (higher is better), and NaN where the capture below
+        cannot recover them -- the ordering is always meaningful even when the
+        numbers are not.
+        """
+        import torch
+        from transformers import GenerationMixin
+
+        beams = num_beams or self.num_beams
+        limits = {}
+        if language:
+            limits["language"] = language
+
+        # The published runs set no_repeat_ngram_size: 3 and nothing else.
+        ngram = (DEFAULT_NO_REPEAT_NGRAM if no_repeat_ngram_size is None
+                 else no_repeat_ngram_size)
+        if ngram:
+            limits["no_repeat_ngram_size"] = ngram
+
+        # `constrain` is what the published dataset runs used: ban numerals,
+        # and bound the length by the utterance rather than by a flat limit.
+        # Letting generation run to 225 regardless lets a half-second clip
+        # come back as a paragraph, which is the decoder's prior talking
+        # rather than the input.
+        if constrain:
+            limits["bad_words_ids"] = [[i] for i in self._digit_token_ids()]
+            max_tokens = max_tokens or self.tokens_for(mel)
+        else:
+            max_tokens = max_tokens or DEFAULT_MAX_NEW_TOKENS
+        max_tokens = self.token_budget(max_tokens, language)
+        features = pad_mel_for_whisper(mel, self.mel_bins, ENCODER_FRAMES)
+        batch = torch.from_numpy(features).unsqueeze(0).to(self.device, self.dtype)
+
+        if beams <= 1:
+            with torch.no_grad():
+                ids = self.model.generate(
+                    batch, num_beams=1, max_new_tokens=max_tokens, **limits
+                )
+            return [(self.tokenizer.batch_decode(
+                ids, skip_special_tokens=True)[0].strip(), float("nan"))]
+
+        # Whisper's generate() wrapper keeps only the top beam and does not
+        # honour num_return_sequences; asking it for all of them inflates the
+        # decoder cache until a 24 GB GPU runs out. The approach here is the
+        # one the reference inference script uses: patch
+        # GenerationMixin.generate, set num_return_sequences on the inner
+        # call, capture every beam, and hand Whisper back the single sequence
+        # it expects.
+        captured: list = []
+        original = GenerationMixin.generate
+
+        def capturing(model_self, input_features, **kwargs):
+            config = kwargs.get("generation_config")
+            width = config.num_beams if config else 1
+            if width <= 1:
+                return original(model_self, input_features, **kwargs)
+
+            config.num_return_sequences = width
+            # sequences_scores only comes back when the inner call is asked
+            # for a dict; Whisper's wrapper does not always pass this down.
+            kwargs["return_dict_in_generate"] = True
+            # ...and, since transformers 5, only when scores are explicitly
+            # requested. Without this the attribute is present but None, so
+            # every hypothesis came back scored NaN -- the ranking survived,
+            # being the beam's own order, but the numbers beside it were
+            # gone and nothing said so.
+            kwargs["output_scores"] = True
+            out = original(model_self, input_features, **kwargs)
+            config.num_return_sequences = 1
+
+            seqs = out.sequences if hasattr(out, "sequences") else out
+            captured.append((seqs, getattr(out, "sequences_scores", None)))
+
+            top1 = seqs[::width]
+            if hasattr(out, "sequences"):
+                out.sequences = top1
+            else:
+                out = top1
+            return out
+
+        GenerationMixin.generate = capturing
+        try:
+            with torch.no_grad():
+                self.model.generate(
+                    batch,
+                    num_beams=beams,
+                    max_new_tokens=max_tokens,
+                    return_dict_in_generate=True,
+                    **limits,
+                )
+        finally:
+            GenerationMixin.generate = original
+
+        if not captured:
+            return []
+        seqs, scores = captured[0]
+        texts = self.tokenizer.batch_decode(seqs, skip_special_tokens=True)
+        values = (
+            scores.tolist() if scores is not None
+            else [float("nan")] * len(texts)
+        )
+        ranked = [(t.strip(), s) for t, s in zip(texts, values)]
+        return ranked[:top_k] if top_k else ranked