Whisper model support & /v1/audio/transcriptions endpoint & benchmark (#16983)
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com> Co-authored-by: MahmoudAshraf97 <hassouna97.ma@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Xinyuan Tong
@@ -545,6 +545,17 @@ class ModelConfig:
|
||||
if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures:
|
||||
loop_num = getattr(self.hf_text_config, "loop_num", 1)
|
||||
self.num_attention_layers = int(self.num_hidden_layers * int(loop_num))
|
||||
if "WhisperForConditionalGeneration" in self.hf_config.architectures:
|
||||
# Whisper has unique layer ID scheme:
|
||||
# - Encoder self-attention: 0 to encoder_layers-1 (no KV cache)
|
||||
# - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache)
|
||||
# - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1
|
||||
# Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist
|
||||
encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0)
|
||||
decoder_layers = getattr(
|
||||
self.hf_text_config, "decoder_layers", self.num_hidden_layers
|
||||
)
|
||||
self.num_attention_layers = encoder_layers + 2 * decoder_layers
|
||||
self.num_nextn_predict_layers = getattr(
|
||||
self.hf_text_config, "num_nextn_predict_layers", None
|
||||
)
|
||||
@@ -1247,6 +1258,7 @@ multimodal_model_archs = [
|
||||
"InternS1ForConditionalGeneration",
|
||||
"InternS1ProForConditionalGeneration",
|
||||
"Phi4MMForCausalLM",
|
||||
"WhisperForConditionalGeneration",
|
||||
"Step3VLForConditionalGeneration",
|
||||
"POINTSV15ChatModel",
|
||||
"DotsVLMForCausalLM",
|
||||
@@ -1285,11 +1297,17 @@ def is_image_gen_model(model_architectures: List[str]):
|
||||
|
||||
|
||||
def is_audio_model(model_architectures: List[str]):
|
||||
return False
|
||||
models = [
|
||||
"WhisperForConditionalGeneration",
|
||||
]
|
||||
return any(model in model_architectures for model in models)
|
||||
|
||||
|
||||
def is_encoder_decoder_model(model_architectures: List[str]):
|
||||
return "MllamaForConditionalGeneration" in model_architectures
|
||||
models = [
|
||||
"WhisperForConditionalGeneration",
|
||||
]
|
||||
return any(model in model_architectures for model in models)
|
||||
|
||||
|
||||
def is_local_attention_model(model_architectures: List[str]):
|
||||
|
||||
@@ -46,7 +46,7 @@ import orjson
|
||||
import requests
|
||||
import uvicorn
|
||||
import uvloop
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from fastapi import Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
||||
@@ -93,6 +93,9 @@ from sglang.srt.entrypoints.openai.serving_tokenize import (
|
||||
OpenAIServingDetokenize,
|
||||
OpenAIServingTokenize,
|
||||
)
|
||||
from sglang.srt.entrypoints.openai.serving_transcription import (
|
||||
OpenAIServingTranscription,
|
||||
)
|
||||
from sglang.srt.entrypoints.warmup import execute_warmups
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
@@ -298,6 +301,9 @@ async def lifespan(fast_api_app: FastAPI):
|
||||
fast_api_app.state.openai_serving_detokenize = OpenAIServingDetokenize(
|
||||
_global_state.tokenizer_manager
|
||||
)
|
||||
fast_api_app.state.openai_serving_transcription = OpenAIServingTranscription(
|
||||
_global_state.tokenizer_manager
|
||||
)
|
||||
|
||||
# Initialize Ollama-compatible serving handler
|
||||
fast_api_app.state.ollama_serving = OllamaServing(_global_state.tokenizer_manager)
|
||||
@@ -1418,6 +1424,38 @@ async def openai_v1_detokenize(request: DetokenizeRequest, raw_request: Request)
|
||||
)
|
||||
|
||||
|
||||
@app.post("/v1/audio/transcriptions")
|
||||
async def openai_v1_audio_transcriptions(
|
||||
raw_request: Request,
|
||||
file: UploadFile = File(...),
|
||||
model: str = Form(default="default"),
|
||||
language: Optional[str] = Form(default=None),
|
||||
response_format: str = Form(default="json"),
|
||||
temperature: float = Form(default=0.0),
|
||||
stream: bool = Form(default=False),
|
||||
):
|
||||
"""OpenAI-compatible audio transcription endpoint."""
|
||||
if response_format not in ["json", "text"]:
|
||||
return ORJSONResponse(
|
||||
content={"error": {"message": "Only 'json' and 'text' formats supported"}},
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
audio_data = await file.read()
|
||||
|
||||
return (
|
||||
await raw_request.app.state.openai_serving_transcription.create_transcription(
|
||||
audio_data=audio_data,
|
||||
model=model,
|
||||
language=language,
|
||||
response_format=response_format,
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
raw_request=raw_request,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@app.get("/v1/models", response_class=ORJSONResponse)
|
||||
async def available_models():
|
||||
"""Show available models. OpenAI-compatible endpoint."""
|
||||
|
||||
@@ -1389,3 +1389,51 @@ class ResponseReasoningTextContent(BaseModel):
|
||||
ResponseInputOutputItem: TypeAlias = Union[
|
||||
ResponseInputItemParam, "ResponseReasoningItem", ResponseFunctionToolCall
|
||||
]
|
||||
|
||||
|
||||
# ================== Transcription API Protocol Definitions ==================
|
||||
|
||||
|
||||
class TranscriptionRequest(BaseModel):
|
||||
"""Request model for audio transcription (OpenAI-compatible)."""
|
||||
|
||||
model: str = DEFAULT_MODEL_NAME
|
||||
language: Optional[str] = None
|
||||
response_format: str = "json"
|
||||
temperature: float = 0.0
|
||||
stream: bool = False
|
||||
# Internal fields (not from API)
|
||||
audio_data: Optional[bytes] = None
|
||||
audio_duration_s: float = 0.0
|
||||
|
||||
|
||||
class TranscriptionUsage(BaseModel):
|
||||
"""Usage info for transcription response (duration-based)."""
|
||||
|
||||
type: Literal["duration"] = "duration"
|
||||
seconds: int # Audio duration in seconds (rounded up)
|
||||
|
||||
|
||||
class TranscriptionResponse(BaseModel):
|
||||
"""Non-streaming transcription response (OpenAI-compatible)."""
|
||||
|
||||
text: str
|
||||
usage: Optional[TranscriptionUsage] = None
|
||||
|
||||
|
||||
class TranscriptionStreamChoice(BaseModel):
|
||||
"""Delta content for streaming transcription."""
|
||||
|
||||
delta: DeltaMessage
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
class TranscriptionStreamResponse(BaseModel):
|
||||
"""Streaming transcription chunk (OpenAI-compatible)."""
|
||||
|
||||
id: str = Field(default_factory=lambda: f"trsc-{uuid.uuid4().hex}")
|
||||
object: Literal["transcription.chunk"] = "transcription.chunk"
|
||||
created: int = Field(default_factory=lambda: int(time.time()))
|
||||
model: str
|
||||
choices: List[TranscriptionStreamChoice]
|
||||
usage: Optional[UsageInfo] = None
|
||||
|
||||
220
python/sglang/srt/entrypoints/openai/serving_transcription.py
Normal file
220
python/sglang/srt/entrypoints/openai/serving_transcription.py
Normal file
@@ -0,0 +1,220 @@
|
||||
# Copyright 2025 SGLang Team
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""
|
||||
OpenAI-compatible transcription endpoint handler for Whisper models.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, AsyncGenerator, Optional, Union
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import (
|
||||
DeltaMessage,
|
||||
ErrorResponse,
|
||||
TranscriptionRequest,
|
||||
TranscriptionResponse,
|
||||
TranscriptionStreamChoice,
|
||||
TranscriptionStreamResponse,
|
||||
TranscriptionUsage,
|
||||
)
|
||||
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpenAIServingTranscription(OpenAIServingBase):
|
||||
"""Handler for /v1/audio/transcriptions requests"""
|
||||
|
||||
def __init__(self, tokenizer_manager: TokenizerManager):
|
||||
super().__init__(tokenizer_manager)
|
||||
|
||||
def _request_id_prefix(self) -> str:
|
||||
return "trsc-"
|
||||
|
||||
def _validate_request(self, request: TranscriptionRequest) -> Optional[str]:
|
||||
"""Validate transcription request."""
|
||||
# Validation is done in the route handler for form data
|
||||
return None
|
||||
|
||||
def _convert_to_internal_request(
|
||||
self,
|
||||
request: TranscriptionRequest,
|
||||
raw_request: Request = None,
|
||||
) -> tuple[GenerateReqInput, TranscriptionRequest]:
|
||||
"""Convert transcription request to internal format."""
|
||||
# Build sampling params - include language for WhisperProcessor
|
||||
sampling_params = {
|
||||
"temperature": request.temperature,
|
||||
"max_new_tokens": 448, # Whisper default max tokens
|
||||
"language": request.language, # Pass to WhisperProcessor for language-specific decoding
|
||||
}
|
||||
|
||||
# For Whisper, we pass audio_data and let the processor handle it
|
||||
adapted_request = GenerateReqInput(
|
||||
text="", # Empty text - Whisper processor will set proper decoder tokens
|
||||
audio_data=request.audio_data,
|
||||
sampling_params=sampling_params,
|
||||
stream=request.stream,
|
||||
modalities=["audio"],
|
||||
routing_key=self.extract_routing_key(raw_request),
|
||||
)
|
||||
|
||||
return adapted_request, request
|
||||
|
||||
def _get_audio_duration(self, audio_data: bytes) -> float:
|
||||
"""Calculate audio duration in seconds."""
|
||||
try:
|
||||
import soundfile as sf
|
||||
|
||||
audio_array, sr = sf.read(io.BytesIO(audio_data))
|
||||
duration = len(audio_array) / sr
|
||||
return duration
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not calculate audio duration: {e}")
|
||||
return 0.0
|
||||
|
||||
async def create_transcription(
|
||||
self,
|
||||
audio_data: bytes,
|
||||
model: str,
|
||||
language: Optional[str],
|
||||
response_format: str,
|
||||
temperature: float,
|
||||
stream: bool,
|
||||
raw_request: Request,
|
||||
) -> Union[TranscriptionResponse, StreamingResponse, Response, ORJSONResponse]:
|
||||
"""Main entry point for transcription requests."""
|
||||
# Calculate audio duration for usage reporting
|
||||
audio_duration_s = self._get_audio_duration(audio_data)
|
||||
|
||||
# Build request
|
||||
request = TranscriptionRequest(
|
||||
audio_data=audio_data,
|
||||
model=model,
|
||||
language=language,
|
||||
response_format=response_format,
|
||||
temperature=temperature,
|
||||
stream=stream,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
|
||||
# Use the base class handle_request pattern
|
||||
return await self.handle_request(request, raw_request)
|
||||
|
||||
async def _handle_non_streaming_request(
|
||||
self,
|
||||
adapted_request: GenerateReqInput,
|
||||
request: TranscriptionRequest,
|
||||
raw_request: Request,
|
||||
) -> Union[TranscriptionResponse, ErrorResponse, ORJSONResponse, Response]:
|
||||
"""Handle non-streaming transcription request."""
|
||||
try:
|
||||
ret = await self.tokenizer_manager.generate_request(
|
||||
adapted_request, raw_request
|
||||
).__anext__()
|
||||
except ValueError as e:
|
||||
return self.create_error_response(str(e))
|
||||
|
||||
text = ret.get("text", "")
|
||||
|
||||
# Build response based on format
|
||||
if request.response_format == "text":
|
||||
return Response(content=text, media_type="text/plain")
|
||||
|
||||
# JSON format
|
||||
usage = TranscriptionUsage(seconds=int(math.ceil(request.audio_duration_s)))
|
||||
|
||||
return TranscriptionResponse(text=text, usage=usage)
|
||||
|
||||
async def _handle_streaming_request(
|
||||
self,
|
||||
adapted_request: GenerateReqInput,
|
||||
request: TranscriptionRequest,
|
||||
raw_request: Request,
|
||||
) -> StreamingResponse:
|
||||
"""Handle streaming transcription request."""
|
||||
return StreamingResponse(
|
||||
self._generate_transcription_stream(adapted_request, request, raw_request),
|
||||
media_type="text/event-stream",
|
||||
background=self.tokenizer_manager.create_abort_task(adapted_request),
|
||||
)
|
||||
|
||||
async def _generate_transcription_stream(
|
||||
self,
|
||||
adapted_request: GenerateReqInput,
|
||||
request: TranscriptionRequest,
|
||||
raw_request: Request,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Generate streaming transcription response."""
|
||||
created_time = int(time.time())
|
||||
request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
|
||||
model = request.model
|
||||
stream_buffer = ""
|
||||
|
||||
try:
|
||||
async for content in self.tokenizer_manager.generate_request(
|
||||
adapted_request, raw_request
|
||||
):
|
||||
finish_reason = content["meta_info"]["finish_reason"]
|
||||
finish_reason_type = finish_reason["type"] if finish_reason else None
|
||||
|
||||
# Calculate delta (new text since last chunk)
|
||||
current_text = content.get("text", "")
|
||||
delta = current_text[len(stream_buffer) :]
|
||||
stream_buffer = current_text
|
||||
|
||||
# Send content delta if there's new text
|
||||
if delta:
|
||||
choice_data = TranscriptionStreamChoice(
|
||||
delta=DeltaMessage(content=delta),
|
||||
finish_reason=None,
|
||||
)
|
||||
chunk = TranscriptionStreamResponse(
|
||||
id=request_id,
|
||||
created=created_time,
|
||||
model=model,
|
||||
choices=[choice_data],
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Send finish reason when done
|
||||
if finish_reason_type:
|
||||
choice_data = TranscriptionStreamChoice(
|
||||
delta=DeltaMessage(),
|
||||
finish_reason=finish_reason_type,
|
||||
)
|
||||
chunk = TranscriptionStreamResponse(
|
||||
id=request_id,
|
||||
created=created_time,
|
||||
model=model,
|
||||
choices=[choice_data],
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
except ValueError as e:
|
||||
error = self.create_streaming_error_response(str(e))
|
||||
yield f"data: {error}\n\n"
|
||||
|
||||
yield "data: [DONE]\n\n"
|
||||
@@ -694,9 +694,15 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
||||
"the engine with skip_tokenizer_init=False."
|
||||
)
|
||||
|
||||
input_ids, token_type_ids = await self._tokenize_texts(
|
||||
input_text, is_cross_encoder_request
|
||||
)
|
||||
# For audio-only requests (e.g., Whisper), text may be empty.
|
||||
# The multimodal processor will provide input_ids later.
|
||||
if not input_text and self.mm_processor and obj.contains_mm_input():
|
||||
# Use empty placeholder - multimodal processor will override
|
||||
input_ids = []
|
||||
else:
|
||||
input_ids, token_type_ids = await self._tokenize_texts(
|
||||
input_text, is_cross_encoder_request
|
||||
)
|
||||
|
||||
if self.mm_processor and obj.contains_mm_input():
|
||||
if obj.image_data is not None and not isinstance(obj.image_data, list):
|
||||
|
||||
543
python/sglang/srt/models/whisper.py
Normal file
543
python/sglang/srt/models/whisper.py
Normal file
@@ -0,0 +1,543 @@
|
||||
from typing import Any, Iterable, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from transformers import WhisperConfig
|
||||
|
||||
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
||||
from sglang.srt.layers.activation import get_act_fn
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
|
||||
from sglang.srt.layers.quantization import QuantizationConfig
|
||||
from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
|
||||
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
from sglang.srt.managers.schedule_batch import MultimodalInputs
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
|
||||
class WhisperAttention(torch.nn.Module):
|
||||
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
num_heads: int,
|
||||
bias: bool = True,
|
||||
layer_id: Optional[int] = None,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
is_cross_attention: bool = False,
|
||||
is_encoder=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.total_num_heads = num_heads
|
||||
head_dim = embed_dim // num_heads
|
||||
self.is_cross_attention = is_cross_attention
|
||||
self.is_encoder = is_encoder
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
assert (
|
||||
num_heads % tp_size == 0
|
||||
), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})"
|
||||
self.num_heads = num_heads // tp_size
|
||||
|
||||
if (head_dim * num_heads) != embed_dim:
|
||||
raise ValueError(
|
||||
f"embed_dim must be divisible by num_heads (got `embed_dim`: {embed_dim}"
|
||||
f" and `num_heads`: {num_heads})."
|
||||
)
|
||||
self.scaling = head_dim**-0.5
|
||||
self.head_dim = head_dim
|
||||
self.kv_size = self.num_heads * head_dim
|
||||
|
||||
if is_cross_attention:
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
embed_dim, embed_dim, quant_config=quant_config
|
||||
)
|
||||
self.kv_proj = QKVParallelLinear(
|
||||
hidden_size=embed_dim,
|
||||
head_size=head_dim,
|
||||
total_num_heads=0,
|
||||
total_num_kv_heads=num_heads,
|
||||
bias=bias,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
else:
|
||||
self.qkv_proj = QKVParallelLinear(
|
||||
embed_dim, head_dim, num_heads, quant_config=quant_config
|
||||
)
|
||||
self.out_proj = RowParallelLinear(
|
||||
embed_dim, embed_dim, bias=bias, quant_config=quant_config
|
||||
)
|
||||
self.attn = RadixAttention(
|
||||
self.num_heads,
|
||||
head_dim,
|
||||
scaling=1.0,
|
||||
num_kv_heads=self.num_heads,
|
||||
layer_id=layer_id,
|
||||
quant_config=quant_config,
|
||||
is_cross_attention=is_cross_attention,
|
||||
attn_type=(
|
||||
AttentionType.ENCODER_ONLY if is_encoder else AttentionType.DECODER
|
||||
),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
cross_hidden_states: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||
"""Input shape: Batch x Time x Channel"""
|
||||
|
||||
if self.is_cross_attention:
|
||||
q, _ = self.q_proj(hidden_states)
|
||||
if cross_hidden_states is not None:
|
||||
kv, _ = self.kv_proj(cross_hidden_states)
|
||||
k, v = kv.split([self.kv_size, self.kv_size], dim=-1)
|
||||
else:
|
||||
k = torch.zeros_like(q)
|
||||
v = torch.zeros_like(q)
|
||||
|
||||
q = q * self.scaling
|
||||
num_heads = self.attn.tp_q_head_num
|
||||
head_dim = self.attn.head_dim
|
||||
|
||||
q = q.view(-1, num_heads, head_dim)
|
||||
k = k.view(-1, num_heads, head_dim)
|
||||
v = v.view(-1, num_heads, head_dim)
|
||||
|
||||
q_len = q.shape[0]
|
||||
kv_len = k.shape[0]
|
||||
|
||||
q = q.transpose(0, 1)
|
||||
k = k.transpose(0, 1)
|
||||
v = v.transpose(0, 1)
|
||||
|
||||
attn_weights = torch.bmm(q, k.transpose(1, 2))
|
||||
|
||||
# Apply block-diagonal mask for batched cross-attention
|
||||
batch_size = forward_batch.batch_size if forward_batch else 1
|
||||
if batch_size > 1 and kv_len > 0:
|
||||
encoder_len_per_request = kv_len // batch_size
|
||||
if encoder_len_per_request * batch_size == kv_len:
|
||||
is_decode = forward_batch.forward_mode.is_decode()
|
||||
if is_decode:
|
||||
mask = torch.zeros(
|
||||
(q_len, kv_len), device=q.device, dtype=torch.bool
|
||||
)
|
||||
for i in range(batch_size):
|
||||
enc_start = i * encoder_len_per_request
|
||||
enc_end = (i + 1) * encoder_len_per_request
|
||||
mask[i, enc_start:enc_end] = True
|
||||
attn_weights = attn_weights.masked_fill(
|
||||
~mask.unsqueeze(0), float("-inf")
|
||||
)
|
||||
else:
|
||||
seq_lens = forward_batch.seq_lens
|
||||
if seq_lens is not None and len(seq_lens) == batch_size:
|
||||
seq_lens_list = seq_lens.tolist()
|
||||
mask = torch.zeros(
|
||||
(q_len, kv_len), device=q.device, dtype=torch.bool
|
||||
)
|
||||
q_start = 0
|
||||
for i, dec_len in enumerate(seq_lens_list):
|
||||
enc_start = i * encoder_len_per_request
|
||||
enc_end = (i + 1) * encoder_len_per_request
|
||||
q_end = q_start + dec_len
|
||||
mask[q_start:q_end, enc_start:enc_end] = True
|
||||
q_start = q_end
|
||||
attn_weights = attn_weights.masked_fill(
|
||||
~mask.unsqueeze(0), float("-inf")
|
||||
)
|
||||
|
||||
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
|
||||
attn_output = torch.bmm(attn_weights, v)
|
||||
attn_output = attn_output.transpose(0, 1)
|
||||
attn_output = attn_output.reshape(q_len, num_heads * head_dim)
|
||||
else:
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
q, k, v = qkv.chunk(chunks=3, dim=-1)
|
||||
q = q * self.scaling
|
||||
|
||||
if self.is_encoder:
|
||||
num_heads = self.attn.tp_q_head_num
|
||||
head_dim = self.attn.head_dim
|
||||
batch_size, seq_len, _ = hidden_states.shape
|
||||
|
||||
q = q.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||
k = k.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||
v = v.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||
|
||||
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
||||
q, k, v, scale=1.0
|
||||
)
|
||||
attn_output = attn_output.permute(0, 2, 1, 3).reshape(
|
||||
batch_size, seq_len, num_heads * head_dim
|
||||
)
|
||||
else:
|
||||
attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=True)
|
||||
|
||||
attn_output, _ = self.out_proj(attn_output)
|
||||
|
||||
return attn_output
|
||||
|
||||
|
||||
class WhisperEncoderLayer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: WhisperConfig,
|
||||
layer_id: Optional[int] = None,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.embed_dim = config.d_model
|
||||
|
||||
self.self_attn = WhisperAttention(
|
||||
embed_dim=self.embed_dim,
|
||||
num_heads=config.encoder_attention_heads,
|
||||
layer_id=layer_id,
|
||||
quant_config=quant_config,
|
||||
is_encoder=True,
|
||||
)
|
||||
self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||
|
||||
self.activation_fn = get_act_fn(
|
||||
config.activation_function, quant_config=quant_config
|
||||
)
|
||||
|
||||
self.fc1 = ColumnParallelLinear(self.embed_dim, config.encoder_ffn_dim)
|
||||
self.fc2 = RowParallelLinear(config.encoder_ffn_dim, self.embed_dim)
|
||||
self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||
hidden_states = self.self_attn(hidden_states, forward_batch)
|
||||
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
hidden_states, _ = self.fc1(hidden_states)
|
||||
hidden_states = self.activation_fn(hidden_states)
|
||||
|
||||
hidden_states, _ = self.fc2(hidden_states)
|
||||
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
if hidden_states.dtype == torch.float16:
|
||||
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
||||
hidden_states = torch.clamp(
|
||||
hidden_states, min=-clamp_value, max=clamp_value
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WhisperDecoderLayer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: WhisperConfig,
|
||||
layer_id: Optional[int] = None,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.embed_dim = config.d_model
|
||||
|
||||
# Offset decoder layer IDs to avoid overlap with encoder layers
|
||||
decoder_self_attn_layer_id = config.encoder_layers + layer_id
|
||||
decoder_cross_attn_layer_id = (
|
||||
config.encoder_layers + config.decoder_layers + layer_id
|
||||
)
|
||||
|
||||
self.self_attn = WhisperAttention(
|
||||
embed_dim=self.embed_dim,
|
||||
num_heads=config.decoder_attention_heads,
|
||||
layer_id=decoder_self_attn_layer_id,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
self.activation_fn = get_act_fn(
|
||||
config.activation_function, quant_config=quant_config
|
||||
)
|
||||
|
||||
self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||
self.encoder_attn = WhisperAttention(
|
||||
embed_dim=self.embed_dim,
|
||||
num_heads=config.decoder_attention_heads,
|
||||
layer_id=decoder_cross_attn_layer_id,
|
||||
quant_config=quant_config,
|
||||
is_cross_attention=True,
|
||||
)
|
||||
self.encoder_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||
self.fc1 = ColumnParallelLinear(self.embed_dim, config.decoder_ffn_dim)
|
||||
self.fc2 = RowParallelLinear(config.decoder_ffn_dim, self.embed_dim)
|
||||
self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
decoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor],
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
|
||||
residual = decoder_hidden_states
|
||||
decoder_hidden_states = self.self_attn_layer_norm(decoder_hidden_states)
|
||||
decoder_hidden_states = self.self_attn(decoder_hidden_states, forward_batch)
|
||||
decoder_hidden_states = residual + decoder_hidden_states
|
||||
|
||||
residual = decoder_hidden_states
|
||||
decoder_hidden_states = self.encoder_attn_layer_norm(decoder_hidden_states)
|
||||
decoder_hidden_states = self.encoder_attn(
|
||||
decoder_hidden_states, forward_batch, encoder_hidden_states
|
||||
)
|
||||
decoder_hidden_states = residual + decoder_hidden_states
|
||||
|
||||
residual = decoder_hidden_states
|
||||
decoder_hidden_states = self.final_layer_norm(decoder_hidden_states)
|
||||
decoder_hidden_states, _ = self.fc1(decoder_hidden_states)
|
||||
decoder_hidden_states = self.activation_fn(decoder_hidden_states)
|
||||
decoder_hidden_states, _ = self.fc2(decoder_hidden_states)
|
||||
|
||||
decoder_hidden_states = residual + decoder_hidden_states
|
||||
|
||||
return decoder_hidden_states
|
||||
|
||||
|
||||
class WhisperEncoder(torch.nn.Module):
|
||||
|
||||
def __init__(
|
||||
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
embed_dim = config.d_model
|
||||
self.embed_scale = embed_dim**-0.5 if config.scale_embedding else 1.0
|
||||
|
||||
self.conv1 = torch.nn.Conv1d(
|
||||
config.num_mel_bins, embed_dim, kernel_size=3, padding=1
|
||||
)
|
||||
self.conv2 = torch.nn.Conv1d(
|
||||
embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
|
||||
)
|
||||
self.embed_positions = torch.nn.Embedding(
|
||||
config.max_source_positions, embed_dim
|
||||
)
|
||||
|
||||
self.layers = torch.nn.ModuleList(
|
||||
[
|
||||
WhisperEncoderLayer(config, id, quant_config)
|
||||
for id in range(config.encoder_layers)
|
||||
]
|
||||
)
|
||||
self.layer_norm = torch.nn.LayerNorm(config.d_model)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_features: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
):
|
||||
inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
|
||||
inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))
|
||||
|
||||
inputs_embeds = inputs_embeds.mT
|
||||
|
||||
hidden_states = inputs_embeds + self.embed_positions(position_ids)
|
||||
|
||||
for encoder_layer in self.layers:
|
||||
hidden_states = encoder_layer(hidden_states, forward_batch)
|
||||
|
||||
hidden_states = self.layer_norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WhisperDecoder(torch.nn.Module):
|
||||
|
||||
def __init__(
|
||||
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||
):
|
||||
super().__init__()
|
||||
self.max_target_positions = config.max_target_positions
|
||||
self.max_source_positions = config.max_source_positions
|
||||
self.embed_scale = config.d_model**-0.5 if config.scale_embedding else 1.0
|
||||
|
||||
self.embed_tokens = torch.nn.Embedding(
|
||||
config.vocab_size, config.d_model, padding_idx=config.pad_token_id
|
||||
)
|
||||
self.embed_positions = torch.nn.Embedding(
|
||||
self.max_target_positions, config.d_model
|
||||
)
|
||||
|
||||
self.layers = torch.nn.ModuleList(
|
||||
[
|
||||
WhisperDecoderLayer(config, layer_idx, quant_config)
|
||||
for layer_idx in range(config.decoder_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.layer_norm = torch.nn.LayerNorm(config.d_model)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor],
|
||||
forward_batch: ForwardBatch,
|
||||
position_ids=None,
|
||||
):
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
positions = self.embed_positions(position_ids)
|
||||
hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
|
||||
|
||||
for decoder_layer in self.layers:
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states, encoder_hidden_states, forward_batch
|
||||
)
|
||||
|
||||
hidden_states = self.layer_norm(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WhisperForConditionalGeneration(torch.nn.Module):
|
||||
|
||||
def __init__(
|
||||
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||
):
|
||||
super().__init__()
|
||||
self.encoder = WhisperEncoder(config, quant_config)
|
||||
self.decoder = WhisperDecoder(config, quant_config)
|
||||
self.proj_out = ParallelLMHead(
|
||||
config.vocab_size, config.d_model, quant_config=quant_config
|
||||
)
|
||||
self.logits_processor = LogitsProcessor(config)
|
||||
self.config = config
|
||||
self._encoder_cache = {}
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||
stacked_params_mapping = [
|
||||
(".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
|
||||
(".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
|
||||
(".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
|
||||
(".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
|
||||
(".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
weights_dict = dict(weights)
|
||||
|
||||
# Whisper has no k_proj bias, create zeros
|
||||
for layer_idx in range(self.config.decoder_layers):
|
||||
layer_prefix = f"model.decoder.layers.{layer_idx}.encoder_attn."
|
||||
k_proj_key = layer_prefix + "k_proj.weight"
|
||||
if k_proj_key in weights_dict:
|
||||
k_proj_weight = weights_dict[k_proj_key]
|
||||
bias_key = layer_prefix + "k_proj.bias"
|
||||
if bias_key not in weights_dict:
|
||||
weights_dict[bias_key] = torch.zeros(k_proj_weight.size(0))
|
||||
|
||||
weights_dict["proj_out.weight"] = weights_dict[
|
||||
"model.decoder.embed_tokens.weight"
|
||||
]
|
||||
|
||||
for name, loaded_weight in weights_dict.items():
|
||||
name = name.replace("model.", "")
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
if name not in params_dict:
|
||||
break
|
||||
param = params_dict[name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
if name not in params_dict:
|
||||
continue
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
|
||||
def pad_input_ids(self, input_ids: List[int], _mm_inputs: MultimodalInputs):
|
||||
return input_ids
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
**kwargs: Any,
|
||||
) -> LogitsProcessorOutput:
|
||||
dtype = self.encoder.conv1.weight.dtype
|
||||
is_decode = forward_batch.forward_mode.is_decode()
|
||||
|
||||
if is_decode:
|
||||
encoder_outputs = None
|
||||
if forward_batch.req_pool_indices is not None:
|
||||
req_indices = forward_batch.req_pool_indices.tolist()
|
||||
encoder_list = []
|
||||
for req_idx in req_indices:
|
||||
if req_idx in self._encoder_cache:
|
||||
encoder_list.append(self._encoder_cache[req_idx])
|
||||
if encoder_list:
|
||||
encoder_outputs = torch.cat(encoder_list, dim=0)
|
||||
else:
|
||||
encoder_list = []
|
||||
mm_inputs_list = forward_batch.mm_inputs if forward_batch.mm_inputs else []
|
||||
req_indices = (
|
||||
forward_batch.req_pool_indices.tolist()
|
||||
if forward_batch.req_pool_indices is not None
|
||||
else []
|
||||
)
|
||||
|
||||
for req_idx, mm_input in zip(req_indices, mm_inputs_list):
|
||||
if mm_input is None or not mm_input.mm_items:
|
||||
continue
|
||||
|
||||
features = mm_input.mm_items[0].feature
|
||||
if features.ndim == 2:
|
||||
features = features.unsqueeze(0)
|
||||
|
||||
encoder_len = features.shape[-1] // 2
|
||||
encoder_position_ids = torch.arange(encoder_len).to(
|
||||
features.device, non_blocking=True
|
||||
)
|
||||
|
||||
req_encoder_outputs = self.encoder(
|
||||
features.to(dtype), encoder_position_ids, forward_batch
|
||||
)
|
||||
req_encoder_outputs = req_encoder_outputs.squeeze(0)
|
||||
|
||||
self._encoder_cache[req_idx] = req_encoder_outputs
|
||||
encoder_list.append(req_encoder_outputs)
|
||||
|
||||
if encoder_list:
|
||||
encoder_outputs = torch.cat(encoder_list, dim=0)
|
||||
else:
|
||||
encoder_outputs = None
|
||||
|
||||
decoder_outputs = self.decoder(
|
||||
input_ids, encoder_outputs, forward_batch, positions
|
||||
)
|
||||
|
||||
logits = self.logits_processor(
|
||||
input_ids=input_ids,
|
||||
lm_head=self.proj_out,
|
||||
hidden_states=decoder_outputs,
|
||||
logits_metadata=forward_batch,
|
||||
)
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
EntryClass = [WhisperForConditionalGeneration]
|
||||
191
python/sglang/srt/multimodal/processors/whisper.py
Normal file
191
python/sglang/srt/multimodal/processors/whisper.py
Normal file
@@ -0,0 +1,191 @@
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.models.whisper import WhisperForConditionalGeneration
|
||||
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||
from sglang.srt.utils import load_audio
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ISO 639-1 supported languages for Whisper
|
||||
# From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
|
||||
# Maps ISO 639-1 code -> Full language name
|
||||
ISO639_1_SUPPORTED_LANGS = {
|
||||
"af": "Afrikaans",
|
||||
"ar": "Arabic",
|
||||
"hy": "Armenian",
|
||||
"az": "Azerbaijani",
|
||||
"be": "Belarusian",
|
||||
"bs": "Bosnian",
|
||||
"bg": "Bulgarian",
|
||||
"ca": "Catalan",
|
||||
"zh": "Chinese",
|
||||
"hr": "Croatian",
|
||||
"cs": "Czech",
|
||||
"da": "Danish",
|
||||
"nl": "Dutch",
|
||||
"en": "English",
|
||||
"et": "Estonian",
|
||||
"fi": "Finnish",
|
||||
"fr": "French",
|
||||
"gl": "Galician",
|
||||
"de": "German",
|
||||
"el": "Greek",
|
||||
"he": "Hebrew",
|
||||
"hi": "Hindi",
|
||||
"hu": "Hungarian",
|
||||
"is": "Icelandic",
|
||||
"id": "Indonesian",
|
||||
"it": "Italian",
|
||||
"ja": "Japanese",
|
||||
"kn": "Kannada",
|
||||
"kk": "Kazakh",
|
||||
"ko": "Korean",
|
||||
"lv": "Latvian",
|
||||
"lt": "Lithuanian",
|
||||
"mk": "Macedonian",
|
||||
"ms": "Malay",
|
||||
"mr": "Marathi",
|
||||
"mi": "Maori",
|
||||
"ne": "Nepali",
|
||||
"no": "Norwegian",
|
||||
"fa": "Persian",
|
||||
"pl": "Polish",
|
||||
"pt": "Portuguese",
|
||||
"ro": "Romanian",
|
||||
"ru": "Russian",
|
||||
"sr": "Serbian",
|
||||
"sk": "Slovak",
|
||||
"sl": "Slovenian",
|
||||
"es": "Spanish",
|
||||
"sw": "Swahili",
|
||||
"sv": "Swedish",
|
||||
"tl": "Tagalog",
|
||||
"ta": "Tamil",
|
||||
"th": "Thai",
|
||||
"tr": "Turkish",
|
||||
"uk": "Ukrainian",
|
||||
"ur": "Urdu",
|
||||
"vi": "Vietnamese",
|
||||
"cy": "Welsh",
|
||||
}
|
||||
|
||||
# Reverse mapping: Full language name (lowercase) -> ISO 639-1 code
|
||||
LANG_NAME_TO_CODE = {
|
||||
name.lower(): code for code, name in ISO639_1_SUPPORTED_LANGS.items()
|
||||
}
|
||||
|
||||
|
||||
def normalize_language_to_code(language: Optional[str]) -> Optional[str]:
|
||||
"""Convert a language input (full name or code) to ISO 639-1 code.
|
||||
|
||||
Args:
|
||||
language: Language as full name (e.g., 'English', 'Spanish') or
|
||||
ISO 639-1 code (e.g., 'en', 'es')
|
||||
|
||||
Returns:
|
||||
ISO 639-1 code or None if input is None
|
||||
"""
|
||||
if language is None:
|
||||
return None
|
||||
|
||||
language_lower = language.lower().strip()
|
||||
|
||||
# Check if it's already a valid ISO code
|
||||
if language_lower in ISO639_1_SUPPORTED_LANGS:
|
||||
return language_lower
|
||||
|
||||
# Check if it's a full language name
|
||||
if language_lower in LANG_NAME_TO_CODE:
|
||||
return LANG_NAME_TO_CODE[language_lower]
|
||||
|
||||
# Not recognized
|
||||
raise ValueError(
|
||||
f"Language '{language}' not recognized. "
|
||||
f"Use full name (e.g., 'English') or ISO 639-1 code (e.g., 'en')."
|
||||
)
|
||||
|
||||
|
||||
class WhisperProcessor(BaseMultimodalProcessor):
|
||||
models = [WhisperForConditionalGeneration]
|
||||
|
||||
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||
# Cache tokenizer for language token lookup
|
||||
self._tokenizer = getattr(self._processor, "tokenizer", None)
|
||||
|
||||
def _extract_language_from_request(self, request_obj) -> Optional[str]:
|
||||
sampling_params = getattr(request_obj, "sampling_params", None) or {}
|
||||
language = sampling_params.pop("language", None)
|
||||
return normalize_language_to_code(language)
|
||||
|
||||
def _get_language_token_id(self, language: Optional[str]) -> int:
|
||||
# Default to English if not specified
|
||||
if language is None:
|
||||
language = "en" # Default to English
|
||||
language_token = f"<|{language}|>"
|
||||
return self._tokenizer.convert_tokens_to_ids(language_token)
|
||||
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
image_data,
|
||||
audio_data,
|
||||
input_text,
|
||||
request_obj,
|
||||
**kwargs,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if not audio_data:
|
||||
return None
|
||||
|
||||
if len(audio_data) != 1:
|
||||
raise ValueError(
|
||||
f"Whisper expects exactly 1 audio input, got {len(audio_data)}"
|
||||
)
|
||||
|
||||
audios = [load_audio(audio) for audio in audio_data]
|
||||
|
||||
# For Whisper, ALWAYS use the proper transcription token sequence
|
||||
# and IGNORE any text prompt - Whisper is a pure speech-to-text model
|
||||
# The decoder_start_token_id and forced_decoder_ids from generation config
|
||||
# set up: <|startoftranscript|> <|lang|> <|task|> [<|notimestamps|>]
|
||||
|
||||
# Extract language from request and get token ID
|
||||
language = self._extract_language_from_request(request_obj)
|
||||
language_token_id = self._get_language_token_id(language)
|
||||
|
||||
# Build decoder input tokens
|
||||
# <|startoftranscript|> + <|lang|> + <|transcribe|> + <|notimestamps|>
|
||||
decoder_start_token_id = getattr(
|
||||
self.hf_config, "decoder_start_token_id", 50258
|
||||
)
|
||||
transcribe_token_id = self._tokenizer.convert_tokens_to_ids("<|transcribe|>")
|
||||
notimestamps_token_id = self._tokenizer.convert_tokens_to_ids(
|
||||
"<|notimestamps|>"
|
||||
)
|
||||
|
||||
input_ids = [
|
||||
decoder_start_token_id,
|
||||
language_token_id,
|
||||
transcribe_token_id,
|
||||
notimestamps_token_id,
|
||||
]
|
||||
|
||||
# Whisper expects input features padded to max_length (3000 frames = 30 seconds)
|
||||
# This is the standard context length for Whisper
|
||||
input_features = self._processor.feature_extractor(
|
||||
audios[0],
|
||||
sampling_rate=16000,
|
||||
padding="max_length", # Pad to 3000 frames
|
||||
return_tensors="pt",
|
||||
)["input_features"][0]
|
||||
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"mm_items": [
|
||||
MultimodalDataItem(
|
||||
feature=input_features,
|
||||
modality=Modality.AUDIO,
|
||||
)
|
||||
],
|
||||
}
|
||||
@@ -1027,6 +1027,23 @@ register_conv_template(
|
||||
)
|
||||
)
|
||||
|
||||
# Whisper speech-to-text template
|
||||
# Whisper uses special tokens: <|startoftranscript|>, <|en|>, <|transcribe|>, etc.
|
||||
# Audio features are processed by encoder separately, not inserted into text
|
||||
# The decoder start tokens (task, language) should be set via generation config
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="whisper",
|
||||
system_template="",
|
||||
system_message="",
|
||||
roles=("", ""),
|
||||
sep_style=SeparatorStyle.NO_COLON_SINGLE,
|
||||
sep="",
|
||||
stop_str=["<|endoftext|>"],
|
||||
audio_token="", # Empty - audio is handled by encoder, not as text token
|
||||
)
|
||||
)
|
||||
|
||||
MODEL_TYPE_TO_TEMPLATE = {
|
||||
"internvl_chat": "internvl-2-5",
|
||||
"deepseek_vl_v2": "deepseek-vl2",
|
||||
@@ -1036,6 +1053,7 @@ MODEL_TYPE_TO_TEMPLATE = {
|
||||
"minicpmo": "minicpmo",
|
||||
"deepseek-ocr": "deepseek-ocr",
|
||||
"paddleocr_vl": "paddle-ocr",
|
||||
"whisper": "whisper",
|
||||
}
|
||||
|
||||
|
||||
@@ -1129,3 +1147,11 @@ def match_paddle_ocr(model_path: str):
|
||||
return "paddle-ocr"
|
||||
model_type = get_model_type(model_path)
|
||||
return MODEL_TYPE_TO_TEMPLATE.get(model_type)
|
||||
|
||||
|
||||
@register_conv_template_matching_function
|
||||
def match_whisper(model_path: str):
|
||||
if "whisper" in model_path.lower():
|
||||
return "whisper"
|
||||
model_type = get_model_type(model_path)
|
||||
return MODEL_TYPE_TO_TEMPLATE.get(model_type)
|
||||
|
||||
@@ -1840,6 +1840,13 @@ class ServerArgs:
|
||||
self.speculative_algorithm is None
|
||||
), "Speculative decoding is currently not supported with Flex Attention backend"
|
||||
|
||||
# Encoder-decoder models (e.g., Whisper)
|
||||
if model_config.is_encoder_decoder:
|
||||
logger.warning(
|
||||
"Cuda graph is disabled for encoder-decoder models (e.g., Whisper)"
|
||||
)
|
||||
self.disable_cuda_graph = True
|
||||
|
||||
# Major NVIDIA platforms backends
|
||||
if (
|
||||
self.attention_backend == "flashmla"
|
||||
|
||||
Reference in New Issue
Block a user