whisper.py

# Copyright 2025 Rebellions Inc. All rights reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import (
    SupportsMultiModal,
    SupportsTranscription,
)
from vllm.model_executor.models.whisper import (
    ISO639_1_SUPPORTED_LANGS,
    WhisperAudioInputs,
)
from vllm.utils.jsontree import json_map_leaves

from .base import ModelInputForRBLN
from .model_base import RBLNOptimumDecoderMixin, RBLNOptimumModelBase

logger = init_logger(__name__)


class RBLNOptimumWhisperForConditionalGeneration(
    RBLNOptimumModelBase,
    RBLNOptimumDecoderMixin,
    SupportsTranscription,
    SupportsMultiModal,
):
    # Whisper only supports audio-conditioned generation.
    supports_transcription_only = True
    supports_segment_timestamp = True
    supported_languages = ISO639_1_SUPPORTED_LANGS

    @classmethod
    def validate_language(cls, language: str | None) -> str | None:
        if language is None:
            # TODO language should be optional and can be guessed.
            # For now we default to en. See
            # https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/generation_whisper.py#L1520
            logger.warning(
                "Defaulting to language='en'. If you wish to transcribe "
                "audio in a different language, pass the `language` field "
                "in the TranscriptionRequest."
            )
            language = "en"
        return super().validate_language(language)

    # FIXME more method needed

    def __init__(
        self,
        vllm_config: VllmConfig,
    ) -> None:
        super().__init__(vllm_config=vllm_config)
        assert self.kv_block_adapter is not None
        self.setup_decoder_mixin(
            attn_impl=self.attn_impl,
            vocab_size=self.model_config.get_vocab_size,
            use_multiple_decoder=False,
            default_batch_size=self.scheduler_config.max_num_seqs,
            decoder_batch_sizes=[self.batch_size],
            num_blocks=self.kv_block_adapter._estimated_num_blocks(),
        )
        self.dec_max_seq_len = self.model_config.max_model_len
        self.dec_lengths = [0] * self.batch_size

    def forward(self, model_input: ModelInputForRBLN, **kwargs) -> torch.Tensor:
        input_ids = model_input.input_tokens
        block_tables = model_input.block_tables
        request_nums = input_ids.shape[0]
        is_prompt = model_input.is_prompt
        valid_block_ids = block_tables.flatten().to(torch.int32)

        if is_prompt:
            if model_input.multi_modal_kwargs:
                audio_input = self._parse_and_validate_audio_input(
                    **model_input.multi_modal_kwargs
                )
                input_features = audio_input["input_features"]
            if input_features is None:
                raise ValueError("Whisper requires `input_features` as an input.")
            _ = self.model.encoder(
                input_features=input_features,
                block_tables=block_tables.squeeze(0).to(torch.int16),
            )

        cache_position = torch.zeros(request_nums, 1, dtype=torch.int32)

        # In whisper model,
        # decoder input is always required in prefill step,
        # so is_prompt=False is set for both prefill and decode step.
        kwargs = self.preprocess_for_decoder(
            is_prompt=False,
            block_tables=block_tables,
            input_ids=input_ids,
            cache_position=cache_position,
            input_block_ids=valid_block_ids,
        )
        decoder_cache_position = kwargs.pop("cache_position")
        decoder_block_tables = kwargs.pop("block_tables")

        # Whisper model does not support bucketing.
        decoder_attention_mask = torch.zeros(
            self.batch_size, self.dec_max_seq_len, dtype=self.dtype
        )
        if is_prompt:
            decoder_input_ids = torch.full(
                (self.batch_size, 1),
                self.model.config.decoder_start_token_id,
                dtype=torch.long,
            )
            for batch_idx in valid_block_ids:
                decoder_cache_position[batch_idx] = 0
                decoder_attention_mask[batch_idx, 0] = 1
                self.dec_lengths[batch_idx] = 1
            decoder_output = self.model.decoder(
                decoder_input_ids=decoder_input_ids.contiguous(),
                decoder_attention_mask=decoder_attention_mask,
                cache_position=decoder_cache_position,
                block_tables=decoder_block_tables,
            )

        else:
            decoder_input_ids = kwargs.pop("input_ids")
            # Generate cache_position using dec_lengths
            for batch_idx in valid_block_ids:
                decoder_cache_position[batch_idx] = self.dec_lengths[batch_idx]
                decoder_attention_mask[
                    batch_idx, : decoder_cache_position[batch_idx] + 1
                ] = 1
                self.dec_lengths[batch_idx] += 1

            decoder_output = self.model.decoder(
                decoder_input_ids=decoder_input_ids.contiguous(),
                decoder_attention_mask=decoder_attention_mask,
                cache_position=decoder_cache_position,
                block_tables=decoder_block_tables,
            )

        lm_logits = decoder_output.logits
        lm_logits = lm_logits[valid_block_ids]
        return lm_logits

    def _parse_and_validate_audio_input(self, **kwargs: object) -> WhisperAudioInputs:
        input_features = kwargs.pop("input_features", None)

        if input_features is not None:
            input_features = json_map_leaves(lambda x: x.to(self.dtype), input_features)
        return WhisperAudioInputs(input_features=input_features)