Skip to content

vllm.entrypoints.openai.serving_embedding

logger module-attribute 

logger = init_logger(__name__)

EmbeddingMixin

Bases: OpenAIServing

Source code in vllm/entrypoints/openai/serving_embedding.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class EmbeddingMixin(OpenAIServing):

    @override
    async def _preprocess(
        self,
        ctx: ServeContext,
    ) -> Optional[ErrorResponse]:
        ctx = cast(EmbeddingServeContext, ctx)
        try:
            ctx.lora_request = self._maybe_get_adapters(ctx.request)

            tokenizer = await self.engine_client.get_tokenizer(ctx.lora_request
                                                               )

            if isinstance(ctx.request, EmbeddingChatRequest):
                (
                    _,
                    ctx.request_prompts,
                    ctx.engine_prompts,
                ) = await self._preprocess_chat(
                    ctx.request,
                    tokenizer,
                    ctx.request.messages,
                    chat_template=ctx.request.chat_template
                    or ctx.chat_template,
                    chat_template_content_format=ctx.
                    chat_template_content_format,
                    # 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,
                    truncate_prompt_tokens=ctx.truncate_prompt_tokens,
                    add_special_tokens=ctx.request.add_special_tokens,
                )
            else:
                (ctx.request_prompts,
                 ctx.engine_prompts) = await self._preprocess_completion(
                     ctx.request,
                     tokenizer,
                     ctx.request.input,
                     truncate_prompt_tokens=ctx.truncate_prompt_tokens,
                     add_special_tokens=ctx.request.add_special_tokens,
                 )
            return None
        except (ValueError, TypeError) as e:
            logger.exception("Error in preprocessing prompt inputs")
            return self.create_error_response(str(e))

    @override
    def _build_response(
        self,
        ctx: ServeContext,
    ) -> Union[EmbeddingResponse, ErrorResponse]:
        items: list[EmbeddingResponseData] = []
        num_prompt_tokens = 0

        final_res_batch_checked = cast(list[PoolingRequestOutput],
                                       ctx.final_res_batch)

        for idx, final_res in enumerate(final_res_batch_checked):
            embedding_res = EmbeddingRequestOutput.from_base(final_res)

            item = EmbeddingResponseData(
                index=idx,
                embedding=_get_embedding(embedding_res.outputs,
                                         ctx.request.encoding_format),
            )
            prompt_token_ids = final_res.prompt_token_ids

            items.append(item)
            num_prompt_tokens += len(prompt_token_ids)

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

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

_build_response 

_build_response(
    ctx: ServeContext,
) -> Union[EmbeddingResponse, ErrorResponse]
Source code in vllm/entrypoints/openai/serving_embedding.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
@override
def _build_response(
    self,
    ctx: ServeContext,
) -> Union[EmbeddingResponse, ErrorResponse]:
    items: list[EmbeddingResponseData] = []
    num_prompt_tokens = 0

    final_res_batch_checked = cast(list[PoolingRequestOutput],
                                   ctx.final_res_batch)

    for idx, final_res in enumerate(final_res_batch_checked):
        embedding_res = EmbeddingRequestOutput.from_base(final_res)

        item = EmbeddingResponseData(
            index=idx,
            embedding=_get_embedding(embedding_res.outputs,
                                     ctx.request.encoding_format),
        )
        prompt_token_ids = final_res.prompt_token_ids

        items.append(item)
        num_prompt_tokens += len(prompt_token_ids)

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

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

_preprocess async 

_preprocess(ctx: ServeContext) -> Optional[ErrorResponse]
Source code in vllm/entrypoints/openai/serving_embedding.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@override
async def _preprocess(
    self,
    ctx: ServeContext,
) -> Optional[ErrorResponse]:
    ctx = cast(EmbeddingServeContext, ctx)
    try:
        ctx.lora_request = self._maybe_get_adapters(ctx.request)

        tokenizer = await self.engine_client.get_tokenizer(ctx.lora_request
                                                           )

        if isinstance(ctx.request, EmbeddingChatRequest):
            (
                _,
                ctx.request_prompts,
                ctx.engine_prompts,
            ) = await self._preprocess_chat(
                ctx.request,
                tokenizer,
                ctx.request.messages,
                chat_template=ctx.request.chat_template
                or ctx.chat_template,
                chat_template_content_format=ctx.
                chat_template_content_format,
                # 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,
                truncate_prompt_tokens=ctx.truncate_prompt_tokens,
                add_special_tokens=ctx.request.add_special_tokens,
            )
        else:
            (ctx.request_prompts,
             ctx.engine_prompts) = await self._preprocess_completion(
                 ctx.request,
                 tokenizer,
                 ctx.request.input,
                 truncate_prompt_tokens=ctx.truncate_prompt_tokens,
                 add_special_tokens=ctx.request.add_special_tokens,
             )
        return None
    except (ValueError, TypeError) as e:
        logger.exception("Error in preprocessing prompt inputs")
        return self.create_error_response(str(e))

OpenAIServingEmbedding

Bases: EmbeddingMixin

Source code in vllm/entrypoints/openai/serving_embedding.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
class OpenAIServingEmbedding(EmbeddingMixin):
    request_id_prefix = "embd"

    def __init__(
        self,
        engine_client: EngineClient,
        model_config: ModelConfig,
        models: OpenAIServingModels,
        *,
        request_logger: Optional[RequestLogger],
        chat_template: Optional[str],
        chat_template_content_format: ChatTemplateContentFormatOption,
    ) -> None:
        super().__init__(engine_client=engine_client,
                         model_config=model_config,
                         models=models,
                         request_logger=request_logger)

        self.chat_template = chat_template
        self.chat_template_content_format: Final = chat_template_content_format

    async def create_embedding(
        self,
        request: EmbeddingRequest,
        raw_request: Optional[Request] = None,
    ) -> Union[EmbeddingResponse, ErrorResponse]:
        """
        Embedding API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/embeddings/create
        for the API specification. This API mimics the OpenAI Embedding API.
        """
        model_name = self._get_model_name(request.model)
        request_id = (f"{self.request_id_prefix}-"
                      f"{self._base_request_id(raw_request)}")

        ctx = EmbeddingServeContext(
            request=request,
            raw_request=raw_request,
            model_name=model_name,
            request_id=request_id,
            chat_template=self.chat_template,
            chat_template_content_format=self.chat_template_content_format,
        )

        return await super().handle(ctx)  # type: ignore

    @override
    def _validate_request(
        self,
        ctx: ServeContext[EmbeddingRequest],
    ) -> Optional[ErrorResponse]:
        if error := super()._validate_request(ctx):
            return error

        ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens

        return None

    @override
    def _create_pooling_params(
        self,
        ctx: ServeContext[EmbeddingRequest],
    ) -> Union[PoolingParams, ErrorResponse]:
        pooling_params = super()._create_pooling_params(ctx)
        if isinstance(pooling_params, ErrorResponse):
            return pooling_params

        try:
            pooling_params.verify("embed", self.model_config)
        except ValueError as e:
            return self.create_error_response(str(e))

        return pooling_params

chat_template instance-attribute 

chat_template = chat_template

chat_template_content_format instance-attribute 

chat_template_content_format: Final = (
    chat_template_content_format
)

request_id_prefix class-attribute instance-attribute 

request_id_prefix = 'embd'

__init__ 

__init__(
    engine_client: EngineClient,
    model_config: ModelConfig,
    models: OpenAIServingModels,
    *,
    request_logger: Optional[RequestLogger],
    chat_template: Optional[str],
    chat_template_content_format: ChatTemplateContentFormatOption,
) -> None
Source code in vllm/entrypoints/openai/serving_embedding.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def __init__(
    self,
    engine_client: EngineClient,
    model_config: ModelConfig,
    models: OpenAIServingModels,
    *,
    request_logger: Optional[RequestLogger],
    chat_template: Optional[str],
    chat_template_content_format: ChatTemplateContentFormatOption,
) -> None:
    super().__init__(engine_client=engine_client,
                     model_config=model_config,
                     models=models,
                     request_logger=request_logger)

    self.chat_template = chat_template
    self.chat_template_content_format: Final = chat_template_content_format

_create_pooling_params 

_create_pooling_params(
    ctx: ServeContext[EmbeddingRequest],
) -> Union[PoolingParams, ErrorResponse]
Source code in vllm/entrypoints/openai/serving_embedding.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@override
def _create_pooling_params(
    self,
    ctx: ServeContext[EmbeddingRequest],
) -> Union[PoolingParams, ErrorResponse]:
    pooling_params = super()._create_pooling_params(ctx)
    if isinstance(pooling_params, ErrorResponse):
        return pooling_params

    try:
        pooling_params.verify("embed", self.model_config)
    except ValueError as e:
        return self.create_error_response(str(e))

    return pooling_params

_validate_request 

_validate_request(
    ctx: ServeContext[EmbeddingRequest],
) -> Optional[ErrorResponse]
Source code in vllm/entrypoints/openai/serving_embedding.py 
180
181
182
183
184
185
186
187
188
189
190
@override
def _validate_request(
    self,
    ctx: ServeContext[EmbeddingRequest],
) -> Optional[ErrorResponse]:
    if error := super()._validate_request(ctx):
        return error

    ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens

    return None

create_embedding async 

create_embedding(
    request: EmbeddingRequest,
    raw_request: Optional[Request] = None,
) -> Union[EmbeddingResponse, ErrorResponse]

Embedding API similar to OpenAI's API.

See https://platform.openai.com/docs/api-reference/embeddings/create for the API specification. This API mimics the OpenAI Embedding API.

Source code in vllm/entrypoints/openai/serving_embedding.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
async def create_embedding(
    self,
    request: EmbeddingRequest,
    raw_request: Optional[Request] = None,
) -> Union[EmbeddingResponse, ErrorResponse]:
    """
    Embedding API similar to OpenAI's API.

    See https://platform.openai.com/docs/api-reference/embeddings/create
    for the API specification. This API mimics the OpenAI Embedding API.
    """
    model_name = self._get_model_name(request.model)
    request_id = (f"{self.request_id_prefix}-"
                  f"{self._base_request_id(raw_request)}")

    ctx = EmbeddingServeContext(
        request=request,
        raw_request=raw_request,
        model_name=model_name,
        request_id=request_id,
        chat_template=self.chat_template,
        chat_template_content_format=self.chat_template_content_format,
    )

    return await super().handle(ctx)  # type: ignore

_get_embedding 

_get_embedding(
    output: EmbeddingOutput,
    encoding_format: Literal["float", "base64"],
) -> Union[list[float], str]
Source code in vllm/entrypoints/openai/serving_embedding.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
def _get_embedding(
    output: EmbeddingOutput,
    encoding_format: Literal["float", "base64"],
) -> Union[list[float], str]:
    if encoding_format == "float":
        return output.embedding
    elif encoding_format == "base64":
        # Force to use float32 for base64 encoding
        # to match the OpenAI python client behavior
        embedding_bytes = np.array(output.embedding, dtype="float32").tobytes()
        return base64.b64encode(embedding_bytes).decode("utf-8")

    assert_never(encoding_format)