A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Migrate whisper inputs to TensorSchema (#23505)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Sep 3, 2025 at 11:04 UTC
731a6940e39e84619bbc8db8a794563bb8cc61a5
1 file changed
+12
-4
vllm/model_executor/models/whisper.py
+12
-4
@@ -4,7 +4,7 @@
4
import math
5
from collections.abc import Iterable, Mapping, Sequence
6
from contextlib import nullcontext
7
-from typing import Literal, Optional, TypedDict, Union, cast
7
+from typing import Annotated, Literal, Optional, Union, cast
8
9
import numpy as np
10
import torch
@@ -40,6 +40,7 @@ from vllm.multimodal.processing import (BaseProcessingInfo,
40
PromptReplacement, PromptUpdate)
41
from vllm.multimodal.profiling import BaseDummyInputsBuilder
42
from vllm.transformers_utils.processor import cached_get_processor
43
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
44
45
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,
46
SupportsTranscription, SupportsV0Only)
@@ -111,9 +112,16 @@ ISO639_1_SUPPORTED_LANGS = {
112
}
113
114
114
-class WhisperAudioInputs(TypedDict):
115
- input_features: NestedTensors
116
- """Shape: `(batch_size, 128, M)`"""
115
+class WhisperAudioInputs(TensorSchema):
116
+ """
117
+ Dimensions:
118
+ - b: Batch size
119
+ - nmb: Number of mel bins
120
+ - t: Time frames (M)
121
+ """
122
+
123
+ input_features: Annotated[Optional[NestedTensors],
124
+ TensorShape("b", "nmb", "t")]
125
126
127
class WhisperPositionalEmbedding(nn.Embedding):