| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import io |
| 5 | import math |
| 6 | import time |
| 7 | import zlib |
| 8 | from collections.abc import AsyncGenerator, Callable, Set |
| 9 | from concurrent.futures import ThreadPoolExecutor |
| 10 | from functools import cached_property |
| 11 | from typing import Final, Literal, TypeAlias, TypeVar, cast |
| 12 | |
| 13 | import numpy as np |
| 14 | from fastapi import Request |
| 15 | from transformers import PreTrainedTokenizerBase |
| 16 | |
| 17 | import vllm.envs as envs |
| 18 | from vllm.engine.protocol import EngineClient |
| 19 | from vllm.entrypoints.generate.base.protocol import ( |
| 20 | DeltaMessage, |
| 21 | RequestResponseMetadata, |
| 22 | ) |
| 23 | from vllm.entrypoints.generate.base.serving import GenerateBaseServing |
| 24 | from vllm.entrypoints.openai.models.serving import OpenAIServingModels |
| 25 | from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo |
| 26 | from vllm.entrypoints.serve.engine.typing import SpeechToTextRequest |
| 27 | from vllm.entrypoints.serve.utils.api_utils import get_max_tokens |
| 28 | from vllm.entrypoints.serve.utils.request_logger import RequestLogger |
| 29 | from vllm.exceptions import VLLMValidationError |
| 30 | from vllm.inputs import EncoderDecoderInput, EngineInput |
| 31 | from vllm.logger import init_logger |
| 32 | from vllm.logprobs import FlatLogprobs, Logprob |
| 33 | from vllm.model_executor.models import SupportsTranscription |
| 34 | from vllm.multimodal.audio import get_audio_duration, split_audio |
| 35 | from vllm.multimodal.media.audio import load_audio |
| 36 | from vllm.outputs import RequestOutput |
| 37 | from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt |
| 38 | from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt |
| 39 | from vllm.sampling_params import BeamSearchParams, SamplingParams |
| 40 | from vllm.tokenizers import get_tokenizer |
| 41 | from vllm.utils.async_utils import make_async_with_semaphore, merge_async_iterators |
| 42 | |
| 43 | from ..transcription.protocol import ( |
| 44 | TranscriptionResponse, |
| 45 | TranscriptionResponseDiarized, |
| 46 | TranscriptionResponseStreamChoice, |
| 47 | TranscriptionResponseVerbose, |
| 48 | TranscriptionSegment, |
| 49 | TranscriptionStreamResponse, |
| 50 | ) |
| 51 | from ..translation.protocol import ( |
| 52 | TranslationResponse, |
| 53 | TranslationResponseStreamChoice, |
| 54 | TranslationResponseVerbose, |
| 55 | TranslationSegment, |
| 56 | TranslationStreamResponse, |
| 57 | ) |
| 58 | |
| 59 | SpeechToTextResponse: TypeAlias = ( |
| 60 | TranscriptionResponse | TranslationResponse | TranscriptionResponseDiarized |
| 61 | ) |
| 62 | SpeechToTextResponseVerbose: TypeAlias = ( |
| 63 | TranscriptionResponseVerbose | TranslationResponseVerbose |
| 64 | ) |
| 65 | SpeechToTextSegment: TypeAlias = TranscriptionSegment | TranslationSegment |
| 66 | T = TypeVar("T", bound=SpeechToTextResponse) |
| 67 | V = TypeVar("V", bound=SpeechToTextResponseVerbose) |
| 68 | S = TypeVar("S", bound=SpeechToTextSegment) |
| 69 | |
| 70 | |
| 71 | ResponseType: TypeAlias = ( |
| 72 | TranscriptionResponse |
| 73 | | TranslationResponse |
| 74 | | TranscriptionResponseVerbose |
| 75 | | TranscriptionResponseDiarized |
| 76 | | TranslationResponseVerbose |
| 77 | ) |
| 78 | |
| 79 | logger = init_logger(__name__) |
| 80 | |
| 81 | |
| 82 | def asr_inter_chunk_separator( |
| 83 | language: str | None, no_space_languages: Set[str] |
| 84 | ) -> str: |
| 85 | """Space to insert between ASR text chunks for streaming and non-streaming join. |
| 86 | |
| 87 | Languages in ``no_space_languages`` (e.g. Chinese, Japanese) use an empty |
| 88 | separator; others use a single ASCII space. |
| 89 | """ |
| 90 | return "" if language and language.lower() in no_space_languages else " " |
| 91 | |
| 92 | |
| 93 | class SpeechToTextBaseServing(GenerateBaseServing): |
| 94 | """Base class for speech-to-text operations like transcription and |
| 95 | translation.""" |
| 96 | |
| 97 | def __init__( |
| 98 | self, |
| 99 | engine_client: EngineClient, |
| 100 | models: OpenAIServingModels, |
| 101 | *, |
| 102 | request_logger: RequestLogger | None, |
| 103 | return_tokens_as_token_ids: bool = False, |
| 104 | task_type: Literal["transcribe", "translate"] = "transcribe", |
| 105 | enable_force_include_usage: bool = False, |
| 106 | ): |
| 107 | super().__init__( |
| 108 | engine_client=engine_client, |
| 109 | models=models, |
| 110 | request_logger=request_logger, |
| 111 | return_tokens_as_token_ids=return_tokens_as_token_ids, |
| 112 | ) |
| 113 | |
| 114 | self.default_sampling_params = self.model_config.get_diff_sampling_param() |
| 115 | self.task_type: Final = task_type |
| 116 | |
| 117 | self.asr_config = self.model_cls.get_speech_to_text_config( |
| 118 | self.model_config, task_type |
| 119 | ) |
| 120 | self.streaming_post_processor_cls = ( |
| 121 | self.model_cls.get_streaming_post_processor_cls() |
| 122 | ) |
| 123 | |
| 124 | self.enable_force_include_usage = enable_force_include_usage |
| 125 | |
| 126 | self.max_audio_decode_duration_s: int = envs.VLLM_MAX_AUDIO_DECODE_DURATION_S |
| 127 | self.max_audio_decode_bytes: int = envs.VLLM_MAX_AUDIO_DECODE_BYTES |
| 128 | if self.model_cls.supports_segment_timestamp: |
| 129 | self.tokenizer = cast( |
| 130 | PreTrainedTokenizerBase, |
| 131 | get_tokenizer( |
| 132 | tokenizer_name=self.model_config.tokenizer, |
| 133 | tokenizer_mode=self.model_config.tokenizer_mode, |
| 134 | ), |
| 135 | ) |
| 136 | |
| 137 | if self.default_sampling_params: |
| 138 | logger.info( |
| 139 | "Overwriting default completion sampling param with: %s", |
| 140 | self.default_sampling_params, |
| 141 | ) |
| 142 | |
| 143 | # setup preprocess resources |
| 144 | # we keep separate thread pool for frontend preprocessing instead |
| 145 | # of reusing the one from Renderer which showed lower throughput |
| 146 | # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 |
| 147 | num_audio_preprocess_workers = envs.VLLM_MAX_AUDIO_PREPROCESS_WORKERS |
| 148 | self._preprocess_executor = ThreadPoolExecutor( |
| 149 | max_workers=num_audio_preprocess_workers, |
| 150 | thread_name_prefix="stt-preprocess", |
| 151 | ) |
| 152 | self._decode_and_chunk_speech_async = make_async_with_semaphore( |
| 153 | self._decode_and_chunk_speech, executor=self._preprocess_executor |
| 154 | ) |
| 155 | |
| 156 | @cached_property |
| 157 | def model_cls(self) -> type[SupportsTranscription]: |
| 158 | from vllm.model_executor.model_loader import get_model_cls |
| 159 | |
| 160 | model_cls = get_model_cls(self.model_config) |
| 161 | return cast(type[SupportsTranscription], model_cls) |
| 162 | |
| 163 | def shutdown(self) -> None: |
| 164 | self._preprocess_executor.shutdown(wait=False) |
| 165 | |
| 166 | def _decode_and_chunk_speech( |
| 167 | self, |
| 168 | audio_data: bytes, |
| 169 | ) -> tuple[list[np.ndarray], float]: |
| 170 | # Decode audio bytes. For container formats (MP4, M4A, WebM) that |
| 171 | # soundfile cannot detect from a BytesIO stream, _load_audio_bytes |
| 172 | # transparently falls back to ffmpeg via an in-memory fd. |
| 173 | # NOTE resample to model SR here for efficiency. This is also a |
| 174 | # pre-requisite for chunking, as it assumes Whisper SR. |
| 175 | try: |
| 176 | with io.BytesIO(audio_data) as buf: |
| 177 | y, sr = load_audio( |
| 178 | buf, |
| 179 | sr=self.asr_config.sample_rate, |
| 180 | mono=True, |
| 181 | max_duration_s=self.max_audio_decode_duration_s, |
| 182 | max_decode_bytes=self.max_audio_decode_bytes, |
| 183 | ) |
| 184 | except ValueError: |
| 185 | raise |
| 186 | except Exception as exc: |
| 187 | raise ValueError("Invalid or unsupported audio file.") from exc |
| 188 | |
| 189 | duration = get_audio_duration(y=y, sr=sr) |
| 190 | do_split_audio = self.asr_config.allow_audio_chunking and ( |
| 191 | self.asr_config.max_audio_clip_s is not None |
| 192 | and duration > self.asr_config.max_audio_clip_s |
| 193 | ) |
| 194 | |
| 195 | if not do_split_audio: |
| 196 | chunks = [y] |
| 197 | else: |
| 198 | assert self.asr_config.max_audio_clip_s is not None |
| 199 | assert self.asr_config.min_energy_split_window_size is not None |
| 200 | chunks = split_audio( |
| 201 | audio_data=y, |
| 202 | sample_rate=int(sr), |
| 203 | max_clip_duration_s=self.asr_config.max_audio_clip_s, |
| 204 | overlap_duration_s=self.asr_config.overlap_chunk_second, |
| 205 | min_energy_window_size=self.asr_config.min_energy_split_window_size, |
| 206 | ) |
| 207 | |
| 208 | return chunks, duration |
| 209 | |
| 210 | async def _detect_language( |
| 211 | self, |
| 212 | audio_chunk: np.ndarray, |
| 213 | request_id: str, |
| 214 | ) -> str: |
| 215 | """Auto-detect the spoken language from an audio chunk. |
| 216 | |
| 217 | Delegates prompt construction and output parsing to the model class |
| 218 | via ``get_language_detection_prompt`` and |
| 219 | ``parse_language_detection_output``. |
| 220 | """ |
| 221 | prompt = self.model_cls.get_language_detection_prompt( |
| 222 | audio_chunk, |
| 223 | self.asr_config, |
| 224 | ) |
| 225 | allowed_token_ids = self.model_cls.get_language_token_ids( |
| 226 | self.tokenizer, |
| 227 | ) |
| 228 | sampling_params = SamplingParams( |
| 229 | max_tokens=1, |
| 230 | temperature=0.0, |
| 231 | allowed_token_ids=allowed_token_ids, |
| 232 | ) |
| 233 | |
| 234 | result_generator = self.engine_client.generate( |
| 235 | prompt, |
| 236 | sampling_params, |
| 237 | request_id, |
| 238 | ) |
| 239 | |
| 240 | try: |
| 241 | final_output: RequestOutput |
| 242 | async for final_output in result_generator: |
| 243 | if final_output.finished: |
| 244 | break |
| 245 | except asyncio.CancelledError: |
| 246 | await asyncio.gather( |
| 247 | self.engine_client.abort(request_id), |
| 248 | return_exceptions=True, |
| 249 | ) |
| 250 | raise |
| 251 | |
| 252 | token_ids = list(final_output.outputs[0].token_ids) |
| 253 | lang = self.model_cls.parse_language_detection_output( |
| 254 | token_ids, |
| 255 | self.tokenizer, |
| 256 | ) |
| 257 | |
| 258 | logger.info("Auto-detected language: '%s'", lang) |
| 259 | return lang |
| 260 | |
| 261 | async def _preprocess_speech_to_text( |
| 262 | self, |
| 263 | request: SpeechToTextRequest, |
| 264 | audio_data: bytes, |
| 265 | request_id: str, |
| 266 | ) -> tuple[list[EngineInput], float, list[float]]: |
| 267 | # Validate request |
| 268 | request.language = self.model_cls.validate_language(request.language) |
| 269 | request.to_language = ( |
| 270 | self.model_cls.validate_language(request.to_language) |
| 271 | if request.to_language |
| 272 | else None |
| 273 | ) |
| 274 | |
| 275 | # Run cpu intensive preprocess step in a separate thread pool executor. |
| 276 | chunks, duration = await self._decode_and_chunk_speech_async(audio_data) |
| 277 | |
| 278 | chunk_start_offsets: list[float] = [0.0] |
| 279 | |
| 280 | for chunk in chunks[:-1]: |
| 281 | chunk_start_offsets.append( |
| 282 | chunk_start_offsets[-1] + chunk.shape[-1] / self.asr_config.sample_rate |
| 283 | ) |
| 284 | |
| 285 | if request.language is None and getattr( |
| 286 | self.model_cls, "supports_explicit_language_detection", False |
| 287 | ): |
| 288 | # Auto-detect language from the first chunk. |
| 289 | request.language = await self._detect_language( |
| 290 | chunks[0], f"{request_id}-lang_detect" |
| 291 | ) |
| 292 | |
| 293 | parsed_prompts: list[DictPrompt] = [] |
| 294 | for chunk in chunks: |
| 295 | stt_params = request.build_stt_params( |
| 296 | audio=chunk, |
| 297 | stt_config=self.asr_config, |
| 298 | model_config=self.model_config, |
| 299 | task_type=self.task_type, |
| 300 | ) |
| 301 | prompt = self.model_cls.get_generation_prompt(stt_params) |
| 302 | |
| 303 | parsed_prompt: DictPrompt |
| 304 | if request.response_format == "verbose_json": |
| 305 | parsed_prompt = parse_enc_dec_prompt(prompt) |
| 306 | parsed_prompt = self._preprocess_verbose_prompt(parsed_prompt) |
| 307 | else: |
| 308 | parsed_prompt = parse_model_prompt(self.model_config, prompt) |
| 309 | |
| 310 | parsed_prompts.append(parsed_prompt) |
| 311 | |
| 312 | engine_inputs = await self.renderer.render_cmpl_async(parsed_prompts) |
| 313 | |
| 314 | return engine_inputs, duration, chunk_start_offsets |
| 315 | |
| 316 | def _preprocess_verbose_prompt(self, prompt: EncoderDecoderDictPrompt): |
| 317 | dec_prompt = prompt["decoder_prompt"] |
| 318 | |
| 319 | if not (isinstance(dec_prompt, dict) and "prompt" in dec_prompt): |
| 320 | raise VLLMValidationError( |
| 321 | "Expected decoder_prompt to contain text", |
| 322 | parameter="decoder_prompt", |
| 323 | value=type(dec_prompt).__name__, |
| 324 | ) |
| 325 | |
| 326 | dec_prompt["prompt"] = dec_prompt["prompt"].replace( |
| 327 | "<|notimestamps|>", "<|0.00|>" |
| 328 | ) |
| 329 | |
| 330 | return prompt |
| 331 | |
| 332 | @staticmethod |
| 333 | def _get_decoder_prompt_len(engine_inputs: list[EngineInput]) -> int: |
| 334 | """Get the length of the decoder prompt. Currently we need to offset |
| 335 | by the decoder prompt length when running beam search because the mm |
| 336 | encoder is not currently cached and runs on decode calls; because of |
| 337 | this, we need to make sure the redundant encoder calls won't exceed |
| 338 | the context :( |
| 339 | |
| 340 | FIXME (Alex) - this will be removed in the very near future once the |
| 341 | encoder/decoder caching is implemented. |
| 342 | """ |
| 343 | input_len = 0 |
| 344 | assert len(engine_inputs) > 0 |
| 345 | first_input = engine_inputs[0] |
| 346 | |
| 347 | if first_input.get("type") == "enc_dec": |
| 348 | first_input = cast(EncoderDecoderInput, first_input) |
| 349 | input_len = len(first_input["decoder_prompt"]["prompt_token_ids"]) |
| 350 | |
| 351 | return input_len |
| 352 | |
| 353 | def _get_verbose_segments( |
| 354 | self, |
| 355 | tokens: tuple, |
| 356 | log_probs: FlatLogprobs | list[dict[int, Logprob]], |
| 357 | request: SpeechToTextRequest, |
| 358 | segment_class: type[SpeechToTextSegment], |
| 359 | start_time: float = 0, |
| 360 | ) -> list[SpeechToTextSegment]: |
| 361 | """Convert tokens to verbose segments. |
| 362 | |
| 363 | This method expects the model to produce |
| 364 | timestamps as tokens (similar to Whisper). |
| 365 | If the tokens do not include timestamp information, |
| 366 | the segments may not be generated correctly. |
| 367 | |
| 368 | Note: No_speech_prob field is not supported |
| 369 | in this implementation and will be None. See docs for details. |
| 370 | """ |
| 371 | BASE_OFFSET = 0.02 |
| 372 | init_token = self.tokenizer.encode("<|0.00|>", add_special_tokens=False)[0] |
| 373 | if tokens[-1] == self.tokenizer.eos_token_id: |
| 374 | tokens = tokens[:-1] |
| 375 | |
| 376 | tokens_with_start = (init_token,) + tokens |
| 377 | segments: list[SpeechToTextSegment] = [] |
| 378 | last_timestamp_start = 0 |
| 379 | |
| 380 | if tokens_with_start[-2] < init_token and tokens_with_start[-1] >= init_token: |
| 381 | tokens_with_start = tokens_with_start + (tokens_with_start[-1],) |
| 382 | avg_logprob = 0.0 |
| 383 | for idx in range(1, len(tokens_with_start)): |
| 384 | # Timestamp tokens (e.g., <|0.00|>) are assumed to be sorted. |
| 385 | # If the ordering is violated, this slicing may produce incorrect results. |
| 386 | token = tokens_with_start[idx] |
| 387 | if token >= init_token and tokens_with_start[idx - 1] >= init_token: |
| 388 | sliced_timestamp_tokens = tokens_with_start[last_timestamp_start:idx] |
| 389 | start_timestamp = sliced_timestamp_tokens[0] - init_token |
| 390 | end_timestamp = sliced_timestamp_tokens[-1] - init_token |
| 391 | text = self.tokenizer.decode(sliced_timestamp_tokens[1:-1]) |
| 392 | text_bytes = text.encode("utf-8") |
| 393 | |
| 394 | casting_segment = cast( |
| 395 | SpeechToTextSegment, |
| 396 | segment_class( |
| 397 | id=len(segments), |
| 398 | seek=int(start_time), |
| 399 | start=start_time + BASE_OFFSET * start_timestamp, |
| 400 | end=start_time + BASE_OFFSET * end_timestamp, |
| 401 | temperature=request.temperature, |
| 402 | text=text, |
| 403 | # The compression ratio measures |
| 404 | # how compressible the generated text is. |
| 405 | # A higher ratio indicates more repetitive content, |
| 406 | # which is a strong sign of hallucination in outputs. |
| 407 | compression_ratio=len(text_bytes) |
| 408 | / len(zlib.compress(text_bytes)), |
| 409 | tokens=sliced_timestamp_tokens[1:-1], |
| 410 | avg_logprob=avg_logprob / (idx - last_timestamp_start), |
| 411 | ), |
| 412 | ) |
| 413 | segments.append(casting_segment) |
| 414 | last_timestamp_start = idx |
| 415 | avg_logprob = 0 |
| 416 | else: |
| 417 | avg_logprob += log_probs[idx - 1][token].logprob |
| 418 | return segments |
| 419 | |
| 420 | async def _create_speech_to_text( |
| 421 | self, |
| 422 | audio_data: bytes, |
| 423 | request: SpeechToTextRequest, |
| 424 | raw_request: Request, |
| 425 | response_class: type[ResponseType], |
| 426 | stream_generator_method: Callable[..., AsyncGenerator[str, None]], |
| 427 | ) -> T | V | AsyncGenerator[str, None] | ErrorResponse: |
| 428 | """Base method for speech-to-text operations like transcription and |
| 429 | translation.""" |
| 430 | if request.stream and request.use_beam_search: |
| 431 | return self.create_error_response( |
| 432 | "Streaming is not currently supported with beam search" |
| 433 | ) |
| 434 | |
| 435 | error_check_ret = await self._check_model(request) |
| 436 | if error_check_ret is not None: |
| 437 | return error_check_ret |
| 438 | |
| 439 | if not request.model: |
| 440 | request.model = self.models.model_name() |
| 441 | |
| 442 | self._preflight() |
| 443 | |
| 444 | if request.response_format not in [ |
| 445 | "text", |
| 446 | "json", |
| 447 | "verbose_json", |
| 448 | "diarized_json", |
| 449 | ]: |
| 450 | return self.create_error_response( |
| 451 | "Currently only support response_format: " |
| 452 | "`text`, `json`, `verbose_json` or `diarized_json`" |
| 453 | ) |
| 454 | |
| 455 | if ( |
| 456 | request.response_format == "verbose_json" |
| 457 | and not self.model_cls.supports_segment_timestamp |
| 458 | ): |
| 459 | return self.create_error_response( |
| 460 | f"Currently do not support verbose_json for {request.model}" |
| 461 | ) |
| 462 | |
| 463 | if ( |
| 464 | request.response_format == "diarized_json" |
| 465 | and not self.model_cls.supports_diarized_transcription |
| 466 | ): |
| 467 | return self.create_error_response( |
| 468 | f"Currently do not support diarized_json for {request.model}" |
| 469 | ) |
| 470 | |
| 471 | if ( |
| 472 | request.response_format in {"verbose_json", "diarized_json"} |
| 473 | and request.stream |
| 474 | ): |
| 475 | return self.create_error_response( |
| 476 | f"{request.response_format} format doesn't support streaming case" |
| 477 | ) |
| 478 | request_id = f"{self.task_type}-{self._base_request_id(raw_request)}" |
| 479 | |
| 480 | request_metadata = RequestResponseMetadata(request_id=request_id) |
| 481 | if raw_request: |
| 482 | raw_request.state.request_metadata = request_metadata |
| 483 | |
| 484 | lora_request = self._maybe_get_adapters(request) |
| 485 | |
| 486 | ( |
| 487 | engine_inputs, |
| 488 | duration_s, |
| 489 | chunk_start_offsets, |
| 490 | ) = await self._preprocess_speech_to_text( |
| 491 | request=request, |
| 492 | audio_data=audio_data, |
| 493 | request_id=request_id, |
| 494 | ) |
| 495 | |
| 496 | # Schedule the request and get the result generator. |
| 497 | max_model_len = self.model_config.max_model_len |
| 498 | list_result_generator: list[AsyncGenerator[RequestOutput, None]] | None = None |
| 499 | |
| 500 | input_len = ( |
| 501 | SpeechToTextBaseServing._get_decoder_prompt_len(engine_inputs) |
| 502 | if request.use_beam_search |
| 503 | else 0 |
| 504 | ) |
| 505 | |
| 506 | # Unlike most decoder-only models, whisper generation length is not |
| 507 | # constrained by the size of the input audio, which is mapped to a |
| 508 | # fixed-size log-mel-spectogram. Still, allow for fewer tokens to be |
| 509 | # generated by respecting the extra completion tokens arg. |
| 510 | max_tokens = get_max_tokens( |
| 511 | max_model_len, |
| 512 | request.max_completion_tokens, |
| 513 | input_len, |
| 514 | self.default_sampling_params, |
| 515 | ) |
| 516 | |
| 517 | if request.use_beam_search: |
| 518 | sampling_params = request.to_beam_search_params( |
| 519 | max_tokens, self.default_sampling_params |
| 520 | ) |
| 521 | else: |
| 522 | sampling_params = request.to_sampling_params( |
| 523 | max_tokens, |
| 524 | self.default_sampling_params, |
| 525 | ) |
| 526 | |
| 527 | if request.response_format == "verbose_json": |
| 528 | sampling_params.logprobs = 1 |
| 529 | |
| 530 | engine_request_ids = [ |
| 531 | request_id if len(engine_inputs) == 1 else f"{request_id}-{idx}" |
| 532 | for idx in range(len(engine_inputs)) |
| 533 | ] |
| 534 | list_result_generator = [] |
| 535 | try: |
| 536 | for request_id_item, engine_input in zip(engine_request_ids, engine_inputs): |
| 537 | self._log_inputs( |
| 538 | request_id_item, |
| 539 | engine_input, |
| 540 | params=sampling_params, |
| 541 | lora_request=lora_request, |
| 542 | ) |
| 543 | |
| 544 | trace_headers = ( |
| 545 | None |
| 546 | if raw_request is None |
| 547 | else await self._get_trace_headers(raw_request.headers) |
| 548 | ) |
| 549 | |
| 550 | if isinstance(sampling_params, BeamSearchParams): |
| 551 | generator = self.beam_search( |
| 552 | prompt=engine_input, |
| 553 | params=sampling_params, |
| 554 | request_id=request_id_item, |
| 555 | lora_request=lora_request, |
| 556 | trace_headers=trace_headers, |
| 557 | ) |
| 558 | else: |
| 559 | generator = self.engine_client.generate( |
| 560 | engine_input, |
| 561 | sampling_params, |
| 562 | request_id_item, |
| 563 | lora_request=lora_request, |
| 564 | trace_headers=trace_headers, |
| 565 | ) |
| 566 | |
| 567 | list_result_generator.append(generator) |
| 568 | except asyncio.CancelledError: |
| 569 | logger.info( |
| 570 | "Request %s cancelled; aborting %d transcription engine request(s).", |
| 571 | request_id, |
| 572 | len(engine_request_ids), |
| 573 | ) |
| 574 | await asyncio.gather( |
| 575 | self.engine_client.abort(engine_request_ids), |
| 576 | return_exceptions=True, |
| 577 | ) |
| 578 | raise |
| 579 | |
| 580 | separator = asr_inter_chunk_separator( |
| 581 | request.language, self.model_cls.no_space_languages |
| 582 | ) |
| 583 | |
| 584 | if request.stream: |
| 585 | return stream_generator_method( |
| 586 | request, |
| 587 | list_result_generator, |
| 588 | request_id, |
| 589 | request_metadata, |
| 590 | duration_s, |
| 591 | separator, |
| 592 | ) |
| 593 | # Non-streaming response. |
| 594 | try: |
| 595 | assert list_result_generator is not None |
| 596 | chunk_segment_parts: list[list[SpeechToTextSegment]] = [ |
| 597 | [] for _ in list_result_generator |
| 598 | ] |
| 599 | chunk_text_parts: list[list[str]] = [[] for _ in list_result_generator] |
| 600 | segments_types: dict[str, type[SpeechToTextSegment]] = { |
| 601 | "transcribe": TranscriptionSegment, |
| 602 | "translate": TranslationSegment, |
| 603 | } |
| 604 | segment_class: type[SpeechToTextSegment] = segments_types[self.task_type] |
| 605 | chunk_size_in_s = self.asr_config.max_audio_clip_s |
| 606 | if chunk_size_in_s is None: |
| 607 | assert len(list_result_generator) == 1, ( |
| 608 | "`max_audio_clip_s` is set to None, audio cannot be chunked" |
| 609 | ) |
| 610 | assert len(chunk_start_offsets) == len(list_result_generator) |
| 611 | result_generator = merge_async_iterators(*list_result_generator) |
| 612 | async for idx, op in result_generator: |
| 613 | start_time = chunk_start_offsets[idx] |
| 614 | if request.response_format == "verbose_json": |
| 615 | assert op.outputs[0].logprobs |
| 616 | segments: list[SpeechToTextSegment] = self._get_verbose_segments( |
| 617 | tokens=tuple(op.outputs[0].token_ids), |
| 618 | segment_class=segment_class, |
| 619 | request=request, |
| 620 | start_time=start_time, |
| 621 | log_probs=op.outputs[0].logprobs, |
| 622 | ) |
| 623 | |
| 624 | chunk_segment_parts[idx].extend(segments) |
| 625 | chunk_text_parts[idx].extend([seg.text for seg in segments]) |
| 626 | else: |
| 627 | raw_text = op.outputs[0].text |
| 628 | chunk_text_parts[idx].append( |
| 629 | self.model_cls.post_process_output(raw_text) |
| 630 | ) |
| 631 | total_segments = [ |
| 632 | segment |
| 633 | for segment_parts in chunk_segment_parts |
| 634 | for segment in segment_parts |
| 635 | ] |
| 636 | text_parts = [text for text_part in chunk_text_parts for text in text_part] |
| 637 | text = separator.join(text_parts) |
| 638 | if self.task_type == "transcribe": |
| 639 | final_response: ResponseType |
| 640 | # add usage in TranscriptionResponse. |
| 641 | usage = { |
| 642 | "type": "duration", |
| 643 | # rounded up as per openAI specs |
| 644 | "seconds": int(math.ceil(duration_s)), |
| 645 | } |
| 646 | if request.response_format == "diarized_json": |
| 647 | diarized_segments = self.model_cls.parse_diarized_transcript(text) |
| 648 | if not diarized_segments: |
| 649 | return self.create_error_response( |
| 650 | "Model output did not contain a valid diarized transcript" |
| 651 | ) |
| 652 | final_response = cast( |
| 653 | T, |
| 654 | TranscriptionResponseDiarized( |
| 655 | duration=duration_s, |
| 656 | text=separator.join( |
| 657 | segment.text for segment in diarized_segments |
| 658 | ), |
| 659 | segments=[ |
| 660 | { |
| 661 | "id": f"seg_{index}", |
| 662 | "start": segment.start, |
| 663 | "end": segment.end, |
| 664 | "text": segment.text, |
| 665 | "speaker": segment.speaker, |
| 666 | } |
| 667 | for index, segment in enumerate(diarized_segments) |
| 668 | ], |
| 669 | usage=usage, |
| 670 | ), |
| 671 | ) |
| 672 | elif request.response_format != "verbose_json": |
| 673 | final_response = cast( |
| 674 | T, TranscriptionResponse(text=text, usage=usage) |
| 675 | ) |
| 676 | else: |
| 677 | final_response = cast( |
| 678 | V, |
| 679 | TranscriptionResponseVerbose( |
| 680 | text=text, |
| 681 | language=request.language, |
| 682 | duration=duration_s, |
| 683 | segments=total_segments, |
| 684 | ), |
| 685 | ) |
| 686 | else: |
| 687 | # no usage in response for translation task |
| 688 | if request.response_format != "verbose_json": |
| 689 | final_response = cast(T, TranslationResponse(text=text)) |
| 690 | else: |
| 691 | final_response = cast( |
| 692 | V, |
| 693 | TranslationResponseVerbose( |
| 694 | text=text, |
| 695 | language=request.language, |
| 696 | duration=duration_s, |
| 697 | segments=total_segments, |
| 698 | ), |
| 699 | ) |
| 700 | return final_response |
| 701 | except asyncio.CancelledError: |
| 702 | logger.info( |
| 703 | "Request %s cancelled; aborting %d transcription engine request(s).", |
| 704 | request_id, |
| 705 | len(engine_request_ids), |
| 706 | ) |
| 707 | await asyncio.gather( |
| 708 | self.engine_client.abort(engine_request_ids), |
| 709 | return_exceptions=True, |
| 710 | ) |
| 711 | raise |
| 712 | |
| 713 | async def _speech_to_text_stream_generator( |
| 714 | self, |
| 715 | request: SpeechToTextRequest, |
| 716 | list_result_generator: list[AsyncGenerator[RequestOutput, None]], |
| 717 | request_id: str, |
| 718 | request_metadata: RequestResponseMetadata, |
| 719 | audio_duration_s: float, |
| 720 | chunk_object_type: Literal["translation.chunk", "transcription.chunk"], |
| 721 | response_stream_choice_class: type[TranscriptionResponseStreamChoice] |
| 722 | | type[TranslationResponseStreamChoice], |
| 723 | stream_response_class: type[TranscriptionStreamResponse] |
| 724 | | type[TranslationStreamResponse], |
| 725 | separator: str, |
| 726 | ) -> AsyncGenerator[str, None]: |
| 727 | created_time = int(time.time()) |
| 728 | model_name = request.model |
| 729 | |
| 730 | completion_tokens = 0 |
| 731 | num_prompt_tokens = 0 |
| 732 | |
| 733 | include_usage = self.enable_force_include_usage or request.stream_include_usage |
| 734 | include_continuous_usage = ( |
| 735 | request.stream_continuous_usage_stats |
| 736 | if include_usage and request.stream_continuous_usage_stats |
| 737 | else False |
| 738 | ) |
| 739 | |
| 740 | try: |
| 741 | for result_generator in list_result_generator: |
| 742 | beginning_of_chunk = True |
| 743 | post_processor = self.streaming_post_processor_cls() |
| 744 | async for res in result_generator: |
| 745 | # On first result. |
| 746 | if res.prompt_token_ids is not None: |
| 747 | num_prompt_tokens = len(res.prompt_token_ids) |
| 748 | if audio_tokens := self.model_cls.get_num_audio_tokens( |
| 749 | audio_duration_s, self.asr_config, self.model_config |
| 750 | ): |
| 751 | num_prompt_tokens += audio_tokens |
| 752 | |
| 753 | # We need to do it here, because if there are exceptions in |
| 754 | # the result_generator, it needs to be sent as the FIRST |
| 755 | # response (by the try...catch). |
| 756 | |
| 757 | # Just one output (n=1) supported. |
| 758 | assert len(res.outputs) == 1 |
| 759 | output = res.outputs[0] |
| 760 | |
| 761 | output_text = post_processor.process_delta( |
| 762 | output.text, output.finish_reason is not None |
| 763 | ) |
| 764 | |
| 765 | # dont add separator to the first chunk |
| 766 | if ( |
| 767 | result_generator is not list_result_generator[0] |
| 768 | and beginning_of_chunk |
| 769 | and output_text |
| 770 | ): |
| 771 | output_text = separator + output_text |
| 772 | beginning_of_chunk = False |
| 773 | |
| 774 | if output.finish_reason is None and not output_text: |
| 775 | completion_tokens += len(output.token_ids) |
| 776 | continue |
| 777 | |
| 778 | delta_message = DeltaMessage(content=output_text) |
| 779 | completion_tokens += len(output.token_ids) |
| 780 | |
| 781 | if output.finish_reason is None: |
| 782 | # Still generating, send delta update. |
| 783 | choice_data = response_stream_choice_class(delta=delta_message) |
| 784 | else: |
| 785 | # Model is finished generating. |
| 786 | choice_data = response_stream_choice_class( |
| 787 | delta=delta_message, |
| 788 | finish_reason=output.finish_reason, |
| 789 | stop_reason=output.stop_reason, |
| 790 | ) |
| 791 | |
| 792 | chunk = stream_response_class( |
| 793 | id=request_id, |
| 794 | object=chunk_object_type, |
| 795 | created=created_time, |
| 796 | choices=[choice_data], |
| 797 | model=model_name, |
| 798 | ) |
| 799 | |
| 800 | # handle usage stats if requested & if continuous |
| 801 | if include_continuous_usage: |
| 802 | chunk.usage = UsageInfo( |
| 803 | prompt_tokens=num_prompt_tokens, |
| 804 | completion_tokens=completion_tokens, |
| 805 | total_tokens=num_prompt_tokens + completion_tokens, |
| 806 | ) |
| 807 | |
| 808 | data = chunk.model_dump_json(exclude_unset=True) |
| 809 | yield f"data: {data}\n\n" |
| 810 | |
| 811 | # Once the final token is handled, if stream_options.include_usage |
| 812 | # is sent, send the usage. |
| 813 | if include_usage: |
| 814 | final_usage = UsageInfo( |
| 815 | prompt_tokens=num_prompt_tokens, |
| 816 | completion_tokens=completion_tokens, |
| 817 | total_tokens=num_prompt_tokens + completion_tokens, |
| 818 | ) |
| 819 | |
| 820 | final_usage_chunk = stream_response_class( |
| 821 | id=request_id, |
| 822 | object=chunk_object_type, |
| 823 | created=created_time, |
| 824 | choices=[], |
| 825 | model=model_name, |
| 826 | usage=final_usage, |
| 827 | ) |
| 828 | final_usage_data = final_usage_chunk.model_dump_json( |
| 829 | exclude_unset=True, exclude_none=True |
| 830 | ) |
| 831 | yield f"data: {final_usage_data}\n\n" |
| 832 | |
| 833 | # report to FastAPI middleware aggregate usage across all choices |
| 834 | request_metadata.final_usage_info = UsageInfo( |
| 835 | prompt_tokens=num_prompt_tokens, |
| 836 | completion_tokens=completion_tokens, |
| 837 | total_tokens=num_prompt_tokens + completion_tokens, |
| 838 | ) |
| 839 | |
| 840 | except Exception as e: |
| 841 | logger.exception("Error in %s stream generator.", self.task_type) |
| 842 | data = self.create_streaming_error_response(e) |
| 843 | yield f"data: {data}\n\n" |
| 844 | # Send the final done message after all response.n are finished |
| 845 | yield "data: [DONE]\n\n" |