serving.py 5.84 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import AsyncGenerator
4
5
6
7
8

from fastapi import Request

from vllm.engine.protocol import EngineClient
from vllm.entrypoints.logger import RequestLogger
9
from vllm.entrypoints.openai.engine.protocol import (
10
11
    ErrorResponse,
    RequestResponseMetadata,
12
)
13
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
14
from vllm.entrypoints.openai.speech_to_text.protocol import (
15
16
17
    TranscriptionRequest,
    TranscriptionResponse,
    TranscriptionResponseStreamChoice,
18
    TranscriptionResponseVerbose,
19
20
21
22
    TranscriptionStreamResponse,
    TranslationRequest,
    TranslationResponse,
    TranslationResponseStreamChoice,
23
    TranslationResponseVerbose,
24
25
    TranslationStreamResponse,
)
26
from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText
27
28
29
30
31
32
from vllm.logger import init_logger
from vllm.outputs import RequestOutput

logger = init_logger(__name__)


33
34
class OpenAIServingTranscription(OpenAISpeechToText):
    """Handles transcription requests."""
35
36
37
38
39
40

    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
        *,
41
        request_logger: RequestLogger | None,
42
        return_tokens_as_token_ids: bool = False,
43
        enable_force_include_usage: bool = False,
44
    ):
45
46
47
48
49
50
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
            task_type="transcribe",
51
            enable_force_include_usage=enable_force_include_usage,
52
        )
53
54

    async def create_transcription(
55
56
57
58
        self,
        audio_data: bytes,
        request: TranscriptionRequest,
        raw_request: Request | None = None,
59
60
61
62
63
64
    ) -> (
        TranscriptionResponse
        | TranscriptionResponseVerbose
        | AsyncGenerator[str, None]
        | ErrorResponse
    ):
65
66
67
68
69
        """Transcription API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/audio/createTranscription
        for the API specification. This API mimics the OpenAI transcription API.
        """
70
71
72
73
        return await self._create_speech_to_text(
            audio_data=audio_data,
            request=request,
            raw_request=raw_request,
74
75
76
77
78
            response_class=(
                TranscriptionResponseVerbose
                if request.response_format == "verbose_json"
                else TranscriptionResponse
            ),
79
80
            stream_generator_method=self.transcription_stream_generator,
        )
81
82

    async def transcription_stream_generator(
83
84
85
86
87
88
89
        self,
        request: TranscriptionRequest,
        result_generator: list[AsyncGenerator[RequestOutput, None]],
        request_id: str,
        request_metadata: RequestResponseMetadata,
        audio_duration_s: float,
    ) -> AsyncGenerator[str, None]:
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
        generator = self._speech_to_text_stream_generator(
            request=request,
            list_result_generator=result_generator,
            request_id=request_id,
            request_metadata=request_metadata,
            audio_duration_s=audio_duration_s,
            chunk_object_type="transcription.chunk",
            response_stream_choice_class=TranscriptionResponseStreamChoice,
            stream_response_class=TranscriptionStreamResponse,
        )
        async for chunk in generator:
            yield chunk


class OpenAIServingTranslation(OpenAISpeechToText):
    """Handles translation requests."""
106

107
108
109
110
111
    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
        *,
112
        request_logger: RequestLogger | None,
113
        return_tokens_as_token_ids: bool = False,
114
        enable_force_include_usage: bool = False,
115
    ):
116
117
118
119
120
121
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
            task_type="translate",
122
            enable_force_include_usage=enable_force_include_usage,
123
        )
124

125
    async def create_translation(
126
127
128
129
        self,
        audio_data: bytes,
        request: TranslationRequest,
        raw_request: Request | None = None,
130
131
132
133
134
135
    ) -> (
        TranslationResponse
        | TranslationResponseVerbose
        | AsyncGenerator[str, None]
        | ErrorResponse
    ):
136
        """Translation API similar to OpenAI's API.
137

138
139
        See https://platform.openai.com/docs/api-reference/audio/createTranslation
        for the API specification. This API mimics the OpenAI translation API.
140
        """
141
142
143
144
        return await self._create_speech_to_text(
            audio_data=audio_data,
            request=request,
            raw_request=raw_request,
145
146
147
148
149
            response_class=(
                TranslationResponseVerbose
                if request.response_format == "verbose_json"
                else TranslationResponse
            ),
150
151
152
153
            stream_generator_method=self.translation_stream_generator,
        )

    async def translation_stream_generator(
154
155
156
157
158
159
160
        self,
        request: TranslationRequest,
        result_generator: list[AsyncGenerator[RequestOutput, None]],
        request_id: str,
        request_metadata: RequestResponseMetadata,
        audio_duration_s: float,
    ) -> AsyncGenerator[str, None]:
161
162
163
164
165
166
167
168
169
170
171
172
        generator = self._speech_to_text_stream_generator(
            request=request,
            list_result_generator=result_generator,
            request_id=request_id,
            request_metadata=request_metadata,
            audio_duration_s=audio_duration_s,
            chunk_object_type="translation.chunk",
            response_stream_choice_class=TranslationResponseStreamChoice,
            stream_response_class=TranslationStreamResponse,
        )
        async for chunk in generator:
            yield chunk