serving_embedding.py 8.7 KB
Newer Older
1
import asyncio
2
import base64
3
import time
4
from typing import AsyncGenerator, Final, List, Literal, Optional, Union, cast
5

6
import numpy as np
7
from fastapi import Request
8
from typing_extensions import assert_never
9
10

from vllm.config import ModelConfig
11
from vllm.engine.protocol import EngineClient
12
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
13
from vllm.entrypoints.logger import RequestLogger
14
15
from vllm.entrypoints.openai.protocol import (EmbeddingChatRequest,
                                              EmbeddingRequest,
16
                                              EmbeddingResponse,
17
18
                                              EmbeddingResponseData,
                                              ErrorResponse, UsageInfo)
19
from vllm.entrypoints.openai.serving_engine import BaseModelPath, OpenAIServing
20
from vllm.logger import init_logger
21
22
from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,
                          PoolingRequestOutput)
23
from vllm.utils import merge_async_iterators
24
25
26
27

logger = init_logger(__name__)


28
def _get_embedding(
29
    output: EmbeddingOutput,
30
31
32
33
34
    encoding_format: Literal["float", "base64"],
) -> Union[List[float], str]:
    if encoding_format == "float":
        return output.embedding
    elif encoding_format == "base64":
35
36
37
        # Force to use float32 for base64 encoding
        # to match the OpenAI python client behavior
        embedding_bytes = np.array(output.embedding, dtype="float32").tobytes()
38
39
40
41
42
        return base64.b64encode(embedding_bytes).decode("utf-8")

    assert_never(encoding_format)


43
def request_output_to_embedding_response(
44
        final_res_batch: List[PoolingRequestOutput], request_id: str,
45
        created_time: int, model_name: str,
46
        encoding_format: Literal["float", "base64"]) -> EmbeddingResponse:
47
    data: List[EmbeddingResponseData] = []
48
49
    num_prompt_tokens = 0
    for idx, final_res in enumerate(final_res_batch):
50
        embedding_res = EmbeddingRequestOutput.from_base(final_res)
51
        prompt_token_ids = final_res.prompt_token_ids
52
53

        embedding = _get_embedding(embedding_res.outputs, encoding_format)
54
        embedding_data = EmbeddingResponseData(index=idx, embedding=embedding)
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
        data.append(embedding_data)

        num_prompt_tokens += len(prompt_token_ids)

    usage = UsageInfo(
        prompt_tokens=num_prompt_tokens,
        total_tokens=num_prompt_tokens,
    )

    return EmbeddingResponse(
        id=request_id,
        created=created_time,
        model=model_name,
        data=data,
        usage=usage,
    )


class OpenAIServingEmbedding(OpenAIServing):

75
76
    def __init__(
        self,
77
        engine_client: EngineClient,
78
        model_config: ModelConfig,
79
        base_model_paths: List[BaseModelPath],
80
81
        *,
        request_logger: Optional[RequestLogger],
82
        chat_template: Optional[str],
83
84
        chat_template_content_format: ChatTemplateContentFormatOption,
    ) -> None:
85
        super().__init__(engine_client=engine_client,
86
                         model_config=model_config,
87
                         base_model_paths=base_model_paths,
88
89
90
                         lora_modules=None,
                         prompt_adapters=None,
                         request_logger=request_logger)
91

92
93
        self.chat_template = chat_template
        self.chat_template_content_format: Final = chat_template_content_format
94

95
96
97
    async def create_embedding(
        self,
        request: EmbeddingRequest,
98
99
        raw_request: Optional[Request] = None,
    ) -> Union[EmbeddingResponse, ErrorResponse]:
100
101
        """
        Embedding API similar to OpenAI's API.
102
103
104
105
106
107
108
109

        See https://platform.openai.com/docs/api-reference/embeddings/create
        for the API specification. This API mimics the OpenAI Embedding API.
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

110
        encoding_format = request.encoding_format
111
112
113
114
115
        if request.dimensions is not None:
            return self.create_error_response(
                "dimensions is currently not supported")

        model_name = request.model
116
        request_id = f"embd-{self._base_request_id(raw_request)}"
117
118
        created_time = int(time.monotonic())

119
120
121
122
123
124
125
126
127
128
129
        truncate_prompt_tokens = None

        if request.truncate_prompt_tokens is not None:
            if request.truncate_prompt_tokens <= self.max_model_len:
                truncate_prompt_tokens = request.truncate_prompt_tokens
            else:
                return self.create_error_response(
                    "truncate_prompt_tokens value is "
                    "greater than max_model_len."
                    " Please, select a smaller truncation size.")

130
        try:
131
132
133
134
135
            (
                lora_request,
                prompt_adapter_request,
            ) = self._maybe_get_adapters(request)

136
            tokenizer = await self.engine_client.get_tokenizer(lora_request)
137

138
139
140
141
142
143
144
145
146
147
148
149
150
151
            if prompt_adapter_request is not None:
                raise NotImplementedError("Prompt adapter is not supported "
                                          "for embedding models")

            if isinstance(request, EmbeddingChatRequest):
                (
                    _,
                    request_prompts,
                    engine_prompts,
                ) = await self._preprocess_chat(
                    request,
                    tokenizer,
                    request.messages,
                    chat_template=request.chat_template or self.chat_template,
152
153
                    chat_template_content_format=self.
                    chat_template_content_format,
154
155
156
157
                    # In embedding requests, we are not generating tokens,
                    # so there is no need to append extra tokens to the input
                    add_generation_prompt=False,
                    continue_final_message=False,
158
159
160
161
                    truncate_prompt_tokens=truncate_prompt_tokens,
                    add_special_tokens=request.add_special_tokens,
                )
            else:
162
163
164
165
166
167
168
169
                (request_prompts,
                 engine_prompts) = await self._preprocess_completion(
                     request,
                     tokenizer,
                     request.input,
                     truncate_prompt_tokens=truncate_prompt_tokens,
                     add_special_tokens=request.add_special_tokens,
                 )
170
171
172
        except ValueError as e:
            logger.exception("Error in preprocessing prompt inputs")
            return self.create_error_response(str(e))
173

174
        # Schedule the request and get the result generator.
175
        generators: List[AsyncGenerator[PoolingRequestOutput, None]] = []
176
177
        try:
            pooling_params = request.to_pooling_params()
178

179
            for i, engine_prompt in enumerate(engine_prompts):
180
181
182
                request_id_item = f"{request_id}-{i}"

                self._log_inputs(request_id_item,
183
                                 request_prompts[i],
184
185
186
187
                                 params=pooling_params,
                                 lora_request=lora_request,
                                 prompt_adapter_request=prompt_adapter_request)

188
189
                trace_headers = (None if raw_request is None else await
                                 self._get_trace_headers(raw_request.headers))
190

191
                generator = self.engine_client.encode(
192
                    engine_prompt,
193
                    pooling_params,
194
195
                    request_id_item,
                    lora_request=lora_request,
196
                    trace_headers=trace_headers,
197
                    priority=request.priority,
198
199
200
                )

                generators.append(generator)
201
202
203
204
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
            return self.create_error_response(str(e))

205
        result_generator = merge_async_iterators(*generators)
206

207
208
        num_prompts = len(engine_prompts)

209
        # Non-streaming response
210
        final_res_batch: List[Optional[PoolingRequestOutput]]
211
        final_res_batch = [None] * num_prompts
212
213
214
        try:
            async for i, res in result_generator:
                final_res_batch[i] = res
215

216
            assert all(final_res is not None for final_res in final_res_batch)
217

218
            final_res_batch_checked = cast(List[PoolingRequestOutput],
219
220
                                           final_res_batch)

221
            response = request_output_to_embedding_response(
222
                final_res_batch_checked, request_id, created_time, model_name,
223
                encoding_format)
224
225
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")
226
227
228
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
            return self.create_error_response(str(e))
229
230

        return response