main
py 89 lines 3.09 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import asyncio
5 from collections.abc import AsyncGenerator
6 from functools import cached_property
7 from typing import Literal, cast
8
9 import numpy as np
10
11 from vllm.engine.protocol import EngineClient, StreamingInput
12 from vllm.entrypoints.generate.base.serving import GenerateBaseServing
13 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
14 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
15 from vllm.inputs import PromptType
16 from vllm.logger import init_logger
17 from vllm.model_executor.models.interfaces import SupportsRealtime
18 from vllm.renderers.inputs.preprocess import parse_model_prompt
19
20 logger = init_logger(__name__)
21
22
23 class OpenAIServingRealtime(GenerateBaseServing):
24 """Realtime audio transcription service via WebSocket streaming.
25
26 Provides streaming audio-to-text transcription by transforming audio chunks
27 into StreamingInput objects that can be consumed by the engine.
28 """
29
30 def __init__(
31 self,
32 engine_client: EngineClient,
33 models: OpenAIServingModels,
34 *,
35 request_logger: RequestLogger | None,
36 ):
37 super().__init__(
38 engine_client=engine_client,
39 models=models,
40 request_logger=request_logger,
41 )
42
43 self.task_type: Literal["realtime"] = "realtime"
44
45 logger.info("OpenAIServingRealtime initialized for task: %s", self.task_type)
46
47 @cached_property
48 def model_cls(self) -> type[SupportsRealtime]:
49 """Get the model class that supports transcription."""
50 from vllm.model_executor.model_loader import get_model_cls
51
52 model_cls = get_model_cls(self.model_config)
53 return cast(type[SupportsRealtime], model_cls)
54
55 async def transcribe_realtime(
56 self,
57 audio_stream: AsyncGenerator[np.ndarray, None],
58 input_stream: asyncio.Queue[list[int]],
59 ) -> AsyncGenerator[StreamingInput, None]:
60 """Transform audio stream into StreamingInput for engine.generate().
61
62 Args:
63 audio_stream: Async generator yielding float32 numpy audio arrays
64 input_stream: Queue containing context token IDs from previous
65 generation outputs. Used for autoregressive multi-turn
66 processing where each generation's output becomes the context
67 for the next iteration.
68
69 Yields:
70 StreamingInput objects containing audio prompts for the engine
71
72 """
73 model_config = self.model_config
74 renderer = self.renderer
75
76 # mypy is being stupid
77 # TODO(Patrick) - fix this
78 stream_input_iter = cast(
79 AsyncGenerator[PromptType, None],
80 self.model_cls.buffer_realtime_audio(
81 audio_stream, input_stream, model_config
82 ),
83 )
84
85 async for prompt in stream_input_iter:
86 parsed_prompt = parse_model_prompt(model_config, prompt)
87 (engine_input,) = await renderer.render_cmpl_async([parsed_prompt])
88
89 yield StreamingInput(prompt=engine_input)