prosody2text-weights / whisper_decoder.py

Filter the max_length warning by message, since unsetting it does not work

5ab9498 6 days ago

16.9 kB

"""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). The weights come from the model repo, so "
                "load through it rather than pointing here by hand:\n"
                "    AutoModel.from_pretrained("
                "'Daporte/prosody2text-weights', trust_remote_code=True)"
            )

        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

        # Silence one warning, per process, by message.
        #
        # Every call here passes max_new_tokens, and Whisper's generate()
        # fills max_length in from config.max_target_positions (448) on its
        # way through, so transformers reports them both set -- once per
        # utterance, which over a few thousand buries the running WER.
        #
        # Clearing generation_config.max_length does not stop it: it is
        # already None after load (checked), because the value is
        # reconstructed inside generate() rather than read from the
        # checkpoint. Nothing on this side can unset it, and max_new_tokens
        # takes precedence regardless -- the warning says so itself. So drop
        # the record rather than pretend to fix the cause.
        import logging

        class _NoMaxLengthClash(logging.Filter):
            def filter(self, record: logging.LogRecord) -> bool:
                return "seem to have been set" not in record.getMessage()

        logging.getLogger("transformers.generation.utils").addFilter(
            _NoMaxLengthClash()
        )

        # 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

	"""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). The weights come from the model repo, so "
	"load through it rather than pointing here by hand:\n"
	" AutoModel.from_pretrained("
	"'Daporte/prosody2text-weights', trust_remote_code=True)"
	)

	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

	# Silence one warning, per process, by message.
	#
	# Every call here passes max_new_tokens, and Whisper's generate()
	# fills max_length in from config.max_target_positions (448) on its
	# way through, so transformers reports them both set -- once per
	# utterance, which over a few thousand buries the running WER.
	#
	# Clearing generation_config.max_length does not stop it: it is
	# already None after load (checked), because the value is
	# reconstructed inside generate() rather than read from the
	# checkpoint. Nothing on this side can unset it, and max_new_tokens
	# takes precedence regardless -- the warning says so itself. So drop
	# the record rather than pretend to fix the cause.
	import logging

	class _NoMaxLengthClash(logging.Filter):
	def filter(self, record: logging.LogRecord) -> bool:
	return "seem to have been set" not in record.getMessage()

	logging.getLogger("transformers.generation.utils").addFilter(
	_NoMaxLengthClash()
	)

	# 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