serving_classification.py 5.43 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8

from http import HTTPStatus
from typing import Optional, Union, cast

import numpy as np
from fastapi import Request
9
from typing_extensions import override
10
11
12
13

from vllm.config import ModelConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.logger import RequestLogger
14
15
16
17
18
19
20
21
22
23
24
25
from vllm.entrypoints.openai.protocol import (
    ClassificationData,
    ClassificationRequest,
    ClassificationResponse,
    ErrorResponse,
    UsageInfo,
)
from vllm.entrypoints.openai.serving_engine import (
    ClassificationServeContext,
    OpenAIServing,
    ServeContext,
)
26
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
27
from vllm.entrypoints.renderer import RenderConfig
28
29
from vllm.logger import init_logger
from vllm.outputs import ClassificationOutput, PoolingRequestOutput
30
from vllm.pooling_params import PoolingParams
31
32
33
34
35

logger = init_logger(__name__)


class ClassificationMixin(OpenAIServing):
36
    @override
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    async def _preprocess(
        self,
        ctx: ServeContext,
    ) -> Optional[ErrorResponse]:
        """
        Process classification inputs: tokenize text, resolve adapters,
        and prepare model-specific inputs.
        """
        ctx = cast(ClassificationServeContext, ctx)
        if isinstance(ctx.request.input, str) and not ctx.request.input:
            return self.create_error_response(
                "Input cannot be empty for classification",
                status_code=HTTPStatus.BAD_REQUEST,
            )

        if isinstance(ctx.request.input, list) and len(ctx.request.input) == 0:
            return None

        try:
56
            ctx.tokenizer = await self.engine_client.get_tokenizer()
57

58
59
60
            renderer = self._get_renderer(ctx.tokenizer)
            ctx.engine_prompts = await renderer.render_prompt(
                prompt_or_prompts=ctx.request.input,
61
62
                config=self._build_render_config(ctx.request),
            )
63
64
65
66
67
68
69

            return None

        except (ValueError, TypeError) as e:
            logger.exception("Error in preprocessing prompt inputs")
            return self.create_error_response(str(e))

70
    @override
71
72
73
74
75
76
77
78
79
80
81
82
    def _build_response(
        self,
        ctx: ServeContext,
    ) -> Union[ClassificationResponse, ErrorResponse]:
        """
        Convert model outputs to a formatted classification response
        with probabilities and labels.
        """
        ctx = cast(ClassificationServeContext, ctx)
        items: list[ClassificationData] = []
        num_prompt_tokens = 0

83
        final_res_batch_checked = cast(list[PoolingRequestOutput], ctx.final_res_batch)
84
85
86
87
88
89

        for idx, final_res in enumerate(final_res_batch_checked):
            classify_res = ClassificationOutput.from_base(final_res.outputs)

            probs = classify_res.probs
            predicted_index = int(np.argmax(probs))
90
91
92
            label = getattr(self.model_config.hf_config, "id2label", {}).get(
                predicted_index
            )
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

            item = ClassificationData(
                index=idx,
                label=label,
                probs=probs,
                num_classes=len(probs),
            )

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

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

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

118
    def _build_render_config(self, request: ClassificationRequest) -> RenderConfig:
119
120
        return RenderConfig(
            max_length=self.max_model_len,
121
122
            truncate_prompt_tokens=request.truncate_prompt_tokens,
        )
123

124
125
126
127
128
129
130
131
132
133
134

class ServingClassification(ClassificationMixin):
    request_id_prefix = "classify"

    def __init__(
        self,
        engine_client: EngineClient,
        model_config: ModelConfig,
        models: OpenAIServingModels,
        *,
        request_logger: Optional[RequestLogger],
135
        log_error_stack: bool = False,
136
137
138
139
140
141
    ) -> None:
        super().__init__(
            engine_client=engine_client,
            model_config=model_config,
            models=models,
            request_logger=request_logger,
142
            log_error_stack=log_error_stack,
143
144
145
146
147
148
149
        )

    async def create_classify(
        self,
        request: ClassificationRequest,
        raw_request: Request,
    ) -> Union[ClassificationResponse, ErrorResponse]:
150
        model_name = self.models.model_name()
151
        request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}"
152
153
154
155
156
157
158
159
160

        ctx = ClassificationServeContext(
            request=request,
            raw_request=raw_request,
            model_name=model_name,
            request_id=request_id,
        )

        return await super().handle(ctx)  # type: ignore
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

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

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

        return pooling_params