main
py 101 lines 3.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import AsyncGenerator
4
5 from fastapi import Request
6
7 from vllm.engine.protocol import EngineClient
8 from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
9 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
10 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
11 from vllm.logger import init_logger
12 from vllm.outputs import RequestOutput
13
14 from ...serve.engine.protocol import ErrorResponse
15 from ..base.serving import SpeechToTextBaseServing
16 from .protocol import (
17 TranscriptionRequest,
18 TranscriptionResponse,
19 TranscriptionResponseDiarized,
20 TranscriptionResponseStreamChoice,
21 TranscriptionResponseVerbose,
22 TranscriptionStreamResponse,
23 )
24
25 logger = init_logger(__name__)
26
27
28 class OpenAIServingTranscription(SpeechToTextBaseServing):
29 """Handles transcription requests."""
30
31 def __init__(
32 self,
33 engine_client: EngineClient,
34 models: OpenAIServingModels,
35 *,
36 request_logger: RequestLogger | None,
37 return_tokens_as_token_ids: bool = False,
38 enable_force_include_usage: bool = False,
39 ):
40 super().__init__(
41 engine_client=engine_client,
42 models=models,
43 request_logger=request_logger,
44 return_tokens_as_token_ids=return_tokens_as_token_ids,
45 task_type="transcribe",
46 enable_force_include_usage=enable_force_include_usage,
47 )
48
49 async def create_transcription(
50 self,
51 audio_data: bytes,
52 request: TranscriptionRequest,
53 raw_request: Request | None = None,
54 ) -> (
55 TranscriptionResponse
56 | TranscriptionResponseVerbose
57 | TranscriptionResponseDiarized
58 | AsyncGenerator[str, None]
59 | ErrorResponse
60 ):
61 """Transcription API similar to OpenAI's API.
62
63 See https://platform.openai.com/docs/api-reference/audio/createTranscription
64 for the API specification. This API mimics the OpenAI transcription API.
65 """
66 return await self._create_speech_to_text(
67 audio_data=audio_data,
68 request=request,
69 raw_request=raw_request,
70 response_class=(
71 TranscriptionResponseVerbose
72 if request.response_format == "verbose_json"
73 else TranscriptionResponseDiarized
74 if request.response_format == "diarized_json"
75 else TranscriptionResponse
76 ),
77 stream_generator_method=self.transcription_stream_generator,
78 )
79
80 async def transcription_stream_generator(
81 self,
82 request: TranscriptionRequest,
83 result_generator: list[AsyncGenerator[RequestOutput, None]],
84 request_id: str,
85 request_metadata: RequestResponseMetadata,
86 audio_duration_s: float,
87 separator: str,
88 ) -> AsyncGenerator[str, None]:
89 generator = self._speech_to_text_stream_generator(
90 request=request,
91 list_result_generator=result_generator,
92 request_id=request_id,
93 request_metadata=request_metadata,
94 audio_duration_s=audio_duration_s,
95 chunk_object_type="transcription.chunk",
96 response_stream_choice_class=TranscriptionResponseStreamChoice,
97 stream_response_class=TranscriptionStreamResponse,
98 separator=separator,
99 )
100 async for chunk in generator:
101 yield chunk